Spaces:

Wardline
/

cloak-api

Paused

cloak-api / docclass.py
Update docclass.py
ba2fd7e verified 14 days ago
78.1 kB
"""
Document classification — the second path of this Space, in one file.

    POST /file/classify   {"text": "...", "column_headers": [...]}
    POST /file/analyze    raw file bytes (Content-Type: application/octet-stream,
                          X-Filename: <original name>)

`/detect` (in app.py) answers "which personal details does this text contain?".
This answers "what KIND of document is this file?" — a contract, a payslip, a
passport scan, a .env full of keys — and suggests what policy should do with
it. The two are independent signals; CLOAK combines them.

The handover's modules, folded into this file in dependency order:

    Layer 1      real file type from the bytes, not the name
    Layer 2      zip / gzip / 7z unpacking, bounded
    Extraction   pdf · docx · xlsx · csv · text · OCR for images
    Layer 3      Layer3DocType: classifier + abstain guard + credentials escalation
    Policy table class -> suggested action (configuration, not ML)
    API          model loading, limits, the /file routes

Everything here is defensive about the PII service it shares a process with:
if this model cannot be fetched, fails its checksum or fails to load, the
/file paths answer 503 and /detect carries on untouched.

ALWAYS go through Layer3DocType, never the raw bundle — the abstain guard and
the credentials escalation rule live in Layer3DocType, not in the pickle.
"""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable
import csv
import gzip
import hashlib
import io
import math
import os
import re
import tempfile
import threading
import time
import zipfile

import joblib
import numpy as np
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel, Field

# ════════════════════════════════════════════════════════════════════════════
# Layer 1 · entropy
# ════════════════════════════════════════════════════════════════════════════
# Shannon entropy check.
#
# Used when no known magic byte signature matches at all. Real text and
# most structured document formats have a distinctly non-random byte
# distribution; encrypted or already-compressed data looks close to
# statistically random. This doesn't identify a format — it distinguishes
# "unrecognised" from "probably encrypted/packed", which is itself a
# finding worth raising rather than silently discarding.

# Empirically, genuine encrypted/compressed data usually lands above ~7.5
# bits/byte (max possible is 8.0, for a perfectly uniform byte distribution).
# Ordinary text sits well below this; most already-compressed formats we
# recognise by magic bytes (zip, gzip, png, jpeg) never reach this check
# because they matched a known signature first.
HIGH_ENTROPY_THRESHOLD = 7.5


def shannon_entropy(data: bytes) -> float:
    if not data:
        return 0.0
    counts = Counter(data)
    length = len(data)
    return -sum(
        (count / length) * math.log2(count / length)
        for count in counts.values()
    )


def looks_encrypted_or_packed(data: bytes, sample_size: int = 65536) -> tuple[bool, float]:
    """Sample the first `sample_size` bytes (entropy over a large file
    converges quickly; no need to hash the whole thing) and report
    whether it looks encrypted/packed alongside the raw entropy value."""
    sample = data[:sample_size]
    ent = shannon_entropy(sample)
    return ent >= HIGH_ENTROPY_THRESHOLD, ent


# ════════════════════════════════════════════════════════════════════════════
# Layer 1 · magic-byte signatures
# ════════════════════════════════════════════════════════════════════════════
# Magic-byte signature table for Layer 1 (True Format).
#
# Each entry maps a format name to one or more byte signatures that must
# appear at a fixed offset from the start of the file. This is the same
# approach libmagic/the Unix `file` command uses, kept in pure Python here
# so this runs identically on the server and the Windows agent with no
# external binary dependency.
#
# Extend this table (or better: load customer additions from the
# FileTypeSignature table described in the architecture doc) rather than
# hardcoding new formats elsewhere in the codebase.

@dataclass(frozen=True)
class Signature:
    format_name: str
    magic: bytes
    offset: int = 0
    # Extensions commonly (honestly) used for this format, for mismatch checks.
    common_extensions: tuple = ()


SIGNATURES: list[Signature] = [
    Signature("pdf", b"%PDF-", 0, (".pdf",)),
    # ZIP and every OOXML/Java/Android format built on top of it share this
    # magic. Layer 1 can only tell you "this is a zip container" — telling
    # docx apart from xlsx apart from a plain zip requires opening it and
    # checking internal structure, done in _refine_zip_based below.
    Signature("zip", b"PK\x03\x04", 0, (".zip", ".docx", ".xlsx", ".pptx", ".jar", ".apk")),
    Signature("zip_empty", b"PK\x05\x06", 0, (".zip",)),  # empty archive
    Signature("zip_spanned", b"PK\x07\x08", 0, (".zip",)),  # spanned archive
    Signature("rar", b"Rar!\x1a\x07\x00", 0, (".rar",)),
    Signature("rar5", b"Rar!\x1a\x07\x01\x00", 0, (".rar",)),
    Signature("7z", b"7z\xbc\xaf\x27\x1c", 0, (".7z",)),
    Signature("gzip", b"\x1f\x8b", 0, (".gz", ".gzip", ".tgz")),
    Signature("png", b"\x89PNG\r\n\x1a\n", 0, (".png",)),
    Signature("jpeg", b"\xff\xd8\xff", 0, (".jpg", ".jpeg")),
    Signature("gif87a", b"GIF87a", 0, (".gif",)),
    Signature("gif89a", b"GIF89a", 0, (".gif",)),
    Signature("bmp", b"BM", 0, (".bmp",)),
    Signature("tiff_le", b"II*\x00", 0, (".tif", ".tiff")),
    Signature("tiff_be", b"MM\x00*", 0, (".tif", ".tiff")),
    # Legacy OLE2 container format — old-style .doc/.xls/.ppt (pre-2007).
    Signature("ole2", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", 0, (".doc", ".xls", ".ppt", ".msg")),
    Signature("sqlite", b"SQLite format 3\x00", 0, (".db", ".sqlite", ".sqlite3")),
    Signature("rtf", b"{\\rtf1", 0, (".rtf",)),
]

# Sorted longest-magic-first so a more specific signature is checked before
# a shorter one that could also match the same bytes.
SIGNATURES.sort(key=lambda s: len(s.magic), reverse=True)


def match_signature(head: bytes) -> Signature | None:
    """Return the first signature whose magic bytes match the start of `head`."""
    for sig in SIGNATURES:
        end = sig.offset + len(sig.magic)
        if len(head) >= end and head[sig.offset:end] == sig.magic:
            return sig
    return None


# ════════════════════════════════════════════════════════════════════════════
# Layer 1 · structure validators
# ════════════════════════════════════════════════════════════════════════════
# Structure validation for Layer 1.
#
# Matching magic bytes only proves the first few bytes look right. A file
# can have a genuine PDF header and still be truncated, broken, or have a
# fake header pasted onto unrelated content. This module opens the file
# far enough to check the structural invariants each format actually
# requires, and is where "header valid but structure won't parse" (treated
# as concealment, not corruption, per the architecture doc) gets detected.

@dataclass
class StructureResult:
    valid: bool
    refined_type: str | None = None       # e.g. "docx" instead of just "zip"
    notes: list[str] = field(default_factory=list)


# OOXML formats are all zip containers that must include this entry, and
# are then distinguished by which top-level folder they carry.
_OOXML_MARKERS = {
    "word/document.xml": "docx",
    "xl/workbook.xml": "xlsx",
    "ppt/presentation.xml": "pptx",
}


def validate_pdf(data: bytes) -> StructureResult:
    """A real PDF has a trailer and an xref table telling a reader where
    every object lives. Missing/unparseable trailer with a valid header
    is the classic "fake header on real different content" pattern."""
    if b"%PDF-" not in data[:1024]:
        return StructureResult(valid=False, notes=["magic bytes present but header not found on re-check"])

    tail = data[-2048:] if len(data) > 2048 else data
    has_eof = b"%%EOF" in tail
    has_startxref = b"startxref" in tail
    has_trailer = b"trailer" in tail or b"/Root" in tail  # cross-ref streams (PDF 1.5+) skip the trailer keyword

    if has_eof and has_startxref and has_trailer:
        return StructureResult(valid=True, refined_type="pdf")

    notes = []
    if not has_eof:
        notes.append("no %%EOF marker near end of file")
    if not has_startxref:
        notes.append("no startxref marker near end of file")
    if not has_trailer:
        notes.append("no trailer or /Root reference near end of file")
    return StructureResult(valid=False, notes=notes)


def validate_zip_based(data: bytes) -> StructureResult:
    """Open the zip and check both that it's structurally sound, and -
    for the OOXML family - which specific document type it actually is.
    A file claiming .docx that is a valid zip but lacks word/document.xml
    is not a genuine docx, whatever the extension says."""
    try:
        zf = zipfile.ZipFile(io.BytesIO(data))
    except zipfile.BadZipFile as exc:
        return StructureResult(valid=False, notes=[f"not a valid zip container: {exc}"])

    notes: list[str] = []
    has_encrypted_entries = any(info.flag_bits & 0x1 for info in zf.infolist())

    if has_encrypted_entries:
        # zf.testzip() opens every entry to check its CRC, which raises
        # RuntimeError on any password-protected entry with no password
        # supplied. An encrypted archive is a valid zip and a finding in
        # its own right (see Layer 2) - it must never crash Layer 1's
        # structure check. Skip the CRC pass entirely in this case; Layer
        # 2 handles encrypted entries explicitly when it unpacks.
        notes.append("archive contains password-protected entries - integrity of those entries not verified")
    else:
        bad_entry = zf.testzip()
        if bad_entry is not None:
            return StructureResult(valid=False, notes=[f"corrupt entry inside archive: {bad_entry}"])

    names = set(zf.namelist())

    if "[Content_Types].xml" not in names:
        # Genuine zip, just not an OOXML document (a plain .zip, a .jar, an .apk, etc).
        return StructureResult(valid=True, refined_type="zip", notes=notes + ["no [Content_Types].xml - not an OOXML document"])

    for marker, kind in _OOXML_MARKERS.items():
        if marker in names:
            return StructureResult(valid=True, refined_type=kind, notes=notes)

    return StructureResult(
        valid=True,
        refined_type="ooxml_unknown",
        notes=notes + ["has [Content_Types].xml but no recognised word/xl/ppt entry"],
    )


def validate_ole2(data: bytes) -> StructureResult:
    """Legacy pre-2007 Office container (.doc/.xls/.ppt). Full OLE2
    directory parsing is a project of its own; at Layer 1 we only
    confirm the container is intact enough to be read at all, which is
    already enough to catch a fake-header/garbage-body mismatch."""
    if len(data) < 512:
        return StructureResult(valid=False, notes=["file shorter than one OLE2 sector (512 bytes) - truncated"])
    return StructureResult(valid=True, refined_type="ole2")


# Dispatch table: which validator applies to which Layer-1-matched format.
VALIDATORS = {
    "pdf": validate_pdf,
    "zip": validate_zip_based,
    "ole2": validate_ole2,
}


def validate_structure(format_name: str, data: bytes) -> StructureResult:
    validator = VALIDATORS.get(format_name)
    if validator is None:
        # No structural invariant defined yet for this format - magic
        # bytes matching is all we can assert. Not a failure, just a gap
        # to close as more formats get their own validator.
        return StructureResult(valid=True, refined_type=format_name, notes=["no structure validator implemented - magic-byte match only"])
    return validator(data)


# ════════════════════════════════════════════════════════════════════════════
# Layer 1 · identify
# ════════════════════════════════════════════════════════════════════════════
# Layer 1 entry point: True Format identification.
#
# Combines magic-byte matching, structure validation, and the entropy
# fallback into a single call. This is what Layer 2 (container unpacking)
# calls on every archive entry, and what the detection pipeline calls
# before a parser is chosen — the whole point being that parser selection
# must never again be driven by the filename extension alone.

# Extension -> the type name we expect structure validation to confirm.
# Deliberately separate from Signature.common_extensions: several
# extensions (.docx/.xlsx/.pptx/.zip) share one magic-byte signature and
# can only be told apart after structure validation opens the container.
EXTENSION_TYPE_MAP = {
    ".pdf": "pdf",
    ".docx": "docx",
    ".xlsx": "xlsx",
    ".pptx": "pptx",
    ".zip": "zip",
    ".jar": "zip",
    ".apk": "zip",
    ".rar": "rar",
    ".7z": "7z",
    ".gz": "gzip",
    ".gzip": "gzip",
    ".tgz": "gzip",
    ".png": "png",
    ".jpg": "jpeg",
    ".jpeg": "jpeg",
    ".gif": "gif",  # covers both gif87a/gif89a signatures
    ".bmp": "bmp",
    ".tif": "tiff",
    ".tiff": "tiff",
    ".doc": "ole2",
    ".xls": "ole2",
    ".ppt": "ole2",
    ".msg": "ole2",
    ".db": "sqlite",
    ".sqlite": "sqlite",
    ".sqlite3": "sqlite",
    ".rtf": "rtf",
}

# gif87a/gif89a are two magic variants of one logical format; normalise both
# refined-type outputs to "gif" so they compare cleanly against the extension map.
_TYPE_ALIASES = {
    "gif87a": "gif",
    "gif89a": "gif",
}


@dataclass
class FileIdentity:
    declared_type: str | None          # from the filename/extension — a claim, not a fact
    actual_type: str                   # from magic bytes + structure validation
    type_mismatch: bool
    structure_valid: bool | None       # None if no validator exists yet for this format
    entropy: float | None = None       # only set when no magic signature matched at all
    likely_encrypted_or_packed: bool = False
    notes: list[str] = field(default_factory=list)


def _declared_type_from_name(filename: str) -> str | None:
    ext = os.path.splitext(filename)[1].lower()
    return EXTENSION_TYPE_MAP.get(ext)


def identify_bytes(data: bytes, filename: str = "") -> FileIdentity:
    declared_type = _declared_type_from_name(filename) if filename else None
    notes: list[str] = []

    sig = match_signature(data[:64])

    if sig is None:
        # No known magic bytes at all. Either a format we haven't added a
        # signature for yet, or something deliberately packed/encrypted.
        encrypted_like, ent = looks_encrypted_or_packed(data)
        notes.append("no known magic-byte signature matched")
        actual_type = "unknown"
        mismatch = declared_type is not None  # claiming a known type but bytes say otherwise
        return FileIdentity(
            declared_type=declared_type,
            actual_type=actual_type,
            type_mismatch=mismatch,
            structure_valid=None,
            entropy=ent,
            likely_encrypted_or_packed=encrypted_like,
            notes=notes,
        )

    structure = validate_structure(sig.format_name, data)
    actual_type = structure.refined_type or sig.format_name
    actual_type = _TYPE_ALIASES.get(actual_type, actual_type)
    notes.extend(structure.notes)

    mismatch = declared_type is not None and declared_type != actual_type

    return FileIdentity(
        declared_type=declared_type,
        actual_type=actual_type,
        type_mismatch=mismatch,
        structure_valid=structure.valid,
        notes=notes,
    )


def identify_file(path: str) -> FileIdentity:
    with open(path, "rb") as f:
        data = f.read()
    return identify_bytes(data, filename=os.path.basename(path))


# ════════════════════════════════════════════════════════════════════════════
# Layer 2 · container unpacking
# ════════════════════════════════════════════════════════════════════════════
# Layer 2: container unpacking.
#
# Layer 1 tells you a file really is a genuine zip. That's a box - the
# sensitive content is what's inside it. This module opens containers and
# runs Layer 1 again on every entry, recursively, with hard caps on depth
# and expansion ratio so a small malicious (or just unlucky) archive can't
# turn the scanner into its own denial-of-service target (a "zip bomb").
#
# An entry this module cannot open - because it's password-protected, or
# because it's a format we don't extract yet (RAR/7z) - is recorded as a
# finding in its own right, never silently skipped as "clean."

try:
    import py7zr
    _HAVE_PY7ZR = True
except ImportError:
    _HAVE_PY7ZR = False


@dataclass
class ScanLimits:
    max_depth: int = 5
    # Effectively unbounded ratio for tiny files can still explode in
    # absolute terms, hence both a ratio cap and an absolute size cap below.
    max_expansion_ratio: float = 100.0
    max_total_uncompressed_bytes: int = 500 * 1024 * 1024  # 500 MB
    max_entries: int = 10_000


@dataclass
class EntryResult:
    path: str                      # slash-joined nested path, e.g. "outer.zip/photos/IMG_0421.jpg"
    depth: int
    size: int
    identity: FileIdentity | None  # None when the entry couldn't be opened/read at all
    encrypted: bool = False
    unsupported_format: bool = False
    notes: list[str] = field(default_factory=list)
    data: bytes | None = None      # only populated when scan_container(..., keep_bytes=True) - see the extraction section


@dataclass
class ScanResult:
    entries: list[EntryResult] = field(default_factory=list)
    truncated: bool = False
    truncation_reason: str | None = None
    total_uncompressed_bytes: int = 0


@dataclass
class _ScanState:
    entries: list[EntryResult] = field(default_factory=list)
    entry_count: int = 0
    total_uncompressed: int = 0
    truncated: bool = False
    truncation_reason: str | None = None
    keep_bytes: bool = False  # if True, EntryResult.data is populated for later extraction


_UNSUPPORTED_ARCHIVE_TYPES = {"rar"} if _HAVE_PY7ZR else {"rar", "7z"}


def _recurse_into(identity: FileIdentity, entry_data: bytes, entry_path: str, depth: int, state: "_ScanState", limits: "ScanLimits") -> None:
    """Shared dispatch: given an entry we've already identified, decide
    whether/how to keep unpacking it. Used after every extraction (zip,
    gzip, 7z) so the recursion rules only live in one place."""
    if depth >= limits.max_depth:
        state.truncated = True
        state.truncation_reason = f"max recursion depth {limits.max_depth} reached at {entry_path}"
        return

    if identity.actual_type == "zip":
        _scan_zip(entry_data, entry_path, depth + 1, state, limits)  # no filter once genuinely nested
    elif identity.actual_type in ("docx", "xlsx", "pptx"):
        _scan_zip(entry_data, entry_path, depth + 1, state, limits, entry_filter=_is_ooxml_embedded_object)
    elif identity.actual_type == "7z" and _HAVE_PY7ZR:
        _scan_7z(entry_data, entry_path, depth + 1, state, limits)
    elif identity.actual_type in _UNSUPPORTED_ARCHIVE_TYPES:
        state.entries[-1].notes.append(
            f"{identity.actual_type} extraction not implemented yet - contents could not be verified"
        )


def _is_ooxml_embedded_object(entry_name: str) -> bool:
    """True for the entries inside a docx/xlsx/pptx worth actually
    inspecting: embedded objects (another spreadsheet, another document)
    and embedded media (images - including a photographed ID card pasted
    into a Word doc). Everything else in an OOXML zip is boilerplate XML
    (styles, theme, relationships) with no security relevance."""
    lname = entry_name.lower()
    return "/embeddings/" in lname or "/media/" in lname


def _scan_zip(
    data: bytes,
    path: str,
    depth: int,
    state: _ScanState,
    limits: ScanLimits,
    entry_filter=None,
) -> None:
    try:
        zf = zipfile.ZipFile(io.BytesIO(data))
    except zipfile.BadZipFile as exc:
        state.entries.append(EntryResult(
            path=path, depth=depth, size=len(data), identity=None,
            notes=[f"failed to open as zip despite matching signature: {exc}"],
        ))
        return

    for info in zf.infolist():
        if state.truncated:
            return
        if info.is_dir():
            continue
        if entry_filter is not None and not entry_filter(info.filename):
            # Not worth surfacing (e.g. boilerplate OOXML XML parts) - skip
            # entirely, including from the entry-count/size caps, since we
            # never read these bytes at all.
            continue

        state.entry_count += 1
        if state.entry_count > limits.max_entries:
            state.truncated = True
            state.truncation_reason = f"entry count exceeded cap of {limits.max_entries}"
            return

        entry_path = f"{path}/{info.filename}"
        encrypted = bool(info.flag_bits & 0x1)

        state.total_uncompressed += info.file_size
        if state.total_uncompressed > limits.max_total_uncompressed_bytes:
            state.truncated = True
            state.truncation_reason = (
                f"total uncompressed size exceeded cap of {limits.max_total_uncompressed_bytes} bytes "
                f"at {entry_path} - possible zip bomb, extraction stopped"
            )
            return

        if info.compress_size > 0:
            ratio = info.file_size / info.compress_size
            if ratio > limits.max_expansion_ratio:
                state.entries.append(EntryResult(
                    path=entry_path, depth=depth, size=info.file_size, identity=None,
                    notes=[f"expansion ratio {ratio:.0f}x exceeds cap of {limits.max_expansion_ratio}x "
                           f"- not extracted, possible zip bomb"],
                ))
                continue

        if encrypted:
            # Per the architecture doc: an archive entry we cannot open is a
            # finding, not a failure. Do not attempt to brute-force it here.
            state.entries.append(EntryResult(
                path=entry_path, depth=depth, size=info.file_size, identity=None,
                encrypted=True,
                notes=["password-protected entry - contents could not be inspected; "
                       "the encryption itself is the finding"],
            ))
            continue

        try:
            entry_data = zf.read(info)
        except (RuntimeError, zipfile.BadZipFile, EOFError) as exc:
            state.entries.append(EntryResult(
                path=entry_path, depth=depth, size=info.file_size, identity=None,
                notes=[f"could not read entry: {exc}"],
            ))
            continue

        identity = identify_bytes(entry_data, filename=info.filename)
        state.entries.append(EntryResult(
            path=entry_path, depth=depth, size=info.file_size, identity=identity,
            unsupported_format=identity.actual_type in _UNSUPPORTED_ARCHIVE_TYPES,
            notes=list(identity.notes),
            data=entry_data if state.keep_bytes else None,
        ))

        _recurse_into(identity, entry_data, entry_path, depth, state, limits)
        if state.truncated:
            return


def _scan_gzip(data: bytes, path: str, depth: int, state: _ScanState, limits: ScanLimits) -> None:
    try:
        decompressed = gzip.decompress(data)
    except OSError as exc:
        state.entries.append(EntryResult(
            path=path, depth=depth, size=len(data), identity=None,
            notes=[f"failed to decompress gzip stream: {exc}"],
        ))
        return

    if len(decompressed) > limits.max_total_uncompressed_bytes:
        state.truncated = True
        state.truncation_reason = f"gzip stream expanded past cap of {limits.max_total_uncompressed_bytes} bytes"
        return

    # gzip carries the original filename in its header when available; fall
    # back to stripping the .gz suffix from the outer name otherwise.
    inner_name = path[:-3] if path.endswith((".gz", ".tgz")) else path
    identity = identify_bytes(decompressed, filename=inner_name)
    entry_path = f"{path}::decompressed"
    state.entries.append(EntryResult(
        path=entry_path, depth=depth, size=len(decompressed), identity=identity,
        notes=list(identity.notes),
        data=decompressed if state.keep_bytes else None,
    ))

    _recurse_into(identity, decompressed, entry_path, depth, state, limits)


def _scan_7z(data: bytes, path: str, depth: int, state: _ScanState, limits: ScanLimits) -> None:
    """7z support via py7zr - a pure-Python implementation, so this works
    without an external 7-Zip binary being installed on the host. Unlike
    zip, 7z typically compresses many entries together in one "solid"
    block, so there's no meaningful per-entry compressed size to compute
    an expansion ratio against - the total-uncompressed-size cap below is
    checked BEFORE extracting anything, which is the guard that matters
    for a 7z bomb."""
    try:
        archive = py7zr.SevenZipFile(io.BytesIO(data), mode="r")
    except Exception as exc:  # noqa: BLE001 - py7zr raises several distinct exception types
        state.entries.append(EntryResult(
            path=path, depth=depth, size=len(data), identity=None,
            notes=[f"failed to open as 7z despite matching signature: {exc}"],
        ))
        return

    if archive.needs_password():
        state.entries.append(EntryResult(
            path=path, depth=depth, size=len(data), identity=None, encrypted=True,
            notes=["password-protected 7z archive - contents could not be inspected; "
                   "the encryption itself is the finding"],
        ))
        return

    try:
        file_infos = [i for i in archive.list() if not i.is_directory]
    except Exception as exc:  # noqa: BLE001
        state.entries.append(EntryResult(
            path=path, depth=depth, size=len(data), identity=None,
            notes=[f"failed to list 7z contents: {exc}"],
        ))
        return

    if state.entry_count + len(file_infos) > limits.max_entries:
        state.truncated = True
        state.truncation_reason = f"7z entry count would exceed cap of {limits.max_entries}"
        return

    total_size = sum(i.uncompressed for i in file_infos)
    if state.total_uncompressed + total_size > limits.max_total_uncompressed_bytes:
        state.truncated = True
        state.truncation_reason = (
            f"7z total uncompressed size would exceed cap of {limits.max_total_uncompressed_bytes} bytes "
            f"at {path} - not extracted, possible archive bomb"
        )
        return

    # This py7zr version has no in-memory read API (only extractall to a
    # real path), so extraction goes through a temp directory that's
    # cleaned up before this function returns either way.
    try:
        with tempfile.TemporaryDirectory(prefix="cloak_7z_") as tmp_dir:
            archive.extractall(path=tmp_dir)
            for info in file_infos:
                if state.truncated:
                    return
                state.entry_count += 1
                state.total_uncompressed += info.uncompressed
                entry_path = f"{path}/{info.filename}"

                extracted_path = Path(tmp_dir) / info.filename
                if not extracted_path.exists():
                    state.entries.append(EntryResult(
                        path=entry_path, depth=depth, size=info.uncompressed, identity=None,
                        notes=["entry present in archive listing but missing after extraction"],
                    ))
                    continue

                entry_data = extracted_path.read_bytes()
                identity = identify_bytes(entry_data, filename=info.filename)
                state.entries.append(EntryResult(
                    path=entry_path, depth=depth, size=len(entry_data), identity=identity,
                    unsupported_format=identity.actual_type in _UNSUPPORTED_ARCHIVE_TYPES,
                    notes=list(identity.notes),
                    data=entry_data if state.keep_bytes else None,
                ))
                _recurse_into(identity, entry_data, entry_path, depth, state, limits)
    except Exception as exc:  # noqa: BLE001
        state.entries.append(EntryResult(
            path=path, depth=depth, size=len(data), identity=None,
            notes=[f"failed to extract 7z contents: {exc}"],
        ))
        return


def scan_container(
    data: bytes,
    filename: str,
    limits: ScanLimits | None = None,
    root_identity: FileIdentity | None = None,
    keep_bytes: bool = False,
) -> ScanResult:
    """Identify the top-level file (Layer 1), then, if it's a container
    format, recurse into it. For a non-container file this simply returns
    an empty entry list - callers should already have the top-level
    FileIdentity from Layer 1 and only call this when it's worth opening.

    Pass `root_identity` if the caller already ran Layer 1 on this data
    (e.g. via `analyze_file`) to avoid re-running magic-byte/structure
    checks on the same bytes a second time.

    Pass `keep_bytes=True` to populate `EntryResult.data` with each
    entry's raw extracted bytes, so a caller (e.g. the extraction step)
    can run extraction against entries found inside a container, not
    just the top-level file. Off by default since it roughly doubles
    memory use for large archives with many entries."""
    limits = limits or ScanLimits()
    root_identity = root_identity or identify_bytes(data, filename=filename)
    state = _ScanState(keep_bytes=keep_bytes)

    if root_identity.actual_type == "zip":
        _scan_zip(data, filename, depth=1, state=state, limits=limits)
    elif root_identity.actual_type in ("docx", "xlsx", "pptx"):
        # OOXML documents are zip containers too, but recursing into every
        # boilerplate XML part (styles.xml, theme1.xml, etc.) would just be
        # noise - only embedded objects/media are worth surfacing, since
        # that's where a customer_list.xlsx pasted inside a Word doc, or a
        # photographed ID card embedded as an image, would actually hide.
        _scan_zip(data, filename, depth=1, state=state, limits=limits, entry_filter=_is_ooxml_embedded_object)
    elif root_identity.actual_type == "gzip":
        _scan_gzip(data, filename, depth=1, state=state, limits=limits)
    elif root_identity.actual_type == "7z" and _HAVE_PY7ZR:
        _scan_7z(data, filename, depth=1, state=state, limits=limits)
    elif root_identity.actual_type in _UNSUPPORTED_ARCHIVE_TYPES:
        state.entries.append(EntryResult(
            path=filename, depth=0, size=len(data), identity=root_identity,
            unsupported_format=True,
            notes=[f"{root_identity.actual_type} extraction not implemented yet - contents could not be verified"],
        ))
    # else: not a container at all - nothing to unpack, caller already has
    # everything it needs from Layer 1's root_identity.

    return ScanResult(
        entries=state.entries,
        truncated=state.truncated,
        truncation_reason=state.truncation_reason,
        total_uncompressed_bytes=state.total_uncompressed,
    )


# ════════════════════════════════════════════════════════════════════════════
# Layers 1 + 2 · analyze
# ════════════════════════════════════════════════════════════════════════════
# Single combined entry point for Layers 1 + 2. Runs Layer 1 once and
# passes the result into Layer 2, instead of each layer independently
# re-running identify_bytes on the same top-level bytes (which is what
# inspect_file.py and run_test_files.py were doing before this existed).
#
# This is what the real detection pipeline, and Layer 3 later, should call
# - not identify_bytes/scan_container directly - so identity is computed
# exactly once per file no matter how many layers need to read it.

@dataclass
class FileAnalysis:
    identity: FileIdentity
    container: ScanResult


def analyze_bytes(
    data: bytes, filename: str, limits: ScanLimits | None = None, keep_bytes: bool = False
) -> FileAnalysis:
    identity = identify_bytes(data, filename=filename)
    container = scan_container(data, filename=filename, limits=limits, root_identity=identity, keep_bytes=keep_bytes)
    return FileAnalysis(identity=identity, container=container)


# ════════════════════════════════════════════════════════════════════════════
# Extraction · OCR
# ════════════════════════════════════════════════════════════════════════════
# OCR for image files - the gap the architecture doc calls out as
# "likely the highest-value missing detection in the Indian market"
# (§5, Layer 3): a photographed Aadhaar card has no extractable text at
# all, so every layer downstream sees nothing and the file passes as
# clean.
#
# Engine: RapidOCR (ONNX runtime). Chosen over Tesseract because:
#   - pure pip install, no system binary required (the Tesseract found on
#     this machine is 3.02 from 2011 - the pre-LSTM engine, poor accuracy
#     on ID cards, and too old for modern pytesseract to even parse)
#   - runs fully locally, nothing leaves the host, which is the whole
#     point of the product's self-hosted design (§10)
#   - ONNX, consistent with the NER side of CLOAK
#
# Cost: roughly 2-3 seconds per image on CPU. That is far too slow to run
# on every file inline, so OCR is OPT-IN (enable_ocr=True) rather than
# automatic - the caller decides, e.g. only for images attached to an
# outbound channel, or only when other signals already raised interest.
#
# IMPORTANT - OCR text is not clean text. Real output from this engine on
# a rendered Aadhaar card looks like:
#
#     IND / GovernmentofIndia / NaveenYadav / /DOB?07/06/2007 / 445482493579
#
# Spaces are frequently lost between words, ':' is misread as '?', and
# short noise fragments appear. A classifier trained only on clean
# extracted text will NOT handle this well - see normalise_ocr_text()
# below, and the OCR-specific training data in dataset_generation/.

_engine = None
_engine_lock = threading.Lock()
_engine_failed = False


def _get_engine():
    """Lazily construct the OCR engine - it loads ONNX models, which is
    slow and pointless for callers that never OCR anything."""
    global _engine, _engine_failed
    if _engine is not None or _engine_failed:
        return _engine
    with _engine_lock:
        if _engine is None and not _engine_failed:
            try:
                from rapidocr_onnxruntime import RapidOCR
                _engine = RapidOCR()
            except Exception:  # noqa: BLE001 - missing dep or model load failure
                _engine_failed = True
    return _engine


def ocr_available() -> bool:
    return _get_engine() is not None


# Glue words back apart when OCR has run them together. Only applied to
# long alphabetic runs that contain internal capitals, which is the
# signature of lost spaces ("GovernmentofIndia"), not of a real word.
_RUNON = re.compile(r"\b(?=\w*[a-z])(?=\w*[A-Z])[A-Za-z]{12,}\b")


def _split_runon(match: re.Match) -> str:
    token = match.group(0)
    # Split before interior capitals, and before a lowercase 'of'/'and'
    # style joiner that lost its spaces.
    parts = re.sub(r"(?<!^)(?=[A-Z])", " ", token)
    return parts


def normalise_ocr_text(text: str) -> str:
    """Make OCR output look more like clean extracted text, so the same
    downstream keyword rules and TF-IDF vocabulary have a chance of
    matching. Conservative on purpose - it is better to leave text
    slightly wrong than to invent words that were never there."""
    if not text:
        return text
    out = _RUNON.sub(_split_runon, text)
    out = out.replace("?", ":")            # ':' is the most common OCR confusion here
    out = re.sub(r"[ \t]{2,}", " ", out)
    return out


def extract_text_from_image(data: bytes, normalise: bool = True) -> tuple[str | None, list[str]]:
    """Returns (text, notes). text is None when OCR is unavailable or
    finds nothing - never a silent empty string, so callers can tell
    'no text in the image' apart from 'OCR did not run'."""
    engine = _get_engine()
    if engine is None:
        return None, ["OCR engine unavailable (rapidocr-onnxruntime not installed or failed to load)"]

    try:
        from PIL import Image
        import numpy as np
        image = Image.open(io.BytesIO(data)).convert("RGB")
        result, _elapsed = engine(np.array(image))
    except Exception as exc:  # noqa: BLE001
        return None, [f"OCR failed: {type(exc).__name__}: {exc}"]

    if not result:
        return None, ["OCR ran but found no text in the image"]

    raw = "\n".join(line[1] for line in result)
    notes = [f"text recovered by OCR from {len(result)} detected regions - "
             f"OCR output is noisier than extracted text (lost spaces, character confusions)"]
    return (normalise_ocr_text(raw) if normalise else raw), notes


# ════════════════════════════════════════════════════════════════════════════
# Extraction · text
# ════════════════════════════════════════════════════════════════════════════
# Content extraction: the handoff point between "what is this file"
# (Layers 1-2, format identification) and "what does it actually say"
# (Layer 3, document-type classification).
#
# This is deliberately built as its own step, not folded into Layer 1/2 -
# it needs a CONFIRMED actual_type (from identify_bytes) before it knows
# which extractor to use, so it's always the last thing that runs, never
# the first. See CLOAK-v3-ML-GenAI-Work-Plan.md for the reasoning.
#
# Every extractor caps how much text it returns (MAX_TEXT_CHARS) -
# classifiers need enough text to see vocabulary and structure, not an
# entire 200-page document.

MAX_TEXT_CHARS = 8000
# Above this size, counting exact CSV rows means a full parse we'd rather
# skip - the classifier needs headers/shape, not an exact row count.
MAX_BYTES_FOR_EXACT_ROW_COUNT = 2 * 1024 * 1024


@dataclass
class ExtractedContent:
    text: str | None = None
    column_headers: list[str] | None = None
    row_count: int | None = None
    column_count: int | None = None
    page_count: int | None = None
    extraction_method: str = "none"
    truncated: bool = False
    notes: list[str] = field(default_factory=list)


def _cap_text(text: str) -> tuple[str, bool]:
    if len(text) <= MAX_TEXT_CHARS:
        return text, False
    return text[:MAX_TEXT_CHARS], True


def _pdfium_text(data: bytes) -> tuple[str | None, int]:
    """Text layer via pdfium (the engine Chrome uses). Reads many PDFs pypdf cannot:
    unusual font encodings, broken cross-reference tables. (text or None, page count)."""
    try:
        import pypdfium2 as pdfium
    except ImportError:
        return None, 0
    try:
        pdf = pdfium.PdfDocument(data)
    except Exception:  # noqa: BLE001
        return None, 0
    try:
        chunks: list[str] = []
        total = 0
        for i in range(len(pdf)):
            page = pdf[i]
            try:
                tp = page.get_textpage()
                t = tp.get_text_bounded() or ""
                tp.close()
            except Exception:  # noqa: BLE001
                t = ""
            finally:
                page.close()
            chunks.append(t)
            total += len(t)
            if total >= MAX_TEXT_CHARS:
                break
        text = "\n".join(chunks)
        return (text if text.strip() else None), len(pdf)
    finally:
        pdf.close()


def _extract_pdf(data: bytes) -> ExtractedContent:
    try:
        import pypdf
    except ImportError:
        pypdf = None

    reader = None
    open_error: Exception | None = None
    if pypdf is not None:
        try:
            reader = pypdf.PdfReader(io.BytesIO(data))
        except Exception as exc:  # noqa: BLE001 - pypdf raises several distinct exception types
            open_error = exc
    if reader is None:
        # pypdf missing or unable to open it: pdfium may still read the text layer.
        text, pages = _pdfium_text(data)
        if text:
            text, cap_truncated = _cap_text(text)
            return ExtractedContent(text=text, page_count=pages, extraction_method="pdfium",
                                    truncated=cap_truncated)
        return ExtractedContent(notes=[f"failed to open PDF for extraction: {open_error}" if open_error
                                       else "pypdf not installed - cannot extract PDF text"])

    page_count = len(reader.pages)
    chunks: list[str] = []
    truncated = False
    total_len = 0
    for page in reader.pages:
        try:
            page_text = page.extract_text() or ""
        except Exception:  # noqa: BLE001
            chunks.append("")
            continue
        chunks.append(page_text)
        total_len += len(page_text)
        if total_len >= MAX_TEXT_CHARS:
            truncated = True
            break

    text, cap_truncated = _cap_text("\n".join(chunks))
    method = "pypdf"
    # Pages pypdf could not read come back as "" and join to bare newlines, which is
    # NOT text: treat whitespace-only as empty, try pdfium, and only then leave it to OCR.
    if not text.strip():
        alt, _ = _pdfium_text(data)
        if alt:
            text, cap_truncated = _cap_text(alt)
            method = "pdfium"
        else:
            text = ""
    return ExtractedContent(
        text=text if text else None,
        page_count=page_count,
        extraction_method=method,
        truncated=truncated or cap_truncated,
        notes=[] if text else ["no extractable text found (scanned/image-only PDF?) - OCR is tried next"],
    )


def _extract_docx(data: bytes) -> ExtractedContent:
    try:
        import docx
    except ImportError:
        return ExtractedContent(notes=["python-docx not installed - cannot extract docx text"])

    try:
        document = docx.Document(io.BytesIO(data))
    except Exception as exc:  # noqa: BLE001
        return ExtractedContent(notes=[f"failed to open docx for extraction: {exc}"])

    paragraphs = [p.text for p in document.paragraphs if p.text.strip()]
    # Table text carries real content in a lot of business documents
    # (invoices, contracts) - worth including, not just prose paragraphs.
    for table in document.tables:
        for row in table.rows:
            row_text = " | ".join(cell.text for cell in row.cells if cell.text.strip())
            if row_text:
                paragraphs.append(row_text)

    text, truncated = _cap_text("\n".join(paragraphs))
    return ExtractedContent(
        text=text if text else None,
        extraction_method="python-docx",
        truncated=truncated,
        notes=[] if text else ["no extractable text found (empty document?)"],
    )


def _extract_xlsx(data: bytes) -> ExtractedContent:
    try:
        import openpyxl
    except ImportError:
        return ExtractedContent(notes=["openpyxl not installed - cannot extract xlsx structure"])

    try:
        workbook = openpyxl.load_workbook(io.BytesIO(data), read_only=True, data_only=True)
    except Exception as exc:  # noqa: BLE001
        return ExtractedContent(notes=[f"failed to open xlsx for extraction: {exc}"])

    notes = []
    if len(workbook.sheetnames) > 1:
        notes.append(f"{len(workbook.sheetnames)} sheets present - only the first sheet was inspected")
    sheet = workbook[workbook.sheetnames[0]]

    headers: list[str] = []
    first_row = next(sheet.iter_rows(min_row=1, max_row=1, values_only=True), None)
    if first_row:
        headers = [str(c) if c is not None else "" for c in first_row]

    # sheet.max_row/max_column read from the file's own declared <dimension>
    # element when present - fast, but several real-world xlsx exporters
    # (seen on an actual test file, not hypothetical) omit it entirely, in
    # which case openpyxl returns None rather than scanning to find out.
    # Fall back to counting by streaming through rows - still cheap in
    # read_only mode since nothing is loaded into memory at once.
    row_count = sheet.max_row
    column_count = sheet.max_column

    if row_count is None:
        row_count = sum(1 for _ in sheet.iter_rows(values_only=True))
        notes.append("sheet had no declared dimension - row count obtained by streaming scan, not from file metadata")
    row_count = max(row_count - 1, 0)  # minus the header row

    if column_count is None:
        column_count = len(headers) if headers else None
        if column_count is not None:
            notes.append("sheet had no declared dimension - column count taken from header row width")

    workbook.close()

    return ExtractedContent(
        column_headers=headers or None,
        row_count=row_count,
        column_count=column_count,
        extraction_method="openpyxl",
        notes=notes,
    )


# csv.Sniffer() alone is unreliable - it happily "detects" a delimiter
# from a single line with no real tabular structure (found via a real
# bug: it read a SQL dump's leading comment line as a one-column CSV
# header and discarded the entire dump's text). Two independent guards
# fix this: (1) an explicit exclusion for text that is unambiguously SQL,
# since SQL genuinely can have comma-consistent lines (repeated `INSERT
# INTO ... VALUES (a, b, c);` rows) that would otherwise pass a naive
# consistency check too; (2) requiring several lines to actually agree on
# a non-trivial column count, not just trusting one sniffed sample.
_SQL_SIGNAL = re.compile(r"\b(create\s+table|drop\s+table|insert\s+into|select\s+.+\bfrom\b)\b", re.IGNORECASE)
_MIN_DELIMITED_LINES = 3
_MIN_DELIMITED_COLUMNS = 2
_MIN_CONSISTENT_FRACTION = 0.8


def _looks_like_delimited_text(sample: str) -> tuple[bool, str | None]:
    if _SQL_SIGNAL.search(sample):
        return False, None

    lines = [line for line in sample.split("\n") if line.strip()][:20]
    if len(lines) < _MIN_DELIMITED_LINES:
        return False, None

    try:
        dialect = csv.Sniffer().sniff(sample, delimiters=",\t;|")
    except csv.Error:
        return False, None

    delimiter = dialect.delimiter
    field_counts = [len(line.split(delimiter)) for line in lines]
    if field_counts[0] < _MIN_DELIMITED_COLUMNS:
        return False, None

    consistent = sum(1 for c in field_counts if c == field_counts[0])
    if consistent / len(field_counts) < _MIN_CONSISTENT_FRACTION:
        return False, None

    return True, delimiter


def _extract_unknown_text(data: bytes, size_for_row_count_check: int) -> ExtractedContent:
    """actual_type == 'unknown' covers plain text, CSV, SQL dumps, and
    anything else with no magic-byte signature (see the README's
    documented gap: Layer 1 can't mismatch-check these, so this is where
    content-based handling for them actually starts)."""
    try:
        text = data.decode("utf-8")
    except UnicodeDecodeError:
        return ExtractedContent(notes=["not decodable as UTF-8 text - likely binary content with no known format"])

    sample = text[:4096]
    is_delimited, delimiter = _looks_like_delimited_text(sample)

    if is_delimited:
        reader = csv.reader(io.StringIO(text.split("\n", 1)[0]), delimiter=delimiter)
        headers = next(reader, [])
        row_count = None
        notes = []
        if size_for_row_count_check <= MAX_BYTES_FOR_EXACT_ROW_COUNT:
            row_count = sum(1 for _ in io.StringIO(text)) - 1  # minus header
        else:
            notes.append("file too large for an exact row count - only headers extracted")
        return ExtractedContent(
            column_headers=headers,
            row_count=row_count,
            column_count=len(headers),
            extraction_method="csv_sniff",
            notes=notes,
        )

    capped_text, truncated = _cap_text(text)
    return ExtractedContent(
        text=capped_text if capped_text.strip() else None,
        extraction_method="plain_text",
        truncated=truncated,
        notes=[] if capped_text.strip() else ["decoded as text but content is empty/whitespace"],
    )


# actual_type values with no extractor yet - same treatment as an image:
# recorded honestly as "not extracted", never silently returned as if
# nothing was there to extract.
_IMAGE_TYPES = {"png", "jpeg", "gif", "bmp", "tiff"}

_NO_EXTRACTOR_NOTES = {
    "png": "image - OCR not enabled (pass enable_ocr=True to read it)",
    "jpeg": "image - OCR not enabled (pass enable_ocr=True to read it)",
    "gif": "image - OCR not enabled (pass enable_ocr=True to read it)",
    "bmp": "image - OCR not enabled (pass enable_ocr=True to read it)",
    "tiff": "image - OCR not enabled (pass enable_ocr=True to read it)",
    "ole2": "legacy OLE2 container (.doc/.xls/.ppt) - no extractor implemented yet",
    "zip": "generic zip container, not a recognised document type on its own - Layer 2's per-entry extraction applies instead",
    "pptx": "pptx extraction not implemented yet",
    "rtf": "rtf extraction not implemented yet",
}


def extract_content(actual_type: str, data: bytes, enable_ocr: bool = False) -> ExtractedContent:
    """Dispatches purely on Layer 1's already-confirmed actual_type - never
    re-derives format from the filename or re-runs magic-byte detection.
    That's the whole point of this being a separate, later step.

    enable_ocr: run OCR on image types. Off by default because OCR costs
    ~2-3 seconds per image on CPU - far too slow to run inline on every
    file. The caller decides when it's worth it (e.g. images leaving on
    an outbound channel)."""
    if actual_type == "pdf":
        return _extract_pdf(data)
    if actual_type == "docx":
        return _extract_docx(data)
    if actual_type == "xlsx":
        return _extract_xlsx(data)
    if actual_type == "unknown":
        return _extract_unknown_text(data, size_for_row_count_check=len(data))
    if actual_type in _IMAGE_TYPES and enable_ocr:
        text, notes = extract_text_from_image(data)
        capped, truncated = _cap_text(text or "")
        return ExtractedContent(
            text=capped if capped.strip() else None,
            extraction_method="ocr_rapidocr",
            truncated=truncated,
            notes=notes,
        )

    note = _NO_EXTRACTOR_NOTES.get(actual_type, f"no extractor implemented for actual_type={actual_type!r}")
    return ExtractedContent(extraction_method="none", notes=[note])


# ════════════════════════════════════════════════════════════════════════════
# Layer 3 · document type
# ════════════════════════════════════════════════════════════════════════════
# Layer 3 document-type classifier - the serving entry point.
#
#     from docclass import Layer3DocType
#
#     model = Layer3DocType.load("models/layer3_25class.joblib")
#     p = model.predict(text=extracted.text,
#                       column_headers=extracted.column_headers)
#
#     p.document_type   one of 25 classes (+ credentials)
#     p.confidence      0.0 - 1.0
#     p.confident       confidence >= threshold; False -> route to review
#     p.action          suggested policy action, from the lookup table
#     p.top_types       [(class, prob), ...] top 3
#     p.abstained       True when there was nothing to classify
#
# That is the whole interface. There is no second concept to learn: the
# model says what the document IS, and a plain lookup table says what should
# happen to it. The table ships inside the bundle but is configuration, not
# machine learning - override it in the policy engine without retraining.
#
# USE THIS, NOT THE RAW BUNDLE. Two behaviours live here rather than in the
# pickle, and both exist because of measured failures:
#
#   * the ABSTAIN guard. With no text and no column headers, an earlier model
#     in this project confidently predicted a sensitive class at ~99.99%.
#   * CREDENTIALS ESCALATION. Under plain argmax, 8.25% of credential
#     documents were under-restricted, because a file scoring 30%
#     credentials and 35% something-else was called something-else. Those
#     two outcomes do not cost the same.
#
# BOUNDARY RULE (architecture doc S7): this model sets a field. It does not
# block anything. The policy engine decides the action, and a model may only
# ADD restriction, never remove it.

@dataclass
class Prediction:
    document_type: str
    confidence: float
    confident: bool
    action: str
    reason: str = ""
    top_types: list[tuple[str, float]] = field(default_factory=list)
    abstained: bool = False
    escalated: bool = False
    notes: list[str] = field(default_factory=list)


class Layer3DocType:
    def __init__(self, bundle: dict):
        self._b = bundle
        self.classes = np.array(bundle["classes"])
        self.threshold = float(bundle.get("confidence_threshold", 0.50))
        self.credentials_escalation = float(
            bundle.get("credentials_escalation", 0.25))
        self._policy = {k: tuple(v) for k, v in bundle.get("policy", {}).items()}
        self._low_conf = bundle.get("low_confidence_action", "ROUTE_TO_REVIEW")
        self._abstain = bundle.get("abstain_action", "ROUTE_TO_REVIEW")

    @classmethod
    def load(cls, path: str) -> "Layer3DocType":
        if not os.path.exists(path):
            raise FileNotFoundError(path)
        return cls(joblib.load(path))

    @property
    def metrics(self) -> dict:
        return self._b.get("metrics", {})

    def action_for(self, document_type: str) -> tuple[str, str]:
        return self._policy.get(document_type,
                                (self._low_conf, "unknown class - fail closed"))

    @staticmethod
    def build_feature_text(text: str | None,
                           column_headers: list | None) -> str:
        """MUST match train_25class.feature_text exactly. Any divergence is
        a silent train/serve skew that nothing else will report."""
        t = (text or "").strip()
        if t:
            return t
        if column_headers:
            return (" ".join(str(h) for h in column_headers) + " ") * 5
        return ""

    def predict(self, text: str | None = None,
                column_headers: list | None = None,
                **_ignored) -> Prediction:
        feat = self.build_feature_text(text, column_headers)
        if not feat:
            return Prediction(
                document_type="unknown", confidence=0.0, confident=False,
                action=self._abstain,
                reason="nothing to classify - fail closed",
                abstained=True,
                notes=["no text and no column headers were extracted"])

        X = self._b["vectorizer"].transform([feat])
        clf = self._b["classifier"]
        order = [list(clf.classes_).index(c) for c in self.classes]
        proba = clf.predict_proba(X)[0][order]

        best = int(proba.argmax())
        chosen, escalated = best, False
        if "credentials" in self.classes:
            ci = int(np.where(self.classes == "credentials")[0][0])
            if proba[ci] >= self.credentials_escalation and ci != best:
                chosen, escalated = ci, True

        label = str(self.classes[chosen])
        conf = float(proba[chosen])
        confident = conf >= self.threshold
        action, reason = self.action_for(label)

        notes = []
        if escalated:
            notes.append(
                f"escalated to credentials: p={conf:.2f} cleared the "
                f"{self.credentials_escalation:.2f} threshold even though "
                f"{self.classes[best]!r} scored higher")
        if not confident:
            action = self._low_conf
            reason = (f"confidence {conf:.2f} below the "
                      f"{self.threshold:.2f} threshold")
            notes.append("routed to review - do NOT auto-apply policy, and "
                         "do not let the document through pending review")

        top = np.argsort(proba)[-3:][::-1]
        return Prediction(
            document_type=label, confidence=conf, confident=confident,
            action=action, reason=reason,
            top_types=[(str(self.classes[i]), round(float(proba[i]), 3))
                       for i in top],
            abstained=False, escalated=escalated, notes=notes)


# ════════════════════════════════════════════════════════════════════════════
# Policy table
# ════════════════════════════════════════════════════════════════════════════
# Class -> suggested action. CONFIGURATION, not machine learning.
#
# This is deliberately a plain table and deliberately NOT inside the model.
# The model answers one question - "what kind of document is this?" - and
# this table answers the other one - "so what should happen?".
#
# Keeping them apart means:
#   * changing policy does not require retraining anything
#   * the model has one job and one output, which is the whole point of
#     dropping the earlier tier concept
#   * your policy engine can override any row without touching the artifact
#
# The actions below are SUGGESTIONS derived from what each class contains.
# The policy engine owns the final decision - architecture doc S7: a model
# may only ADD restriction, never remove it.

# class -> (suggested action, one-line reason)
POLICY: dict[str, tuple[str, str]] = {
    # ---- secrets: worst case is loss of ACCESS, not just disclosure ----
    "credentials": ("BLOCK_ALL",
                    "API keys, passwords and private keys grant access to "
                    "the systems holding everything else"),

    # ---- personal data: block external transfer ----
    "government_identity": ("BLOCK_EXTERNAL", "government identity proof"),
    "travel_immigration": ("BLOCK_EXTERNAL", "identity plus right-to-remain status"),
    "healthcare_medical": ("BLOCK_EXTERNAL", "health data, own regulatory regime"),
    "insurance": ("BLOCK_EXTERNAL", "often carries diagnosis and financial detail"),
    "customer_client": ("BLOCK_EXTERNAL", "bulk personal data about customers"),
    "employment_hr": ("BLOCK_EXTERNAL", "salary, performance and identity of staff"),
    "education": ("BLOCK_EXTERNAL", "student records, protected in most regimes"),
    "financial": ("BLOCK_EXTERNAL", "a named person's money"),
    "tax": ("BLOCK_EXTERNAL", "income plus a tax identifier"),
    "data_records": ("BLOCK_EXTERNAL", "bulk extracts, risk scales with row count"),
    "property_real_estate": ("BLOCK_EXTERNAL", "names, addresses and consideration"),

    # ---- business confidential: warn ----
    "legal": ("WARN_EXTERNAL", "commercial terms, often under an NDA"),
    "corporate_business": ("WARN_EXTERNAL", "board and strategy material"),
    "intellectual_property": ("WARN_EXTERNAL", "pre-filing disclosure can destroy rights"),
    "cybersecurity": ("WARN_EXTERNAL", "a map of how to attack us"),
    "procurement_vendor": ("WARN_EXTERNAL", "pricing and supplier terms"),
    "sales_marketing": ("WARN_EXTERNAL", "unreleased pricing and plans"),
    "accounting_billing": ("LOG_ONLY", "routine AP/AR traffic moves all day"),

    # ---- low sensitivity ----
    "operations_internal": ("LOG_ONLY", "internal process material"),
    "communication": ("LOG_ONLY", "judge on the attachments, not the envelope"),
    "government_regulatory": ("LOG_ONLY", "permits and filings are often public"),
    "certificates_licenses": ("LOG_ONLY", "attestations, not identity proof"),
    "technical_it": ("LOG_ONLY", "code and docs; IP, but not personal data"),
    "research_scientific": ("LOG_ONLY", "review if unpublished"),
    "general_public": ("LOG_ONLY", "already published"),
}

# What to do when the model is not confident enough, or abstained.
# FAIL CLOSED. Measured on the previous model: below the confidence
# threshold, credential documents were 78.6% under-restricted. A permissive
# fallback undoes every other control here.
LOW_CONFIDENCE_ACTION = "ROUTE_TO_REVIEW"
ABSTAIN_ACTION = "ROUTE_TO_REVIEW"

ACTIONS = ("BLOCK_ALL", "BLOCK_EXTERNAL", "WARN_EXTERNAL", "LOG_ONLY",
           "ROUTE_TO_REVIEW")


def policy_action_for(document_type: str) -> tuple[str, str]:
    return POLICY.get(document_type,
                      (LOW_CONFIDENCE_ACTION, "unknown class - fail closed"))


def validate_policy(classes) -> list[str]:
    problems = []
    for c in classes:
        if c not in POLICY:
            problems.append(f"class {c!r} has no policy row - would fail closed")
    for c, (a, _) in POLICY.items():
        if a not in ACTIONS:
            problems.append(f"class {c!r} has unknown action {a!r}")
        if c not in classes:
            problems.append(f"policy row {c!r} is not a class the model predicts")
    return problems


# ════════════════════════════════════════════════════════════════════════════
# API
# ════════════════════════════════════════════════════════════════════════════

# ── Config ────────────────────────────────────────────────────────────────────

# Same repository as the PII model unless told otherwise.
DOCCLASS_REPO = os.getenv("DOCCLASS_REPO") or os.getenv("MODEL_REPO", "Wardline/CLOAK_V3.12_onnx")
DOCCLASS_FILE = os.getenv("DOCCLASS_FILE", "layer3_25class.joblib")
# The bundle is a pickle: loading it runs code. Refuse anything but the exact
# artifact that was reviewed. Rotate this together with the file, never alone.
DOCCLASS_SHA256 = os.getenv(
    "DOCCLASS_SHA256", "37be0e164060fa1f6ce7f3c627e588b60bb896902f1cc55523ee100df28ac20f")
# Explicit path. Without it, a copy uploaded with the Space (beside this file
# or in models/) is used, and only then the Hub.
DOCCLASS_MODEL_PATH = os.getenv("DOCCLASS_MODEL_PATH")
# Returned with every answer, so a caller can tell results of an older model
# from current ones and reclassify after an upgrade.
MODEL_VERSION = f"{DOCCLASS_FILE}@{(DOCCLASS_SHA256 or 'unpinned')[:12]}"

MAX_FILE_BYTES = int(os.getenv("DOCCLASS_MAX_FILE_BYTES", str(25 * 1024 * 1024)))
MAX_HEADERS = 500
# Inputs longer than what the extractors hand the classifier are cut to the
# same length, so text CLOAK extracted itself is seen the way training saw it.
MAX_CLASSIFY_CHARS = 200_000

# Archives are opened inside a shared CPU Space: far tighter than the library
# defaults (500 MB, 10,000 entries), which assume a dedicated host.
SCAN_LIMITS = ScanLimits(max_depth=3, max_expansion_ratio=100.0,
                         max_total_uncompressed_bytes=150 * 1024 * 1024,
                         max_entries=500)
MAX_NESTED_CLASSIFIED = 200

# OCR costs ~2-3 s of CPU per image. Bound how many run at once so a batch of
# photos cannot starve text classification and PII detection of CPU.
IMAGE_TYPES = {"png", "jpeg", "gif", "bmp", "tiff"}
_ocr_slots = threading.BoundedSemaphore(int(os.getenv("DOCCLASS_OCR_CONCURRENCY", "2")))
OCR_WAIT_SECONDS = 60
# A scanned PDF has no text layer; its first pages are rendered and OCR'd.
PDF_OCR_PAGES = int(os.getenv("DOCCLASS_PDF_OCR_PAGES", "3"))

# ── Model loading ─────────────────────────────────────────────────────────────

MODEL: Layer3DocType | None = None
LOAD_ERROR: str | None = None


def _sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def load(hf_token: str | None = None) -> None:
    """Fetch, verify and load the model. Never raises — a failure is recorded
    and reported by /health, and only the /file paths go dark."""
    global MODEL, LOAD_ERROR
    try:
        # Uploaded with the Space itself (next to app.py, or in models/), an
        # explicit path, or else the Hub. The checksum below applies to all.
        here = os.path.dirname(os.path.abspath(__file__))
        local = [os.path.join(here, DOCCLASS_FILE), os.path.join(here, "models", DOCCLASS_FILE)]
        path = DOCCLASS_MODEL_PATH or next((c for c in local if os.path.isfile(c)), None)
        if not path:
            from huggingface_hub import hf_hub_download
            path = hf_hub_download(DOCCLASS_REPO, DOCCLASS_FILE, token=hf_token,
                                   local_dir=os.getenv("DOCCLASS_DIR", "/tmp/cloak_docclass"))
        digest = _sha256(path)
        if DOCCLASS_SHA256 and digest != DOCCLASS_SHA256.lower():
            raise RuntimeError(f"checksum mismatch: got {digest}, expected {DOCCLASS_SHA256} — refusing to load")
        MODEL = Layer3DocType.load(path)
        LOAD_ERROR = None
        print(f"[docclass] loaded {DOCCLASS_FILE}: {len(MODEL.classes)} classes, threshold {MODEL.threshold}")
    except Exception as exc:  # noqa: BLE001
        MODEL = None
        LOAD_ERROR = f"{type(exc).__name__}: {exc}"
        print(f"[docclass] NOT loaded — /file paths will answer 503: {LOAD_ERROR}")


def status() -> dict:
    ocr = None
    try:
        ocr = ocr_available()
    except Exception:  # noqa: BLE001
        ocr = False
    return {
        "loaded": MODEL is not None,
        "error": LOAD_ERROR,
        "file": DOCCLASS_FILE,
        "model_version": MODEL_VERSION,
        "classes": len(MODEL.classes) if MODEL else 0,
        "threshold": MODEL.threshold if MODEL else None,
        "credentials_escalation": MODEL.credentials_escalation if MODEL else None,
        "ocr": ocr,
        "max_file_bytes": MAX_FILE_BYTES,
    }


def _model() -> Layer3DocType:
    if MODEL is None:
        raise HTTPException(status_code=503, detail=f"Document model unavailable: {LOAD_ERROR or 'not loaded'}")
    return MODEL


# ── Pipeline ──────────────────────────────────────────────────────────────────

def _prediction(p) -> dict:
    return {
        "document_type": p.document_type,
        "confidence":    round(float(p.confidence), 4),
        "confident":     bool(p.confident),
        "action":        p.action,
        "reason":        p.reason,
        "top_types":     [[t, float(c)] for t, c in p.top_types],
        "abstained":     bool(p.abstained),
        "escalated":     bool(p.escalated),
        "model_version": MODEL_VERSION,
    }


def _ocr_pdf(data: bytes) -> tuple[str | None, list[str]]:
    """Render the first pages of a scanned PDF and OCR them."""
    try:
        import pypdfium2 as pdfium
    except ImportError:
        return None, ["pypdfium2 not installed - cannot OCR a scanned PDF"]
    texts: list[str] = []
    notes: list[str] = []
    try:
        pdf = pdfium.PdfDocument(data)
        try:
            for i in range(min(len(pdf), PDF_OCR_PAGES)):
                img = pdf[i].render(scale=2).to_pil()
                buf = io.BytesIO()
                img.save(buf, format="PNG")
                t, n = extract_text_from_image(buf.getvalue())
                if t:
                    texts.append(t)
                notes.extend(n[:1])
        finally:
            pdf.close()
    except Exception as exc:  # noqa: BLE001
        return None, [f"could not render PDF for OCR: {type(exc).__name__}"]
    text = "\n".join(texts)[:MAX_TEXT_CHARS]
    return (text or None), notes


def _extract(actual_type: str, data: bytes) -> tuple[object, bool]:
    """Extract text, with OCR decided by the VERIFIED type: an image renamed
    .pdf is still OCR'd, a PDF renamed .jpg still gets the PDF parser."""
    if actual_type in IMAGE_TYPES:
        if not _ocr_slots.acquire(timeout=OCR_WAIT_SECONDS):
            raise HTTPException(status_code=503, detail="OCR busy - retry later")
        try:
            return extract_content(actual_type, data, enable_ocr=True), True
        finally:
            _ocr_slots.release()
    ex = extract_content(actual_type, data, enable_ocr=False)
    if actual_type == "pdf" and not (ex.text or "").strip():
        if not _ocr_slots.acquire(timeout=OCR_WAIT_SECONDS):
            raise HTTPException(status_code=503, detail="OCR busy - retry later")
        try:
            text, notes = _ocr_pdf(data)
        finally:
            _ocr_slots.release()
        if text:
            ex.text = text
            ex.extraction_method = "pdf_ocr_rapidocr"
            ex.notes = list(ex.notes) + notes
        return ex, True
    return ex, False


def analyze_file(data: bytes, filename: str, model: Layer3DocType, return_text: bool = False) -> dict:
    """All three layers for one file — the handover's classify_file, bounded.

    return_text: also return the extracted / OCR'd text (at most MAX_TEXT_CHARS),
    so CLOAK can run its PII scan on images and scanned PDFs without a second OCR
    pass. Off by default: a caller that only wants the document type is not sent
    the contents back."""
    analysis = analyze_bytes(data, filename=filename, limits=SCAN_LIMITS, keep_bytes=True)
    identity = analysis.identity
    ex, ocr_used = _extract(identity.actual_type, data)
    p = model.predict(text=ex.text, column_headers=ex.column_headers)

    nested = []
    for entry in analysis.container.entries[:MAX_NESTED_CLASSIFIED]:
        if entry.identity is None or entry.data is None:
            # An encrypted member leaving the company is a finding on its own.
            nested.append({"path": entry.path, "actual_type": None,
                           "encrypted": bool(entry.encrypted), "document_type": None})
            continue
        e_ex, e_ocr = _extract(entry.identity.actual_type, entry.data)
        e_p = model.predict(text=e_ex.text, column_headers=e_ex.column_headers)
        nested.append({"path": entry.path, "actual_type": entry.identity.actual_type,
                       "type_mismatch": bool(entry.identity.type_mismatch),
                       "encrypted": False, "ocr_used": e_ocr, **_prediction(e_p)})

    return {
        **_prediction(p),
        "declared_type":     identity.declared_type,
        "actual_type":       identity.actual_type,
        "type_mismatch":     bool(identity.type_mismatch),
        "likely_encrypted":  bool(identity.likely_encrypted_or_packed),
        "ocr_used":          ocr_used,
        "extraction_method": getattr(ex, "extraction_method", "none"),
        "truncated":         bool(getattr(ex, "truncated", False)),
        "nested_entries":    nested,
        "nested_truncated":  bool(analysis.container.truncated) or len(analysis.container.entries) > MAX_NESTED_CLASSIFIED,
        **({"text": (getattr(ex, "text", None) or "")[:MAX_TEXT_CHARS]} if return_text else {}),
    }


# ── HTTP ──────────────────────────────────────────────────────────────────────

class ClassifyRequest(BaseModel):
    text: str | None = Field(default=None, max_length=MAX_CLASSIFY_CHARS)
    column_headers: list[str] | None = Field(default=None, max_length=MAX_HEADERS)


def build_router(verify: Callable[[Request], None]) -> APIRouter:
    router = APIRouter()

    @router.post("/file/classify")
    def classify(body: ClassifyRequest, request: Request):
        verify(request)
        model = _model()
        t0 = time.time()
        text = (body.text or "")[:MAX_TEXT_CHARS]
        headers = [str(h)[:200] for h in (body.column_headers or [])]
        p = model.predict(text=text, column_headers=headers)
        return {**_prediction(p), "processing_time_ms": int((time.time() - t0) * 1000)}

    @router.post("/file/analyze")
    async def analyze(request: Request):
        verify(request)
        model = _model()
        declared = request.headers.get("content-length")
        if declared and declared.isdigit() and int(declared) > MAX_FILE_BYTES:
            raise HTTPException(status_code=413, detail=f"File larger than {MAX_FILE_BYTES} bytes")
        data = await request.body()
        if len(data) > MAX_FILE_BYTES:
            raise HTTPException(status_code=413, detail=f"File larger than {MAX_FILE_BYTES} bytes")
        from urllib.parse import unquote
        filename = os.path.basename(unquote(request.headers.get("x-filename", "") or "file"))[:255]
        t0 = time.time()
        from starlette.concurrency import run_in_threadpool
        try:
            return_text = request.headers.get("x-return-text", "") == "1"
            result = await run_in_threadpool(analyze_file, data, filename, model, return_text)
        except HTTPException:
            raise
        except Exception as exc:  # noqa: BLE001
            # Never echo file contents; the type name is enough to debug.
            raise HTTPException(status_code=422, detail=f"Could not analyze file: {type(exc).__name__}")
        result["processing_time_ms"] = int((time.time() - t0) * 1000)
        return result

    return router