PEE/src/pngv3.ts

import { Buffer } from "buffer";
import type { EmbeddedFile, ImageProcessor } from "./main";
import { PNGDecoder, PNGEncoder } from "./png";
import { decodeCoom3Payload } from "./utils";
import { settings } from "./stores";
import { filehosts } from "./filehosts";

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

settings.subscribe(b => {
    csettings = b;
});
const CUM3 = Buffer.from("doo\0" + "m");
const CUM4 = Buffer.from("voo\0" + "m");
const CUM5 = Buffer.from("boo\0");
const CUM6 = Buffer.from("Creation Time\0");

const BufferReadStream = (b: Buffer) => {
    const ret = new ReadableStream<Buffer>({
        pull(cont) {
            cont.enqueue(b);
            cont.close();
        }
    });
    return ret;
};

const password = Buffer.from("NOA");

const xor = (a: Buffer, p: Buffer) => {
    let n = 0;
    for (let i = 0; i < a.byteLength; ++i) {
        a[i] ^= p[n];
        n++;
        n %= p.byteLength;
    }
};

const prefs: any = {
    'files.catbox.moe': 'c',
    'a.pomf.cat': 'p',
    'take-me-to.space': 't',
    'z.zz.fo': 'z'
};

const rprefs: any = {
    'c': 'files.catbox.moe',
    'p': 'a.pomf.cat',
    't': 'take-me-to.space',
    'z': 'z.zz.fo',
};

const extract = async (png: Buffer) => {
    const reader = BufferReadStream(png).getReader();
    const sneed = new PNGDecoder(reader);
    const ret: EmbeddedFile[] = [];

    try {
        for await (const [name, chunk, crc, offset] of sneed.chunks()) {
            let buff: Buffer;
            switch (name) {
                // should exist at the beginning of file to signal decoders if the file indeed has an embedded chunk
                case 'tEXt':
                    buff = await chunk();
                    if (buff.slice(4, 4 + CUM3.length).equals(CUM3)) {
                        const k = await decodeCoom3Payload(buff.slice(4 + CUM3.length));
                        ret.push(...k.filter(e => e).map(e => e as EmbeddedFile));
                    }
                    if (buff.slice(4, 4 + CUM4.length).equals(CUM4)) {
                        const passed = buff.slice(4 + CUM4.length);
                        xor(passed, password);
                        const k = await decodeCoom3Payload(passed);
                        ret.push(...k.filter(e => e).map(e => e as EmbeddedFile));
                    }
                    if (buff.slice(4, 4 + CUM5.length).equals(CUM5)) {
                        const passed = buff.slice(4 + CUM5.length);
                        const decoded = Buffer.from(passed.toString(), 'base64').toString().split(' ').map(e => {
                            return `https://${rprefs[e[0]]}/${e.slice(1)}`;
                        }).join(' ');
                        const k = await decodeCoom3Payload(Buffer.from(decoded));
                        ret.push(...k.filter(e => e).map(e => e as EmbeddedFile));
                    }
                    if (buff.slice(4, 4 + CUM6.length).equals(CUM6)) {
                        const passed = buff.slice(4 + CUM6.length);
                        if (!passed.toString().match(/^[0-9a-zA-Z+/=]+$/g)) continue;
                        try {
                            const decoded = Buffer
                                .from(passed.toString(), 'base64')
                                .toString()
                                .split(' ')
                                .map(e => {
                                    if (!(e[0] in rprefs))
                                        throw "Uhh";
                                    // should also check if the id has a len of 6-8 or ends in .pee
                                    return `https://${rprefs[e[0]]}/${e.slice(1)}`;
                                }).join(' ');
                            const k = await decodeCoom3Payload(Buffer.from(decoded));
                            ret.push(...k.filter(e => e).map(e => e as EmbeddedFile));
                        } finally {
                            //
                        }
                    }
                    break;
                case 'IDAT':
                // eslint-disable-next-line no-fallthrough
                case 'IEND':
                    return ret.slice(0, csettings.maxe);
                // eslint-disable-next-line no-fallthrough
                default:
                    break;
            }
        }
    } catch (e) {
        console.error(e);
    } finally {
        reader.releaseLock();
    }
};

const buildChunk = (tag: string, data: Buffer) => {
    const ret = Buffer.alloc(data.byteLength + 4);
    ret.write(tag.slice(0, 4), 0);
    data.copy(ret, 4);
    return ret;
};

export const BufferWriteStream = () => {
    let b = Buffer.from([]);
    const ret = new WritableStream<Buffer>({
        write(chunk) {
            b = Buffer.concat([b, chunk]);
        }
    });
    return [ret, () => b] as [WritableStream<Buffer>, () => Buffer];
};

export const inject_data = async (container: File, injb: Buffer) => {
    let magic = false;
    const [writestream, extract] = BufferWriteStream();
    const encoder = new PNGEncoder(writestream);
    const decoder = new PNGDecoder(container.stream().getReader());

    for await (const [name, chunk, crc, offset] of decoder.chunks()) {
        if (magic && name != "IDAT")
            break;
        if (!magic && name == "IDAT") {
            const passed = Buffer.from(injb);
            //xor(passed, password2);
            await encoder.insertchunk(["tEXt", async () => buildChunk("tEXt", Buffer.concat([CUM6, passed])), () => Promise.resolve(0), 0]);
            magic = true;
        }
        await encoder.insertchunk([name, chunk, crc, offset]);
    }
    await encoder.insertchunk(["IEND",
        async () => Promise.resolve(buildChunk("IEND", Buffer.from([]))),
        async () => Promise.resolve(0),
        0]);
    return extract();

};

const inject = async (container: File, links: string[]) => {
    links = links.map(link => {
        for (const h of filehosts) {
            if (link.includes(h.serving)) {
                const end = link.split('/').slice(-1)[0];
                return `${prefs[h.serving]}${end}`;
            }
        }
        return '';
    });
    const injb = Buffer.from(Buffer.from(links.join(' ')).toString("base64"));
    return inject_data(container, injb);
};

const has_embed = async (png: Buffer) => {
    const reader = BufferReadStream(png).getReader();
    const sneed = new PNGDecoder(reader);
    try {
        for await (const [name, chunk, crc, offset] of sneed.chunks()) {
            let buff: Buffer;
            switch (name) {
                // should exist at the beginning of file to signal decoders if the file indeed has an embedded chunk
                case 'tEXt':
                    buff = await chunk();
                    if (buff.slice(4, 4 + CUM3.length).equals(CUM3))
                        return true;
                    if (buff.slice(4, 4 + CUM4.length).equals(CUM4))
                        return true;
                    if (buff.slice(4, 4 + CUM5.length).equals(CUM5))
                        return true;
                    if (buff.slice(4, 4 + CUM6.length).equals(CUM6)) {
                        const passed = buff.slice(4 + CUM6.length).toString();
                        if (passed.match(/^[0-9a-zA-Z+/=]+$/g)) {
                            if (Buffer.from(passed, "base64").toString().split(" ").every(l => l[0] in rprefs))
                                return true;
                        }
                    }
                    break;
                case 'IDAT':
                // eslint-disable-next-line no-fallthrough
                case 'IEND':
                    return false; // Didn't find tExt Chunk; Definite no
                // eslint-disable-next-line no-fallthrough
                default:
                    break;
            }
        }
        // stream ended on chunk boundary, so no unexpected EOF was fired, need more data anyway
    } catch (e) {
        return; // possibly unexpected EOF, need more data to decide
    } finally {
        reader.releaseLock();
    }
};

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