quak/src/client.ts

import { ApiClient, type ApiClientOptions } from "./api/client.js";
import {
    beginLogin,
    submitTOTP,
    requestEmailOTP,
    submitEmailOTP,
} from "./auth/login.js";
import { unwrapAuth } from "./auth/unwrap.js";
import { init, fromBase64, toBase64 } from "./crypto/index.js";
import { fetchMLDataBatch, type MLData } from "./mldata-fetch.js";
import { decryptCollection, decryptFile } from "./model/index.js";
import {
    downloadFile as dlFile,
    downloadThumbnail as dlThumb,
} from "./download/index.js";
import {
    makeDownloadContentSource,
    type ContentSource,
} from "./library/content.js";
import type {
    Collection,
    EnteFile,
    RawCollection,
    RawEnteFile,
} from "./model/types.js";
import type { DownloadResult } from "./download/index.js";

export interface LoginOptions {
    email: string;
    password: string;
    totp?: () => Promise<string>;
    emailOTP?: () => Promise<string>;
    apiOptions?: ApiClientOptions;
}

export interface ClientSnapshot {
    email: string;
    userID: number;
    token: string;
    masterKey: string;
    secretKey: string;
    publicKey: string;
}

// The result of a resumable enumeration. Live decrypted records and deleted
// ids are kept apart on purpose: a tombstone carries no key or metadata to
// decrypt, so it is a bare id rather than a hollowed-out record. `cursor` is
// the max `updationTime` seen, to pass back into the next call.
export interface CollectionsPage {
    collections: Collection[];
    deleted: number[];
    cursor: number;
}

export interface FilesPage {
    files: EnteFile[];
    deleted: number[];
    cursor: number;
}

export class Client {
    private readonly api: ApiClient;
    private readonly email: string;
    private readonly userID: number;
    private readonly masterKey: Uint8Array;
    private readonly secretKey: Uint8Array;
    private readonly publicKey: Uint8Array;
    private loggedOut = false;

    private constructor(
        api: ApiClient,
        email: string,
        userID: number,
        masterKey: Uint8Array,
        secretKey: Uint8Array,
        publicKey: Uint8Array,
    ) {
        this.api = api;
        this.email = email;
        this.userID = userID;
        this.masterKey = masterKey;
        this.secretKey = secretKey;
        this.publicKey = publicKey;
    }

    static async login(opts: LoginOptions): Promise<Client> {
        await init();
        const api = new ApiClient(opts.apiOptions);
        const challenge = await beginLogin(api, opts.email, opts.password);

        let response;
        if (challenge.kind === "complete") {
            response = challenge.response;
        } else if (challenge.kind === "totp") {
            if (!opts.totp)
                throw new Error(
                    "Account requires TOTP but no totp callback provided",
                );
            const code = await opts.totp();
            response = await submitTOTP(api, challenge.sessionID, code);
        } else if (challenge.kind === "emailOTP") {
            if (!opts.emailOTP)
                throw new Error(
                    "Account requires email OTP but no emailOTP callback provided",
                );
            await requestEmailOTP(api, opts.email);
            const code = await opts.emailOTP();
            response = await submitEmailOTP(api, opts.email, code);
        } else if (challenge.kind === "passkey") {
            throw new Error("Passkey authentication is not supported by quak");
        } else {
            throw new Error(`Unknown login challenge kind`);
        }

        const unwrapped = await unwrapAuth(response, opts.password);
        api.setAuthToken(unwrapped.token);

        return new Client(
            api,
            opts.email,
            response.id,
            unwrapped.masterKey,
            unwrapped.secretKey,
            unwrapped.publicKey,
        );
    }

    static fromJSON(
        snapshot: ClientSnapshot,
        apiOptions?: ApiClientOptions,
    ): Client {
        const api = new ApiClient({ ...apiOptions, authToken: snapshot.token });
        return new Client(
            api,
            snapshot.email,
            snapshot.userID,
            fromBase64(snapshot.masterKey),
            fromBase64(snapshot.secretKey),
            fromBase64(snapshot.publicKey),
        );
    }

    getApiClient(): ApiClient {
        this.assertLoggedIn();
        return this.api;
    }

    // The content-cache byte source over this client's API: each fetch is the
    // download layer's request + streaming decrypt + atomic write. `Library`
    // calls this to enable the on-disk content cache.
    contentSource(): ContentSource {
        this.assertLoggedIn();
        return makeDownloadContentSource(this.api);
    }

    private assertLoggedIn(): void {
        if (this.loggedOut) throw new Error("Client has been logged out");
    }

    whoami(): { email: string; userID: number } {
        this.assertLoggedIn();
        return { email: this.email, userID: this.userID };
    }

    toJSON(): ClientSnapshot {
        this.assertLoggedIn();
        return {
            email: this.email,
            userID: this.userID,
            token: this.api["token"]!,
            masterKey: toBase64(this.masterKey),
            secretKey: toBase64(this.secretKey),
            publicKey: toBase64(this.publicKey),
        };
    }

    logout(): void {
        this.loggedOut = true;
        this.api.clearAuthToken();
    }

    // Enumerate collections changed since `sinceTime`. Live collections are
    // decrypted; tombstoned ones (isDeleted) are surfaced as bare ids. The
    // returned cursor is the max `updationTime` seen — including tombstones, so
    // the next sync resumes past them — falling back to `sinceTime` when the
    // response is empty. `/collections/v2` returns the whole changed set in one
    // response, so there is no pagination here.
    async collectionsSince(args: {
        sinceTime: number;
    }): Promise<CollectionsPage> {
        this.assertLoggedIn();
        const { collections: raws } = await this.api.getJSON<{
            collections: RawCollection[];
        }>("/collections/v2", { sinceTime: args.sinceTime });

        const collections: Collection[] = [];
        const deleted: number[] = [];
        let cursor = args.sinceTime;
        for (const raw of raws) {
            if (raw.isDeleted) {
                deleted.push(raw.id);
            } else {
                collections.push(
                    decryptCollection(
                        raw,
                        {
                            masterKey: this.masterKey,
                            publicKey: this.publicKey,
                            secretKey: this.secretKey,
                        },
                        this.userID,
                    ),
                );
            }
            if (raw.updationTime > cursor) cursor = raw.updationTime;
        }
        return { collections, deleted, cursor };
    }

    // Enumerate a collection's files changed since `sinceTime`, paginating the
    // diff from that cursor. Live rows are decrypted; tombstoned ones are
    // surfaced as bare ids. Returns the final cursor to resume from.
    async filesSince(args: {
        collectionID: number;
        collectionKey: Uint8Array;
        sinceTime: number;
    }): Promise<FilesPage> {
        this.assertLoggedIn();
        const { collectionID, collectionKey } = args;
        const files: EnteFile[] = [];
        const deleted: number[] = [];
        let cursor = args.sinceTime;
        for (;;) {
            const { diff, hasMore } = await this.api.getJSON<{
                diff: RawEnteFile[];
                hasMore: boolean;
            }>("/collections/v2/diff", { collectionID, sinceTime: cursor });

            let pageMax = cursor;
            for (const raw of diff) {
                if (raw.isDeleted) {
                    deleted.push(raw.id);
                } else {
                    files.push(decryptFile(raw, collectionKey));
                }
                if (raw.updationTime > pageMax) pageMax = raw.updationTime;
            }

            if (!hasMore) {
                cursor = pageMax;
                break;
            }
            // The server says there is more, but this page did not advance the
            // cursor: following hasMore would refetch the same page forever
            // (#7). Stop with a clear error instead of looping.
            if (pageMax <= cursor) {
                throw new Error(
                    `/collections/v2/diff for collection ${collectionID} ` +
                        `returned hasMore with a cursor that did not advance ` +
                        `(stuck at ${cursor}); refusing to loop`,
                );
            }
            cursor = pageMax;
        }
        return { files, deleted, cursor };
    }

    // Whole-account listing: every live collection, deletions hidden. A thin
    // wrapper over `collectionsSince` from the beginning of time.
    async listCollections(): Promise<Collection[]> {
        const { collections } = await this.collectionsSince({ sinceTime: 0 });
        return collections;
    }

    // Every live file in a collection, deletions hidden. A thin wrapper over
    // `filesSince` from the beginning of time.
    async listFiles(
        collectionID: number,
        collectionKey: Uint8Array,
    ): Promise<EnteFile[]> {
        const { files } = await this.filesSince({
            collectionID,
            collectionKey,
            sinceTime: 0,
        });
        return files;
    }

    // Fetch machine-learning data (face detections + CLIP embeddings) for up
    // to a batch of files, each decrypted with its own key. One request; the
    // library batches at `MLDATA_BATCH_SIZE` and schedules each batch through
    // its metadata request pool.
    async fetchMLData(args: {
        fileIDs: number[];
        fileKeys: Map<number, Uint8Array>;
    }): Promise<Map<number, MLData>> {
        this.assertLoggedIn();
        return fetchMLDataBatch(this.api, args.fileIDs, args.fileKeys);
    }

    async downloadFile(
        file: EnteFile,
        outPath?: string,
    ): Promise<DownloadResult> {
        this.assertLoggedIn();
        return dlFile(this.api, file, outPath);
    }

    async downloadThumbnail(
        file: EnteFile,
        outPath?: string,
    ): Promise<DownloadResult> {
        this.assertLoggedIn();
        return dlThumb(this.api, file, outPath);
    }
}