quak/src/metadata-backup.ts

import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import * as jpeg from "jpeg-js";
import exifReader from "exif-reader";
import type { Client } from "./client.js";
import type { Library, Photo } from "./library/index.js";
import { sanitizeFileName } from "./filename.js";
import { fetchMLData } from "./mldata-fetch.js";
import type { EnteFile } from "./model/types.js";

export type ProgressCallback = (message: string) => void;

export interface MetadataBackupOptions {
    exif?: boolean;
    onProgress?: ProgressCallback;
}

// Extract the raw EXIF APP1 segment from JPEG bytes. Returns the EXIF
// data buffer (starting after the APP1 length field, at the "Exif\0\0"
// header) or undefined if no APP1 marker is found.
const extractExifFromJpeg = (buf: Uint8Array): Buffer | undefined => {
    if (buf[0] !== 0xff || buf[1] !== 0xd8) return undefined;
    let offset = 2;
    while (offset < buf.length - 1) {
        if (buf[offset] !== 0xff) return undefined;
        const marker = buf[offset + 1]!;
        if (marker === 0xda) break; // start of scan, no more markers
        if (offset + 3 >= buf.length) break;
        const len = (buf[offset + 2]! << 8) | buf[offset + 3]!;
        if (marker === 0xe1) {
            // APP1 — check for "Exif\0\0" header
            if (
                buf[offset + 4] === 0x45 &&
                buf[offset + 5] === 0x78 &&
                buf[offset + 6] === 0x69 &&
                buf[offset + 7] === 0x66
            ) {
                return Buffer.from(
                    buf.buffer,
                    buf.byteOffset + offset + 4,
                    len - 2,
                );
            }
        }
        offset += 2 + len;
    }
    return undefined;
};

const extractImageMetadata = (
    fileBytes: Uint8Array,
): Record<string, unknown> | undefined => {
    try {
        const result: Record<string, unknown> = {};

        // Try to get dimensions from JPEG decode
        try {
            const decoded = jpeg.decode(fileBytes, {
                useTArray: true,
                formatAsRGBA: false,
            });
            result.format = "jpeg";
            result.width = decoded.width;
            result.height = decoded.height;
        } catch {
            // Not a JPEG or corrupt; still try EXIF extraction
        }

        const exifBuf = extractExifFromJpeg(fileBytes);
        if (exifBuf) {
            try {
                result.exif = exifReader(exifBuf);
            } catch {
                result.exifRaw = exifBuf.toString("base64");
            }
        }

        // Extract XMP (look for "http://ns.adobe.com/xap" in the bytes)
        const xmpStart = Buffer.from(fileBytes).indexOf("<?xpacket begin");
        if (xmpStart !== -1) {
            const xmpEnd = Buffer.from(fileBytes).indexOf(
                "<?xpacket end",
                xmpStart,
            );
            if (xmpEnd !== -1) {
                const end = Buffer.from(fileBytes).indexOf("?>", xmpEnd);
                result.xmp = Buffer.from(fileBytes)
                    .subarray(xmpStart, end !== -1 ? end + 2 : xmpEnd + 50)
                    .toString("utf-8");
            }
        }

        return Object.keys(result).length > 0 ? result : undefined;
    } catch {
        return undefined;
    }
};

// Read a file's original bytes through the library's content cache and extract
// its embedded image metadata. The bytes come from `photo.original()` — the
// same on-disk cache the rest of the library fills — rather than a fresh
// per-call download to a throwaway temp file.
const extractExif = async (
    photo: Photo,
): Promise<Record<string, unknown> | undefined> => {
    try {
        const { path } = await photo.original();
        const fileBytes = new Uint8Array(readFileSync(path));
        return extractImageMetadata(fileBytes);
    } catch {
        return undefined;
    }
};

// Dump every decrypted metadata layer the account holds into a directory tree
// of plain JSON: account, per-collection, and per-file records including the
// private and public magic metadata and (by default) the ML data. Collections
// and files are enumerated from the library's cache rather than a fresh server
// scan; the ML fetch and EXIF extraction are unchanged.
export const runMetadataBackup = async (
    lib: Library,
    client: Client,
    outDir: string,
    opts?: MetadataBackupOptions,
): Promise<void> => {
    const log = opts?.onProgress ?? (() => {});
    const wantExif = opts?.exif ?? false;

    mkdirSync(outDir, { recursive: true });
    mkdirSync(join(outDir, "collections"), { recursive: true });

    const { email, userID } = client.whoami();
    writeFileSync(
        join(outDir, "account.json"),
        JSON.stringify({ email, userID }, null, 2),
    );

    log("Fetching collections...");

    // Enumerate through the library's read surface. Each album carries its
    // photos, but the full decrypted `Collection`/`EnteFile` records (with the
    // magic-metadata layers this dump exists to preserve) come from the
    // library's by-id accessors.
    const allFiles: { file: EnteFile; photo: Photo; colDirName: string }[] = [];
    const fileKeys = new Map<number, Uint8Array>();
    const seenFileIDs = new Set<number>();

    for (const album of lib.albums.list()) {
        const col = lib.getCollection(album.collectionID);
        if (!col) continue;

        const dirName = `${col.id}-${sanitizeFileName(col.name, "unnamed")}`;
        const colDir = join(outDir, "collections", dirName);
        mkdirSync(colDir, { recursive: true });

        const collectionMeta: Record<string, unknown> = {
            id: col.id,
            name: col.name,
            type: col.type,
            ownerID: col.ownerID,
            isShared: col.isShared,
            updationTime: col.updationTime,
        };
        if (col.magicMetadata) collectionMeta.magicMetadata = col.magicMetadata;
        if (col.pubMagicMetadata)
            collectionMeta.pubMagicMetadata = col.pubMagicMetadata;
        if (col.sharedMagicMetadata)
            collectionMeta.sharedMagicMetadata = col.sharedMagicMetadata;

        writeFileSync(
            join(colDir, "_collection.json"),
            JSON.stringify(collectionMeta, null, 2),
        );

        log(`[${col.name}] Fetching files...`);
        const photos = album.photos.list();
        log(`[${col.name}] ${photos.length} file(s)`);

        for (const photo of photos) {
            const file = lib.getFile(col.id, photo.fileID);
            if (!file) continue;
            allFiles.push({ file, photo, colDirName: dirName });
            if (!seenFileIDs.has(file.id)) {
                fileKeys.set(file.id, file.key);
                seenFileIDs.add(file.id);
            }
        }
    }

    log("Fetching ML data (face detections, CLIP embeddings)...");
    const mlDataMap = await fetchMLData(
        client.getApiClient(),
        [...fileKeys.keys()],
        fileKeys,
    );
    log(`Got ML data for ${mlDataMap.size} file(s)`);

    const writtenFileIDs = new Set<number>();
    for (const { file, photo, colDirName } of allFiles) {
        const colDir = join(outDir, "collections", colDirName);

        const fileMeta: Record<string, unknown> = {
            id: file.id,
            collectionID: file.collectionID,
            ownerID: file.ownerID,
            metadata: file.metadata,
            updationTime: file.updationTime,
        };
        if (file.magicMetadata) fileMeta.magicMetadata = file.magicMetadata;
        if (file.pubMagicMetadata)
            fileMeta.pubMagicMetadata = file.pubMagicMetadata;

        const ml = mlDataMap.get(file.id);
        if (ml) fileMeta.mlData = ml;

        if (wantExif && !writtenFileIDs.has(file.id)) {
            log(`[${file.metadata.title}] Extracting EXIF...`);
            const exifData = await extractExif(photo);
            if (exifData) fileMeta.imageMetadata = exifData;
        }
        writtenFileIDs.add(file.id);

        writeFileSync(
            join(colDir, `${file.id}.json`),
            JSON.stringify(fileMeta, null, 2),
        );
    }

    log("Metadata backup complete.");
};