lmexif-sniffer

import * as fs from "node:fs";
import { deflateSync } from "node:zlib";

import { PNG_SIGNATURE } from "./pngText";

const STRIP_KEYWORDS = new Set(["chara", "ccv3", "character"]);

function crc32Png(buf: Buffer): number {
  let c = 0xffffffff;
  for (let i = 0; i < buf.length; i++) {
    c ^= buf[i];
    for (let j = 0; j < 8; j++) {
      c = (c >>> 1) ^ (0xedb88320 & -(c & 1));
    }
  }
  return (c ^ 0xffffffff) >>> 0;
}

function encodePngChunk(type: string, data: Buffer): Buffer {
  const typeBuf = Buffer.from(type, "ascii");
  if (typeBuf.length !== 4) throw new Error(`Invalid PNG chunk type: ${type}`);
  const len = data.length;
  const crc = crc32Png(Buffer.concat([typeBuf, data]));
  const out = Buffer.alloc(4 + 4 + len + 4);
  out.writeUInt32BE(len, 0);
  typeBuf.copy(out, 4);
  data.copy(out, 8);
  out.writeUInt32BE(crc, 8 + len);
  return out;
}

function keywordFromTextChunkData(data: Buffer): string | null {
  const z = data.indexOf(0);
  if (z < 0) return null;
  return data.subarray(0, z).toString("latin1");
}

function shouldStripChunk(type: string, data: Buffer): boolean {
  if (type !== "tEXt" && type !== "zTXt") return false;
  const kw = keywordFromTextChunkData(data);
  return kw != null && STRIP_KEYWORDS.has(kw);
}

type PngChunk = { type: string; data: Buffer };

function readPngChunks(buf: Buffer): PngChunk[] {
  if (buf.length < 8 || !buf.subarray(0, 8).equals(PNG_SIGNATURE)) {
    throw new Error("Not a PNG file");
  }
  const out: PngChunk[] = [];
  let pos = 8;
  while (pos + 8 <= buf.length) {
    const length = buf.readUInt32BE(pos);
    const type = buf.subarray(pos + 4, pos + 8).toString("ascii");
    const chunkStart = pos + 8;
    const chunkEnd = chunkStart + length;
    if (chunkEnd + 4 > buf.length) throw new Error("Truncated PNG");
    const data = buf.subarray(chunkStart, chunkEnd);
    pos = chunkEnd + 4;
    out.push({ type, data });
    if (type === "IEND") break;
  }
  if (out.length === 0 || out[out.length - 1]!.type !== "IEND") {
    throw new Error("PNG missing IEND chunk");
  }
  const ihdr = out.find((c) => c.type === "IHDR");
  if (!ihdr) throw new Error("PNG missing IHDR chunk");
  return out;
}

function cardJsonUtf8(card: string | Record<string, unknown>): string {
  if (typeof card === "string") {
    JSON.parse(card);
    return card;
  }
  return JSON.stringify(card);
}

export type PngCharacterKeyword = "chara" | "ccv3";

export type EmbedCharacterCardOptions = {
  keyword: PngCharacterKeyword;
  /** If true, write zTXt (zlib); otherwise tEXt. */
  compress?: boolean;
};

/**
 * Strip chara/ccv3/character text chunks, embed a new card chunk immediately after IHDR.
 */
export function embedCharacterCardInPngBuffer(
  pngBytes: Buffer,
  card: string | Record<string, unknown>,
  options: EmbedCharacterCardOptions,
): Buffer {
  const keyword = options.keyword;
  if (keyword.length > 79) throw new Error("PNG keyword must be at most 79 bytes");
  const jsonUtf8 = cardJsonUtf8(card);
  const b64 = Buffer.from(jsonUtf8, "utf8").toString("base64");
  const compress = Boolean(options.compress);
  let textChunkData: Buffer;
  let chunkType: "tEXt" | "zTXt";
  if (compress) {
    chunkType = "zTXt";
    const compressed = deflateSync(Buffer.from(b64, "utf8"));
    textChunkData = Buffer.concat([Buffer.from(`${keyword}\x00\x00`, "binary"), compressed]);
  } else {
    chunkType = "tEXt";
    textChunkData = Buffer.concat([Buffer.from(`${keyword}\x00`, "binary"), Buffer.from(b64, "latin1")]);
  }

  const newChunk = encodePngChunk(chunkType, textChunkData);
  const chunks = readPngChunks(pngBytes);
  const filtered = chunks.filter((c) => !shouldStripChunk(c.type, c.data));

  const pieces: Buffer[] = [PNG_SIGNATURE];
  let inserted = false;
  for (const c of filtered) {
    pieces.push(encodePngChunk(c.type, c.data));
    if (!inserted && c.type === "IHDR") {
      pieces.push(newChunk);
      inserted = true;
    }
  }
  if (!inserted) throw new Error("IHDR not found after filter (invalid PNG)");
  return Buffer.concat(pieces);
}

export function embedCharacterCardInPngFile(
  filePath: string,
  card: string | Record<string, unknown>,
  options: EmbedCharacterCardOptions,
): Record<string, unknown> {
  if (!fs.existsSync(filePath) || !fs.statSync(filePath).isFile()) {
    throw new Error(`Not a file: ${filePath}`);
  }
  const buf = fs.readFileSync(filePath);
  const next = embedCharacterCardInPngBuffer(buf, card, options);
  fs.writeFileSync(filePath, next);
  return {
    path: fs.realpathSync(filePath),
    keyword: options.keyword,
    compress: Boolean(options.compress),
    bytes_written: next.length,
  };
}