safetynet/apps/inspector/engine.py

"""Private content inspection. No external inference, raw-content logs or remote URLs."""
import base64
import hashlib
import io
import math
import os
import re
import shutil
import subprocess
import tempfile
import time
import unicodedata
import zipfile
from pathlib import Path
from stage_timings import timed, stage

os.environ['HF_HUB_OFFLINE']='1'
os.environ['TRANSFORMERS_OFFLINE']='1'
os.environ['HF_HUB_DISABLE_TELEMETRY']='1'

MAX_BYTES = 20 * 1024 * 1024
MAX_TEXT = 1_000_000
_MODEL = None
_OCR = None


@timed("ocr.load")
def local_ocr():
    global _OCR
    if _OCR is None:
        import onnxruntime
        onnxruntime.disable_telemetry_events()
        import rapidocr
        from rapidocr import RapidOCR
        directory=Path(rapidocr.__file__).parent/"models"
        names={"Det":"ch_PP-OCRv4_det_infer.onnx","Rec":"ch_PP-OCRv4_rec_infer.onnx","Cls":"ch_ppocr_mobile_v2.0_cls_infer.onnx"}
        checksums={"Det":"d2a7720d45a54257208b1e13e36a8479894cb74155a5efe29462512d42f49da9",
                   "Rec":"48fc40f24f6d2a207a2b1091d3437eb3cc3eb6b676dc3ef9c37384005483683b",
                   "Cls":"e47acedf663230f8863ff1ab0e64dd2d82b838fceb5957146dab185a89d6215c"}
        params={"Global.log_level":"critical","EngineConfig.onnxruntime.intra_op_num_threads":2,
                "EngineConfig.onnxruntime.inter_op_num_threads":1}
        for task,name in names.items():
            model=directory/name
            if not model.is_file(): raise ValueError("Bundled offline OCR model unavailable")
            if hashlib.sha256(model.read_bytes()).hexdigest()!=checksums[task]: raise ValueError("Offline OCR model checksum mismatch")
            params[task+".model_path"]=str(model)
        # Explicit paths bypass RapidOCR's download path entirely.
        _OCR=RapidOCR(params=params)
    return _OCR


@timed("yara")
def scan_download(data, policy):
    """Malware gate only: never implies document/DLP coverage."""
    result={"decision":"block","ruleId":"yara-unavailable","engine":"yara",
            "sha256":hashlib.sha256(data).hexdigest(),"bytes":len(data),
            "matchedRules":[],"coverage":{"complete":False}}
    try:
        config=policy.get("security",{}).get("yara",{})
        if config.get("enabled") is not True:
            result["ruleId"]="yara-required";return result
        maximum=config.get("maxScanBytes",MAX_BYTES)
        if not isinstance(maximum,int) or maximum<=0 or len(data)>min(maximum,MAX_BYTES):
            result["ruleId"]="yara-size-limit";return result
        entries=config.get("rules",[])
        if not isinstance(entries,list) or not 1<=len(entries)<=128: raise ValueError()
        sources={};ids={};total=0
        for index,entry in enumerate(entries):
            source=entry["source"];identifier=entry["id"]
            if not isinstance(source,str) or not isinstance(identifier,str) or not identifier: raise ValueError()
            total+=len(source.encode("utf-8"))
            if total>262144: raise ValueError()
            # No imported modules: console can leak scanned content; others expand attack surface.
            if re.search(r'\bimport\s+"',source): raise ValueError()
            namespace="policy_"+str(index);sources[namespace]=source;ids[namespace]=identifier
        import yara
        rules=yara.compile(sources=sources,includes=False,error_on_warning=True)
        matches=rules.match(data=data,timeout=5)
        if config.get("scanOfficeMacros",True) is not False:
            from office_macros import macro_payloads
            try:
                with stage("macros"):
                    summary,payloads=macro_payloads(data)
                result["officeMacros"]=summary
                deadline=time.monotonic()+10
                for payload in payloads:
                    if time.monotonic()>=deadline: raise ValueError('macro-scan-timeout')
                    matches.extend(rules.match(data=payload,timeout=5))
            except Exception:
                result["ruleId"]="office-macro-inspection-failed"
                result["officeMacros"]={"office":data.startswith(b'\xd0\xcf\x11\xe0') or data.startswith(b'PK'),
                                       "present":False,"modules":0,"complete":False}
                return result
        else:
            result["officeMacros"]={"office":False,"present":False,"modules":0,"complete":False}
        result["coverage"]["complete"]=True
        result["matchedRules"]=[{"id":identifier,"rule":rule} for identifier,rule in
                                dict.fromkeys((ids[item.namespace],item.rule) for item in matches)]
        if matches: result["ruleId"]=ids[matches[0].namespace]
        else: result["decision"]="allow";result["ruleId"]=None
    except ImportError: pass
    except Exception: result["ruleId"]="yara-inspection-failed"
    return result


@timed("inspect")
def reference_document(data):
    """Return privacy-preserving registration features, without extracted plaintext."""
    result={"sha256":hashlib.sha256(data).hexdigest(),"bytes":len(data),
            "contentType":content_kind(data),"fragmentHashes":[],"coverage":{"complete":False}}
    text="";media=[];hashes=set()
    try:
        if len(data)>MAX_BYTES: raise ValueError()
        text,complete,media=extract(data,result["contentType"])
        if len(text)>MAX_TEXT: raise ValueError()
        hashes=fragment_hashes(text)
        result["fragmentHashes"]=sorted(hashes)[:1000]
        truncated=len(hashes)>1000
        result["coverage"]={"complete":complete and not truncated,"extractionComplete":complete,
                            "featuresComplete":not truncated,"fragmentHashesTruncated":truncated,
                            "fragmentHashesTotal":len(hashes),"fragmentHashesRetained":len(result["fragmentHashes"]),
                            "textCharacters":len(text),"mediaItems":len(media)}
        result["registrationReady"]=bool(complete and not truncated and (hashes or text.strip()))
    except Exception: result["registrationReady"]=False
    finally:
        # Release transient content references; immutable Python strings cannot be securely wiped.
        text="";media.clear();hashes.clear();data=None
    return result


def fragment_hashes(text):
    normalized=unicodedata.normalize("NFKC",text).lower().replace("ß","ss").replace("ς","σ")
    words=re.findall(r"\w+",normalized)
    # Every shingle matches excerpts even when their starting offset changes.
    hashes={hashlib.sha256(" ".join(words[index:index+8]).encode()).hexdigest()
            for index in range(max(0,len(words)-7))}
    for paragraph in re.split(r"\n\s*\n",normalized):
        paragraph=" ".join(paragraph.split())
        if len(paragraph)>=32: hashes.add(hashlib.sha256(paragraph.encode()).hexdigest())
    return hashes


def cosine(left, right):
    if len(left) != len(right) or not left:
        raise ValueError("Embedding dimensions do not match")
    if not all(math.isfinite(float(v)) for v in [*left, *right]):
        raise ValueError("Invalid embedding")
    denominator = math.sqrt(sum(v*v for v in left)) * math.sqrt(sum(v*v for v in right))
    if not denominator:
        raise ValueError("Empty embedding")
    return sum(a*b for a, b in zip(left, right)) / denominator


def content_kind(data):
    if data.startswith(b"%PDF-"): return "pdf"
    if data.startswith(b"PK\x03\x04"): return "office"
    if data.startswith(b"\x89PNG\r\n\x1a\n") or data.startswith(b"\xff\xd8\xff") or data.startswith((b"GIF87a", b"GIF89a")): return "image"
    if data.startswith(b"RIFF"):
        return {b"WAVE": "audio", b"WEBP": "image", b"AVI ": "video"}.get(data[8:12], "unknown")
    if data.startswith((b"ID3", b"fLaC", b"OggS")): return "audio"
    if len(data)>12 and data[4:8] == b"ftyp": return "video"
    try:
        text=data.decode("utf-8", errors="strict")
        if "\x00" in text: return "unknown"
        return "text"
    except UnicodeError: return "unknown"


@timed("ocr")
def ocr_image(data):
    from PIL import Image, ImageSequence
    image=Image.open(io.BytesIO(data))
    frames=getattr(image,"n_frames",1)
    if image.width*image.height>20_000_000 or frames>32 or image.width*image.height*frames>50_000_000:
        raise ValueError("Image processing budget exceeded")
    if os.environ.get("SAFETYNET_OCR_ENGINE","rapidocr")=="rapidocr":
        import numpy as np
        engine=local_ocr();parts=[]
        for frame in ImageSequence.Iterator(image):
            output=engine(np.asarray(frame.convert("RGB")))
            parts.extend(output.txts or [])
        return "\n".join(parts),True
    configured=os.environ.get("TESSERACT_PATH", "")
    executable=configured if configured and Path(configured).is_file() else shutil.which("tesseract")
    if not executable: return "", False
    parts=[]
    with tempfile.TemporaryDirectory(prefix="safetynet-ocr-") as temp:
        for index, frame in enumerate(ImageSequence.Iterator(image)):
            file=Path(temp)/f"frame-{index}.png"
            frame.convert("RGB").save(file)
            result=subprocess.run([executable,str(file),"stdout","--psm","3"],capture_output=True,timeout=20,check=True)
            parts.append(result.stdout.decode("utf-8"))
    return "\n".join(parts), True


@timed("parse")
def extract(data, kind):
    if kind=="text": return data.decode("utf-8"), True, []
    if kind=="image":
        text,complete=ocr_image(data)
        return text,complete,[("image",data)]
    if kind=="pdf":
        from pypdf import PdfReader
        reader=PdfReader(io.BytesIO(data),strict=True)
        if reader.is_encrypted: raise ValueError("Encrypted PDF cannot be inspected")
        if len(reader.pages)>100: raise ValueError("PDF page budget exceeded")
        text=[];complete=True;media=[]
        for page in reader.pages:
            value=page.extract_text() or ""
            text.append(value)
            images=list(page.images)
            if not value.strip() and not images: complete=False
            for image in images:
                if len(image.data)>MAX_BYTES: raise ValueError("PDF image budget exceeded")
                scanned,ok=ocr_image(image.data);text.append(scanned);complete=complete and ok
                media.append(("image",image.data))
                if len(media)>100: raise ValueError("PDF media budget exceeded")
            # Attachments, JavaScript, forms, actions and vector drawings need deeper inspection.
            if page.get("/Annots"): complete=False
        root=reader.trailer["/Root"]
        if root.get("/AcroForm") or root.get("/OpenAction") or root.get("/Names"): complete=False
        # Text extraction is useful, but does not prove all visible/vector PDF content was covered.
        return "\n".join(text),False,media
    if kind=="office":
        from defusedxml.ElementTree import fromstring
        text=[];media=[];complete=True
        with zipfile.ZipFile(io.BytesIO(data)) as archive:
            members=archive.infolist()
            if len(members)>2000 or sum(x.file_size for x in members)>MAX_BYTES:
                raise ValueError("Archive expansion budget exceeded")
            if not any(x.filename in ["word/document.xml","ppt/presentation.xml","xl/workbook.xml"] for x in members):
                raise ValueError("Only OOXML document archives are supported")
            for entry in members:
                if entry.flag_bits&1 or entry.file_size>MAX_BYTES or entry.file_size>max(entry.compress_size,1)*200:
                    raise ValueError("Unsafe document archive")
                name=entry.filename
                if any(s in name.lower() for s in ["vbaproject","embeddings/","activex/","customxml/"]): complete=False
                if name.endswith(".xml"):
                    root=fromstring(archive.read(entry))
                    text.extend(element.text for element in root.iter() if element.tag.rsplit("}",1)[-1] in ["t","v"] and element.text)
                elif "/media/" in name:
                    payload=archive.read(entry)
                    if content_kind(payload)!="image": complete=False; continue
                    scanned,ok=ocr_image(payload);text.append(scanned);complete=complete and ok;media.append(("image",payload))
                    if len(media)>100: raise ValueError("Office media budget exceeded")
        # Charts, relationships, equations and rendered layout are not yet fully inspected.
        return "\n".join(text),False,media
    if kind in ["audio","video"]: return "",False,[(kind,data)]
    raise ValueError("Unsupported content type")


@timed("semantic.load")
def get_model():
    global _MODEL
    directory=os.environ.get("EMBEDDING_MODEL_DIR", "")
    worker=os.environ.get("EMBEDDING_WORKER_PATH", "")
    if not directory or not Path(directory).is_dir() or not worker: raise ValueError("Verified local EmbeddingGemma 2 runtime is not configured")
    if _MODEL is None:
        from semantic import NativeSemanticModel
        _MODEL=NativeSemanticModel(directory,worker)
    return _MODEL


@timed("semantic")
def embed(data, kind):
    model=get_model()
    if kind=="text": vectors=model.text_vectors(data.decode("utf-8"))
    elif kind=="image": vectors=model.image_vectors(data)
    elif kind=="audio": vectors=model.audio_vectors(data)
    else: raise ValueError("Native semantic modality is not configured")
    return aggregate_vectors(vectors)


def aggregate_vectors(vectors):
    from semantic import validate_vectors,DIMENSIONS
    validate_vectors(vectors)
    vector=[sum(item[index] for item in vectors)/len(vectors) for index in range(DIMENSIONS)]
    norm=math.sqrt(sum(value*value for value in vector))
    if norm<1e-12: raise ValueError("Degenerate semantic aggregate")
    return [value/norm for value in vector]


@timed("semantic")
def semantic_vector_features(text,media):
    model=get_model();vectors=model.text_vectors(text) if text.strip() else []
    text_chunks=len(vectors)
    for kind,payload in media:
        if kind=="image": vectors.extend(model.image_vectors(payload))
        elif kind=="audio": vectors.extend(model.audio_vectors(payload))
        else: raise ValueError("Native semantic modality is not configured")
        if len(vectors)>64: raise ValueError("Semantic vector budget exceeded")
    if not vectors: raise ValueError("No semantic inputs")
    return vectors,text_chunks


def semantic_vectors(text,media):
    return semantic_vector_features(text,media)[0]


@timed("semantic")
def embedding_features(data):
    kind=content_kind(data)
    if len(data)>MAX_BYTES: raise ValueError("Embedding input budget exceeded")
    text,complete,media=extract(data,kind)
    if len(text)>MAX_TEXT: raise ValueError("Extracted text budget exceeded")
    vectors,text_chunks=semantic_vector_features(text,media)
    return {**get_model().info(),"vector":aggregate_vectors(vectors),"vectors":vectors,
            "coverage":{"complete":complete,"extractionComplete":complete,"textChunks":text_chunks,"mediaItems":len(media)}}


def deterministic(text, data, policy):
    digest=hashlib.sha256(data).hexdigest()
    normalized=unicodedata.normalize("NFKC",text).casefold()
    patterns=[("private-key",r"-----BEGIN (?:RSA |EC |OPENSSH |DSA )?PRIVATE KEY-----"),
              ("aws-access-key",r"\b(?:AKIA|ASIA)[A-Z0-9]{16}\b"),
              ("api-secret",r"\b(?:sk-[A-Za-z0-9_-]{20,}|gh[pousr]_[A-Za-z0-9]{20,})\b")]
    for rule,pattern in patterns:
        if re.search(pattern,text): return rule
    fragments=None
    for rule in policy["dlp"]["rules"]:
        if rule["kind"]=="sha256" and rule["value"].lower()==digest: return rule["id"]
        if rule["kind"]=="literal" and unicodedata.normalize("NFKC",rule["value"]).casefold() in normalized: return rule["id"]
        if rule["kind"]=="fragment-sha256" and text:
            if fragments is None: fragments=fragment_hashes(text)
            if rule["value"].lower() in fragments: return rule["id"]
    return None


@timed("inspect")
def inspect_content(data, policy, encoder=embed):
    kind=content_kind(data)
    result={"decision":"block","ruleId":"uninspectable","contentType":kind,"coverage":{"complete":False},"model":None}
    if len(data)>min(MAX_BYTES,policy["dlp"]["maxBytes"]): result["ruleId"]="size-limit";return result
    if policy.get("security",{}).get("yara",{}).get("enabled"):
        malware=scan_download(data,policy)
        if malware["decision"]!="allow": return {**result,"ruleId":malware["ruleId"],"malware":malware}
    # Fingerprints can block any bytes before invoking parsers.
    rule=deterministic("",data,policy)
    if rule: result["ruleId"]=rule;return result
    try:
        text,complete,media=extract(data,kind)
        if len(text)>MAX_TEXT: raise ValueError("Extracted text budget exceeded")
        result["coverage"]={"complete":complete,"textCharacters":len(text),"mediaItems":len(media)}
        rule=deterministic(text,data,policy)
        if rule: result["ruleId"]=rule;return result
        semantic=policy["dlp"].get("semantic",{})
        if semantic.get("enabled"):
            if encoder is embed:
                info=get_model().info()
                if any(semantic.get(key)!=info[key] for key in ["model","revision","profile","dimensions","artifactFormat","quantization","runtimeRevision"]): raise ValueError("Pinned model profile unavailable")
            elif semantic.get("model")!="google/embeddinggemma-2" or semantic.get("revision")!=os.environ.get("EMBEDDING_MODEL_REVISION"): raise ValueError("Pinned test model revision unavailable")
            if not semantic.get("categories"): raise ValueError("Semantic policy has no reference categories")
            result["model"]=semantic["model"];result["similarity"]=[]
            if encoder is embed: vectors=semantic_vectors(text,media)
            else:
                inputs=[("text",text[i:i+4000].encode()) for i in range(0,len(text),3500)]+media
                if not inputs or len(inputs)>64: raise ValueError("Inference budget exceeded")
                vectors=[encoder(payload,modality) for modality,payload in inputs]
            if sum(len(category.get("prototypes") or [category["prototype"]]) for category in semantic["categories"])>256: raise ValueError("Semantic reference budget exceeded")
            for category in semantic["categories"]:
                prototypes=category.get("prototypes") or [category["prototype"]]
                if not isinstance(category.get("threshold"),(int,float)) or not math.isfinite(category["threshold"]) or not 0<category["threshold"]<=1: raise ValueError("Invalid semantic threshold")
                scores=[cosine(vector,prototype) for vector in vectors for prototype in prototypes]
                score=max(scores);result["similarity"].append({"categoryId":category["id"],"score":round(score,6)})
                if score>=category["threshold"]: result["ruleId"]=category["id"];return result
        if complete: result["decision"]="allow";result["ruleId"]=None
        return result
    except Exception:
        # Parser/model errors never release content and never serialize raw exception text.
        result["ruleId"]="inspection-incomplete";return result