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

// Find the raw EXIF APP1 segment in JPEG bytes. Returns `exif` (the segment
// data, starting at the "Exif\0\0" header) when there is one, nothing when the
// bytes are not a JPEG or carry no EXIF, and `error` when the segment layout is
// malformed. Each segment length is checked against the bytes that remain and
// each step moves forward by at least 4 bytes, so the scan ends on any input.
export const extractExifFromJpeg = (
    buf: Uint8Array,
): { exif?: Buffer; error?: string } => {
    if (buf[0] !== 0xff || buf[1] !== 0xd8) return {};
    let offset = 2;
    while (offset < buf.length) {
        if (offset + 2 > buf.length)
            return { error: `truncated segment marker at byte ${offset}` };
        if (buf[offset] !== 0xff)
            return { error: `no segment marker at byte ${offset}` };
        const marker = buf[offset + 1]!;
        if (marker === 0xda) return {}; // start of scan, no more markers
        if (offset + 4 > buf.length)
            return { error: `truncated segment length at byte ${offset}` };
        const len = (buf[offset + 2]! << 8) | buf[offset + 3]!;
        // The length counts its own two bytes, so anything under 2 is invalid.
        if (len < 2)
            return {
                error: `segment length ${len} at byte ${offset} is too small`,
            };
        if (offset + 2 + len > buf.length)
            return {
                error: `segment length ${len} at byte ${offset} runs past the end of the file`,
            };
        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 {
                    exif: Buffer.from(
                        buf.buffer,
                        buf.byteOffset + offset + 4,
                        len - 2,
                    ),
                };
            }
        }
        offset += 2 + len;
    }
    return { error: "file ends before the image data" };
};

// Extract dimensions, EXIF and XMP from a file's bytes. When the EXIF segment
// is malformed or cannot be parsed, the record carries the reason in
// `exifError`.
export const extractImageMetadata = (
    fileBytes: Uint8Array,
): Record<string, unknown> | undefined => {
    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 every original is a JPEG (PNG, HEIC, video), so a failed decode
        // is expected and only means no dimensions; a malformed JPEG is still
        // reported below through `exifError`.
    }

    const { exif, error } = extractExifFromJpeg(fileBytes);
    if (error) result.exifError = error;
    if (exif) {
        try {
            result.exif = exifReader(exif);
        } catch (err) {
            result.exifRaw = exif.toString("base64");
            result.exifError = err instanceof Error ? err.message : String(err);
        }
    }

    // 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;
};

// 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> => {
    const { path } = await photo.original();
    const fileBytes = new Uint8Array(readFileSync(path));
    return extractImageMetadata(fileBytes);
};

// 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...`);
            try {
                const exifData = await extractExif(photo);
                if (exifData) fileMeta.imageMetadata = exifData;
            } catch (err) {
                fileMeta.imageMetadataError =
                    err instanceof Error ? err.message : String(err);
            }
        }
        writtenFileIDs.add(file.id);

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

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