371 lines
19 KiB
Python
371 lines
19 KiB
Python
"""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
|