""" 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: ) `/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"(? 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 # 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