PEE/src/jpg.ts

import { Buffer } from "buffer";
import type { ImageProcessor } from "./main";
import pngv3 from "./pngv3";
import f5 from 'f5stegojs';
import { settings } from "./stores";
import { decodeCoom3Payload } from "./utils";

export let csettings: Parameters<typeof settings['set']>[0];

settings.subscribe(b => {
    csettings = b;
});

export const convertToPng = async (f: File): Promise<Blob | undefined> => {
    const can = document.createElement("canvas");

    const url = URL.createObjectURL(f);
    try {
        let dims: [number, number];
        let source: CanvasImageSource;
        if (f.type.startsWith("image")) {
            const imgElem = document.createElement('img');
            imgElem.src = url;
            await new Promise(_ => imgElem.onload = _);
            dims = [imgElem.naturalWidth, imgElem.naturalHeight];
            source = imgElem;
        } else if (f.type.startsWith("video")) {
            const vidElem = document.createElement('video');
            vidElem.src = url;
            await new Promise(_ => vidElem.onloadedmetadata = _);
            vidElem.currentTime = 0;
            await new Promise(_ => vidElem.onloadeddata = _);
            await new Promise(requestAnimationFrame);
            await new Promise(requestAnimationFrame);
            await new Promise(requestAnimationFrame);
            dims = [vidElem.videoWidth, vidElem.videoHeight];
            source = vidElem;
        } else
            return;
        can.width = dims[0];
        can.height = dims[1];
        const ctx = can.getContext("2d");

        if (!ctx)
            return;
        ctx.drawImage(source, 0, 0, dims[0], dims[1]);

        const blob = await new Promise<Blob | null>(_ => can.toBlob(_, "image/png"));
        if (!blob)
            return;
        return blob;
    } finally {
        URL.revokeObjectURL(url);
    }
};

const key = Buffer.from("CUNNYCUNNYCUNNY");
const f5inst = new f5(key);

const injectTrue = async (b: File, links: string[]) => {
    // TODO: maybe do a lossless crop/embed/concat?
    if (b.size / 20 < links.join(' ').length)
        throw "Image too small to embed.";
    const arr = new Uint8Array(new Uint8Array(await b.arrayBuffer()));
    const buff = f5inst.embed(arr, new TextEncoder().encode(links.join(' ')));
    return Buffer.from(buff);
};

const inject = async (b: File, links: string[]) => {
    if (csettings.jpeg)
        return injectTrue(b, links);
    const pngfile = await convertToPng(b);
    if (!pngfile || pngfile.size > 3000 * 1024) {
        throw new Error("Couldn't convert file to PNG: resulting filesize too big.");
    }
    return pngv3.inject!(new File([pngfile], b.name), links);
};

// unfortunately, because of the way f5 work, we can't determine
// if there's an embedded message until we have the complete file
// but the way PEE was designed forces us to just try to extract something until it works
const has_embed = (b: Buffer) => {
    if (!csettings.jpeg)
        return false;
    try {
        const res = f5inst.extract(b);
        if (!res)
            return; // unsure
        if (res.length > 1024) // probably garbage, allows for ~20 links from take-me-to.space, should be enough
            return; // unsure
        const str = Buffer.from(res).toString();
        if (!str.match(/^[a-zA-Z0-9:/.\-_ ]+$/))
            return; // unsure
        return str; // sure
    } catch {
        return; // unsure
    }
};

const extract = (b: Buffer, ex: string) => {
    // if we reached here then ex is heckin cute and valid
    return decodeCoom3Payload(Buffer.from(ex));
};

export default {
    match: fn => !!fn.match(/\.jpe?g$/),
    has_embed,
    extract,
    inject
} as ImageProcessor;