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. A live row
// that cannot be decrypted, or whose metadata cannot be parsed, is in
// `undecryptable` with the reason instead, so one bad row does not fail the
// call; absent means there was none. `cursor` is the max `updationTime` seen,
// those rows included, to pass back into the next call.
export interface CollectionsPage {
    collections: Collection[];
    deleted: number[];
    undecryptable?: { id: number; error: string }[];
    cursor: number;
}

export interface FilesPage {
    files: EnteFile[];
    deleted: number[];
    undecryptable?: { id: number; error: string }[];
    cursor: number;
}

const errorMessage = (err: unknown): string =>
    err instanceof Error ? err.message : String(err);

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,
        );
    }

    // Restore a client from a `toJSON()` snapshot. The snapshot usually comes
    // straight from `JSON.parse` of a file on disk, so every field is checked
    // before use; a bad one throws an error naming it. Needs `init()` first.
    static fromJSON(snapshot: unknown, apiOptions?: ApiClientOptions): Client {
        const invalid = (field: string, problem: string): Error =>
            new Error(`Invalid session data: ${field} ${problem}`);

        if (typeof snapshot !== "object" || snapshot === null) {
            throw new Error("Invalid session data: not a JSON object");
        }
        const s = snapshot as Record<string, unknown>;
        for (const field of ["email", "token"]) {
            if (typeof s[field] !== "string" || s[field] === "") {
                throw invalid(field, "must be a non-empty string");
            }
        }
        if (!Number.isInteger(s.userID)) {
            throw invalid("userID", "must be an integer");
        }
        const key = (field: string): Uint8Array => {
            const value = s[field];
            if (typeof value !== "string") {
                throw invalid(field, "must be a base64 string");
            }
            let bytes: Uint8Array;
            try {
                bytes = fromBase64(value);
            } catch {
                throw invalid(field, "is not valid base64");
            }
            // The master key (secretbox) and the key pair (box) are all 32 bytes.
            if (bytes.length !== 32) {
                throw invalid(
                    field,
                    `must decode to 32 bytes, got ${bytes.length}`,
                );
            }
            return bytes;
        };

        const api = new ApiClient({
            ...apiOptions,
            authToken: s.token as string,
        });
        return new Client(
            api,
            s.email as string,
            s.userID as number,
            key("masterKey"),
            key("secretKey"),
            key("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();
        const token = this.api.getAuthToken();
        if (!token) {
            throw new Error("Cannot serialize client: it has no auth token");
        }
        return {
            email: this.email,
            userID: this.userID,
            token,
            masterKey: toBase64(this.masterKey),
            secretKey: toBase64(this.secretKey),
            publicKey: toBase64(this.publicKey),
        };
    }

    // Ends this client's session on the server (`POST /users/logout`), so the
    // token stops working everywhere, including in any saved copy of it. This
    // client is left as it was; call `logout()` to clear it.
    async logoutOnServer(): Promise<void> {
        this.assertLoggedIn();
        await this.api.postJSON("/users/logout", {});
    }

    // Zeroes the key buffers in place, so any copy of the reference held
    // elsewhere is wiped too. Every method checks `assertLoggedIn` before
    // touching the keys, so nothing decrypts with the zeroed keys.
    logout(): void {
        this.loggedOut = true;
        this.api.clearAuthToken();
        this.masterKey.fill(0);
        this.secretKey.fill(0);
        this.publicKey.fill(0);
    }

    // Enumerate collections changed since `sinceTime`. Live collections are
    // decrypted; tombstoned ones (isDeleted) are surfaced as bare ids, and
    // ones that do not decrypt as ids with the reason. 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 });
        // logout() may have zeroed the keys while the request was in flight.
        this.assertLoggedIn();

        const collections: Collection[] = [];
        const deleted: number[] = [];
        const undecryptable: { id: number; error: string }[] = [];
        let cursor = args.sinceTime;
        for (const raw of raws) {
            if (raw.isDeleted) {
                deleted.push(raw.id);
            } else {
                try {
                    collections.push(
                        decryptCollection(
                            raw,
                            {
                                masterKey: this.masterKey,
                                publicKey: this.publicKey,
                                secretKey: this.secretKey,
                            },
                            this.userID,
                        ),
                    );
                } catch (err) {
                    undecryptable.push({
                        id: raw.id,
                        error: errorMessage(err),
                    });
                }
            }
            if (raw.updationTime > cursor) cursor = raw.updationTime;
        }
        return { collections, deleted, undecryptable, 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, and ones that do not decrypt or parse as ids with
    // the reason. 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[] = [];
        const undecryptable: { id: number; error: string }[] = [];
        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 {
                    try {
                        files.push(decryptFile(raw, collectionKey));
                    } catch (err) {
                        undecryptable.push({
                            id: raw.id,
                            error: errorMessage(err),
                        });
                    }
                }
                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, undecryptable, cursor };
    }

    // Whole-account listing: every live collection that decrypts, 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 that decrypts, 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);
    }
}