Redact Sensitive Data From a Document

import { readFileSync, writeFileSync } from "node:fs";
import { createCanvas } from "@napi-rs/canvas";
import { PDFDocument } from "pdf-lib";
import * as pdfjsLib from "pdfjs-dist/legacy/build/pdf.mjs";

const API_KEY = process.env.UNSILOED_API_KEY;
const BASE_URL = "https://prod.visionapi.unsiloed.ai";

// The kinds of information we want gone. Each description tells the extractor
// what to look for, so adapt these fields to your own document.
const SCHEMA = {
  type: "object",
  properties: {
    full_name:      { type: "string", description: "The full name of the account holder" },
    account_number: { type: "string", description: "The bank account number" },
    email:          { type: "string", description: "The account holder's email address" },
    phone:          { type: "string", description: "The account holder's phone number" },
  },
  required: ["full_name", "account_number", "email", "phone"],
  additionalProperties: false,
};

async function waitFor(url) {
  for (let i = 0; i < 90; i++) { // roughly 6 minutes at 4 seconds per poll
    const job = await (await fetch(url, { headers: { "api-key": API_KEY } })).json();
    if (["completed", "review", "Succeeded"].includes(job.status)) return job;
    // The extractor reports its reason in "error"; the parser uses "message".
    if (["failed", "Failed", "Cancelled"].includes(job.status)) throw new Error(job.error || job.message || `job ${job.status}`);
    await new Promise((r) => setTimeout(r, 4000));
  }
  throw new Error(`Job at ${url} did not finish in time`);
}

const bytes = readFileSync("document.pdf");

// Step 1: ask the extractor what the sensitive values are.
const extractForm = new FormData();
extractForm.append("pdf_file", new Blob([bytes], { type: "application/pdf" }), "document.pdf");
extractForm.append("schema_data", JSON.stringify(SCHEMA));
extractForm.append("model", "gamma");
const extractRes = await fetch(`${BASE_URL}/v2/extract`, { method: "POST", headers: { "api-key": API_KEY }, body: extractForm });
if (!extractRes.ok) throw new Error(`extract submit failed: HTTP ${extractRes.status} ${await extractRes.text()}`);
const extractJob = await extractRes.json();
const result = (await waitFor(`${BASE_URL}/extract/${extractJob.job_id}`)).result;
const sensitive = [];
for (const field of Object.keys(SCHEMA.properties)) {
  const value = result[field].value;
  if (value == null) {
    console.log(`Field '${field}' not found in the document, skipping.`);
    continue;
  }
  sensitive.push(value);
}
console.log("Will redact:", sensitive);

// Step 2: parse to get every word with a page-relative box.
const parseForm = new FormData();
parseForm.append("file", new Blob([bytes], { type: "application/pdf" }), "document.pdf");
const parseRes = await fetch(`${BASE_URL}/parse`, { method: "POST", headers: { "api-key": API_KEY }, body: parseForm });
if (!parseRes.ok) throw new Error(`parse submit failed: HTTP ${parseRes.status} ${await parseRes.text()}`);
const parseJob = await parseRes.json();
const parsed = await waitFor(`${BASE_URL}/parse/${parseJob.job_id}`);

const normalize = (text) => String(text).toLowerCase().replace(/[^a-z0-9]/g, "");
const words = [];
for (const chunk of parsed.chunks) {
  for (const seg of chunk.segments) {
    const { left: sx, top: sy } = seg.bbox;
    for (const word of seg.ocr || []) {
      const norm = normalize(word.text);
      if (!norm) continue;
      const b = word.bbox;
      words.push({
        norm,
        page: seg.page_number,
        pageWidth: seg.page_width,
        x: sx + b.left,
        width: b.width,
        // y + height lands on the baseline reliably; the height alone does not.
        baseline: sy + b.top + b.height,
      });
    }
  }
}

// With no words there is nothing to match, and the script would "succeed"
// by writing an unredacted copy. Fail loudly instead.
if (words.length === 0) throw new Error("Parse returned no OCR words; refusing to write an unredacted PDF");

// Step 3: find every occurrence by normalized concatenation, so a value matches
// however OCR split it (hyphens, "@") or wrapped across lines.
function findRuns(value) {
  const target = normalize(value);
  const runs = [];
  if (target.length < 3) return runs;
  for (let i = 0; i < words.length; i++) {
    let acc = "";
    for (let j = i; j < words.length && words[j].page === words[i].page && j - i < 10; j++) {
      acc += words[j].norm;
      if (acc === target) { runs.push(words.slice(i, j + 1)); break; }
      if (!target.startsWith(acc)) break;
    }
  }
  return runs;
}

// Step 4: turn each run into one box per line, anchored to the baseline.
const BAND_ABOVE = 19, BAND_BELOW = 6, PAD = 2;
function lineBoxes(run) {
  const lines = [];
  let current = [run[0]];
  for (const word of run.slice(1)) {
    const sameLine = Math.abs(word.baseline - current[current.length - 1].baseline) <= 12;
    if (sameLine && word.x >= current[current.length - 1].x - 2) current.push(word);
    else { lines.push(current); current = [word]; }
  }
  lines.push(current);
  return lines.map((line) => {
    const baseline = Math.max(...line.map((w) => w.baseline));
    return [
      Math.min(...line.map((w) => w.x)) - PAD,
      baseline - BAND_ABOVE,
      Math.max(...line.map((w) => w.x + w.width)) + PAD,
      baseline + BAND_BELOW,
    ];
  });
}

const boxesByPage = new Map();
for (const value of sensitive) {
  for (const run of findRuns(value)) {
    const page = run[0].page;
    if (!boxesByPage.has(page)) boxesByPage.set(page, []);
    boxesByPage.get(page).push(...lineBoxes(run));
  }
}

// Step 5: render each page to an image, paint the boxes, and rebuild the PDF.
// Flattening to an image removes the text for good; a drawn box alone would not.
const pageWidths = new Map(words.map((w) => [w.page, w.pageWidth]));
const fonts = "./node_modules/pdfjs-dist/standard_fonts/";
const pdf = await pdfjsLib.getDocument({ data: new Uint8Array(bytes), standardFontDataUrl: fonts }).promise;
const out = await PDFDocument.create();
for (let pageNo = 1; pageNo <= pdf.numPages; pageNo++) {
  const page = await pdf.getPage(pageNo);
  const base = page.getViewport({ scale: 1.0 }); // page size in PDF points
  // Render at whatever scale matches parse's pixel grid for this page.
  const scale = (pageWidths.get(pageNo) ?? base.width * 2) / base.width;
  const viewport = page.getViewport({ scale });
  const canvas = createCanvas(viewport.width, viewport.height);
  const ctx = canvas.getContext("2d");
  await page.render({ canvasContext: ctx, viewport, canvas }).promise;
  ctx.fillStyle = "#000";
  for (const [x0, y0, x1, y1] of boxesByPage.get(pageNo) || []) {
    ctx.fillRect(x0, y0, x1 - x0, y1 - y0);
  }
  const png = await out.embedPng(canvas.toBuffer("image/png"));
  const outPage = out.addPage([base.width, base.height]);
  outPage.drawImage(png, { x: 0, y: 0, width: base.width, height: base.height });
}
writeFileSync("redacted.pdf", await out.save());
console.log("Wrote redacted.pdf");