Files
lserp_cs_6.0/插件库/astrbot_plugin_lserp/attachment_provenance.py
2026-08-14 14:28:28 +08:00

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
)
)