359 lines
12 KiB
Python
359 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import hmac
|
|
import json
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Callable
|
|
|
|
from .attachment_extract import (
|
|
AttachmentProvenanceError,
|
|
MAX_SOURCE_BYTES,
|
|
PREPROCESS_CONTRACTS,
|
|
bind_extraction_receipt,
|
|
untrusted_image_payload,
|
|
)
|
|
from .purchase_vision_binding import normalize_purchase_vision_documents
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VerifiedAttachmentBundle:
|
|
"""Immutable, one-use attachment evidence released to the ERP planning Tool."""
|
|
|
|
receipts: tuple[tuple[str, str, str, int, str, str], ...]
|
|
purchase_vision_documents: tuple[str, ...]
|
|
purchase_file_documents: tuple[str, ...]
|
|
|
|
def source_documents(self) -> list[dict[str, object]]:
|
|
return [
|
|
{
|
|
"kind": kind,
|
|
"filename": filename,
|
|
"sha256": sha256,
|
|
"sizeBytes": size_bytes,
|
|
"extractionSha256": extraction_sha256,
|
|
"preprocessContract": preprocess_contract,
|
|
}
|
|
for (
|
|
kind,
|
|
filename,
|
|
sha256,
|
|
size_bytes,
|
|
extraction_sha256,
|
|
preprocess_contract,
|
|
)
|
|
in self.receipts
|
|
]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ReceiptState:
|
|
recorded_at: float
|
|
receipts: tuple[tuple[str, str, str, int, str, str], ...]
|
|
purchase_vision_documents: tuple[str, ...]
|
|
purchase_file_documents: tuple[str, ...]
|
|
error_code: str | None
|
|
|
|
|
|
class AttachmentReceiptStateStore:
|
|
"""One-use, bounded handoff from attachment preprocessing to the ERP Tool."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
maximum_sessions: int = 128,
|
|
lifetime_seconds: int = 300,
|
|
time_source: Callable[[], float] = time.monotonic,
|
|
) -> None:
|
|
if maximum_sessions < 1 or maximum_sessions > 1024:
|
|
raise ValueError("attachment_receipt_state_configuration_invalid")
|
|
if lifetime_seconds < 10 or lifetime_seconds > 900:
|
|
raise ValueError("attachment_receipt_state_configuration_invalid")
|
|
self._maximum_sessions = maximum_sessions
|
|
self._lifetime_seconds = lifetime_seconds
|
|
self._time_source = time_source
|
|
self._states: dict[str, _ReceiptState] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def record(
|
|
self,
|
|
session_id: str,
|
|
receipts: list[dict[str, object]],
|
|
error_code: str | None = None,
|
|
purchase_vision_documents: list[str] | None = None,
|
|
purchase_file_documents: list[str] | None = None,
|
|
) -> None:
|
|
key = _session_key(session_id)
|
|
normalized = _normalize_bound_receipts(receipts)
|
|
vision_documents = normalize_purchase_vision_documents(
|
|
purchase_vision_documents or []
|
|
)
|
|
file_documents = _normalize_file_documents(
|
|
purchase_file_documents or []
|
|
)
|
|
_verify_extraction_receipt_links(
|
|
normalized,
|
|
vision_documents,
|
|
file_documents,
|
|
)
|
|
if error_code is not None and not _safe_code(error_code):
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
now = self._time_source()
|
|
with self._lock:
|
|
self._prune(now)
|
|
if key not in self._states and len(self._states) >= self._maximum_sessions:
|
|
oldest = min(
|
|
self._states,
|
|
key=lambda item: self._states[item].recorded_at,
|
|
)
|
|
del self._states[oldest]
|
|
self._states[key] = _ReceiptState(
|
|
now,
|
|
normalized,
|
|
vision_documents,
|
|
file_documents,
|
|
error_code,
|
|
)
|
|
|
|
def consume_verified(
|
|
self,
|
|
session_id: str,
|
|
current_receipts: list[dict[str, object]],
|
|
) -> VerifiedAttachmentBundle:
|
|
key = _session_key(session_id)
|
|
current = _normalize_source_receipts(current_receipts)
|
|
now = self._time_source()
|
|
with self._lock:
|
|
self._prune(now)
|
|
state = self._states.pop(key, None)
|
|
if state is None:
|
|
raise AttachmentProvenanceError(
|
|
"attachment_provenance_state_missing"
|
|
)
|
|
if state.error_code:
|
|
raise AttachmentProvenanceError(state.error_code)
|
|
if tuple(item[:4] for item in state.receipts) != current:
|
|
raise AttachmentProvenanceError(
|
|
"attachment_changed_after_preprocess"
|
|
)
|
|
return VerifiedAttachmentBundle(
|
|
state.receipts,
|
|
state.purchase_vision_documents,
|
|
state.purchase_file_documents,
|
|
)
|
|
|
|
def _prune(self, now: float) -> None:
|
|
expired = [
|
|
key
|
|
for key, state in self._states.items()
|
|
if now - state.recorded_at > self._lifetime_seconds
|
|
]
|
|
for key in expired:
|
|
del self._states[key]
|
|
|
|
|
|
def _normalize_source_receipts(
|
|
receipts: list[dict[str, object]],
|
|
) -> tuple[tuple[str, str, str, int], ...]:
|
|
if not isinstance(receipts, list) or len(receipts) > 3:
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
result: list[tuple[str, str, str, int]] = []
|
|
for receipt in receipts:
|
|
if not isinstance(receipt, dict) or set(receipt) != {
|
|
"kind",
|
|
"filename",
|
|
"sha256",
|
|
"sizeBytes",
|
|
}:
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
kind = receipt.get("kind")
|
|
filename = receipt.get("filename")
|
|
sha256 = receipt.get("sha256")
|
|
size_bytes = receipt.get("sizeBytes")
|
|
if (
|
|
kind not in {"image", "file"}
|
|
or not isinstance(filename, str)
|
|
or not filename
|
|
or len(filename) > 128
|
|
or not isinstance(sha256, str)
|
|
or len(sha256) != 64
|
|
or any(item not in "0123456789abcdef" for item in sha256)
|
|
or isinstance(size_bytes, bool)
|
|
or not isinstance(size_bytes, int)
|
|
or size_bytes <= 0
|
|
or size_bytes > MAX_SOURCE_BYTES
|
|
):
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
result.append((kind, filename, sha256, size_bytes))
|
|
return tuple(result)
|
|
|
|
|
|
def _normalize_bound_receipts(
|
|
receipts: list[dict[str, object]],
|
|
) -> tuple[tuple[str, str, str, int, str, str], ...]:
|
|
if not isinstance(receipts, list) or len(receipts) > 3:
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
result: list[tuple[str, str, str, int, str, str]] = []
|
|
for receipt in receipts:
|
|
if not isinstance(receipt, dict) or set(receipt) != {
|
|
"kind",
|
|
"filename",
|
|
"sha256",
|
|
"sizeBytes",
|
|
"extractionSha256",
|
|
"preprocessContract",
|
|
}:
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
source = {
|
|
"kind": receipt.get("kind"),
|
|
"filename": receipt.get("filename"),
|
|
"sha256": receipt.get("sha256"),
|
|
"sizeBytes": receipt.get("sizeBytes"),
|
|
}
|
|
normalized = _normalize_source_receipts([source])[0]
|
|
extraction_sha256 = receipt.get("extractionSha256")
|
|
preprocess_contract = receipt.get("preprocessContract")
|
|
if (
|
|
not isinstance(extraction_sha256, str)
|
|
or len(extraction_sha256) != 64
|
|
or any(item not in "0123456789abcdef" for item in extraction_sha256)
|
|
or preprocess_contract not in PREPROCESS_CONTRACTS
|
|
):
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
result.append(normalized + (extraction_sha256, preprocess_contract))
|
|
return tuple(result)
|
|
|
|
|
|
def _verify_extraction_receipt_links(
|
|
receipts: tuple[tuple[str, str, str, int, str, str], ...],
|
|
vision_documents: tuple[str, ...],
|
|
file_documents: tuple[str, ...],
|
|
) -> None:
|
|
image_receipts = [item for item in receipts if item[0] == "image"]
|
|
file_receipts = [item for item in receipts if item[0] == "file"]
|
|
if len(image_receipts) != len(vision_documents):
|
|
raise ValueError("purchase_vision_binding_invalid")
|
|
if len(file_receipts) != len(file_documents):
|
|
raise ValueError("purchase_file_binding_invalid")
|
|
for receipt, document in zip(image_receipts, vision_documents):
|
|
(
|
|
kind,
|
|
filename,
|
|
sha256,
|
|
size_bytes,
|
|
extraction_sha256,
|
|
preprocess_contract,
|
|
) = receipt
|
|
source = {
|
|
"kind": kind,
|
|
"filename": filename,
|
|
"sha256": sha256,
|
|
"sizeBytes": size_bytes,
|
|
}
|
|
payload = untrusted_image_payload(filename, document, source)
|
|
expected = bind_extraction_receipt(source, payload)
|
|
if (
|
|
not hmac.compare_digest(
|
|
str(expected["extractionSha256"]), extraction_sha256
|
|
)
|
|
or expected["preprocessContract"] != preprocess_contract
|
|
):
|
|
raise ValueError("purchase_vision_binding_invalid")
|
|
for receipt, payload in zip(file_receipts, file_documents):
|
|
(
|
|
kind,
|
|
filename,
|
|
sha256,
|
|
size_bytes,
|
|
extraction_sha256,
|
|
preprocess_contract,
|
|
) = receipt
|
|
source = {
|
|
"kind": kind,
|
|
"filename": filename,
|
|
"sha256": sha256,
|
|
"sizeBytes": size_bytes,
|
|
}
|
|
if _file_payload_source(payload) != (filename, source):
|
|
raise ValueError("purchase_file_binding_invalid")
|
|
expected = bind_extraction_receipt(source, payload)
|
|
if (
|
|
not hmac.compare_digest(
|
|
str(expected["extractionSha256"]), extraction_sha256
|
|
)
|
|
or expected["preprocessContract"] != preprocess_contract
|
|
):
|
|
raise ValueError("purchase_file_binding_invalid")
|
|
|
|
|
|
def _normalize_file_documents(documents: list[str]) -> tuple[str, ...]:
|
|
if not isinstance(documents, list) or len(documents) > 3:
|
|
raise ValueError("purchase_file_binding_invalid")
|
|
result: list[str] = []
|
|
for document in documents:
|
|
if (
|
|
not isinstance(document, str)
|
|
or not document.startswith("UNTRUSTED_BUSINESS_ATTACHMENT_JSON=")
|
|
or len(document.encode("utf-8", errors="strict")) > 1024 * 1024
|
|
):
|
|
raise ValueError("purchase_file_binding_invalid")
|
|
result.append(document)
|
|
return tuple(result)
|
|
|
|
|
|
def _file_payload_source(
|
|
payload: str,
|
|
) -> tuple[str, dict[str, object]]:
|
|
marker = "UNTRUSTED_BUSINESS_ATTACHMENT_JSON="
|
|
try:
|
|
envelope = json.loads(
|
|
payload[len(marker):],
|
|
object_pairs_hook=_unique_object,
|
|
parse_constant=_reject_json_constant,
|
|
)
|
|
except (TypeError, ValueError, json.JSONDecodeError) as error:
|
|
raise ValueError("purchase_file_binding_invalid") from error
|
|
if (
|
|
not isinstance(envelope, dict)
|
|
or set(envelope) != {"kind", "filename", "sourceDocument", "content"}
|
|
or envelope.get("kind") != "file"
|
|
or not isinstance(envelope.get("filename"), str)
|
|
or not isinstance(envelope.get("content"), str)
|
|
or not isinstance(envelope.get("sourceDocument"), dict)
|
|
):
|
|
raise ValueError("purchase_file_binding_invalid")
|
|
return envelope["filename"], envelope["sourceDocument"]
|
|
|
|
|
|
def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
|
|
result: dict[str, object] = {}
|
|
for key, value in pairs:
|
|
if key in result:
|
|
raise ValueError("duplicate_json_property")
|
|
result[key] = value
|
|
return result
|
|
|
|
|
|
def _reject_json_constant(_: str) -> None:
|
|
raise ValueError("non_finite_json_number")
|
|
|
|
|
|
def _session_key(value: str) -> str:
|
|
key = (value or "").strip()
|
|
if not key or len(key) > 512 or any(character.isspace() for character in key):
|
|
raise ValueError("attachment_receipt_state_invalid")
|
|
return key
|
|
|
|
|
|
def _safe_code(value: str) -> bool:
|
|
return (
|
|
1 <= len(value) <= 128
|
|
and all(
|
|
character.islower()
|
|
or character.isdigit()
|
|
or character in "_.-"
|
|
for character in value
|
|
)
|
|
)
|