from __future__ import annotations import hashlib import hmac import secrets from datetime import datetime, timezone from typing import Any, Callable, Mapping from urllib.parse import urlsplit, urlunsplit from .constants import EXTENSION_ID, EXTENSION_NAME, EXTENSION_VERSION from .durable import DurabilityError, PendingStore from .identity import VerifiedLocalIdentity from .strict_json import canonical_bytes class ProtocolError(RuntimeError): def __init__(self, code: str, message: str) -> None: super().__init__(message) self.code = code def _timestamp(value: str) -> datetime: try: parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError as exc: raise ProtocolError("E_DEADLINE", "deadline is invalid") from exc if parsed.tzinfo is None: raise ProtocolError("E_DEADLINE", "deadline must be offset-aware") return parsed.astimezone(timezone.utc) def _sign(secret: bytes, frame: Mapping[str, Any]) -> str: return hmac.new(secret, canonical_bytes(dict(frame)), hashlib.sha256).hexdigest() def _target_url(value: str) -> str: parsed = urlsplit(value) parts = [part for part in parsed.path.split("/") if part] if ( parsed.scheme != "https" or parsed.hostname != "space.bilibili.com" or parsed.query or parsed.fragment or len(parts) != 2 or not parts[0].isdigit() or parts[0].startswith("0") or parts[1] != "dynamic" ): raise ProtocolError("E_PAGE_IDENTITY", "target URL is not a canonical creator dynamic page") return urlunsplit(("https", "space.bilibili.com", f"/{parts[0]}/dynamic", "", "")) class HostSession: def __init__( self, identity: VerifiedLocalIdentity, pending: PendingStore, commit_observation: Callable[[Mapping[str, Any]], Mapping[str, Any]], target_url: str, *, now: Callable[[], datetime] | None = None, ) -> None: if type(identity) is not VerifiedLocalIdentity: raise ProtocolError("E_TRUSTED_ADAPTER_UNAVAILABLE", "verified local identity is required") self.identity = identity self.pending = pending self.commit_observation = commit_observation self.target_url = _target_url(target_url) self.now = now or (lambda: datetime.now(timezone.utc)) self.secret = secrets.token_bytes(32) self.run_id: str | None = None self.request_id: str | None = None self.deadline_at: str | None = None self.sequence = 0 self.action: dict[str, Any] | None = None self.permit_id: str | None = None self.action_result: Mapping[str, Any] | None = None def start(self, hello: Mapping[str, Any], *, run_id: str, request_id: str, deadline_at: str) -> dict[str, Any]: if set(hello) != {"schema_version", "type", "sequence", "extension_id", "version", "manifest_name"}: raise ProtocolError("E_LOCAL_EXTENSION_IDENTITY", "extension hello shape differs") if ( hello["schema_version"] != 1 or hello["type"] != "EXTENSION_HELLO" or hello["sequence"] != 1 or hello["extension_id"] != EXTENSION_ID or hello["version"] != EXTENSION_VERSION or hello["manifest_name"] != EXTENSION_NAME ): raise ProtocolError("E_LOCAL_EXTENSION_IDENTITY", "extension hello identity differs") deadline = _timestamp(deadline_at) if self.now().astimezone(timezone.utc) >= deadline: raise ProtocolError("E_OVERALL_DEADLINE", "original deadline expired before challenge") self.run_id, self.request_id, self.deadline_at, self.sequence = run_id, request_id, deadline_at, 2 self.pending.initialize(run_id=run_id, request_id=request_id, deadline_at=deadline_at) return self._host_frame("HOST_CHALLENGE", 2, {"challenge_id": secrets.token_hex(16), "secret": self.secret.hex()}) def accept(self, frame: Mapping[str, Any]) -> dict[str, Any]: self._verify_extension(frame) kind = frame["type"] if kind == "EXTENSION_CHALLENGE_ACCEPTED" and frame["sequence"] == 3: self.sequence = 4 self.action = {"action_id": secrets.token_hex(16), "kind": "reload", "url": self.target_url} return self._host_frame("HOST_ACTION_PREPARE", 4, {"action": self.action}) if kind == "EXTENSION_READY_TO_DISPATCH" and frame["sequence"] == 5: if self.action is None or frame.get("action_id") != self.action["action_id"]: raise ProtocolError("E_PROTOCOL", "prepared action identity differs") if self.now().astimezone(timezone.utc) >= _timestamp(str(self.deadline_at)): raise ProtocolError("E_OVERALL_DEADLINE", "original deadline expired before dispatch") candidate_permit_id = secrets.token_hex(16) permit = {"permit_id": candidate_permit_id, "action_id": self.action["action_id"], "deadline_at": self.deadline_at} permit_payload = canonical_bytes(permit) current = self.pending.load() if current is None: raise DurabilityError("E_DISPATCH_DURABILITY_AMBIGUOUS", "pending disappeared") self.pending.consume(current, permit_id=candidate_permit_id, permit_payload=permit_payload) if self.now().astimezone(timezone.utc) >= _timestamp(str(self.deadline_at)): raise ProtocolError("E_OVERALL_DEADLINE", "original deadline expired after durable budget") self.permit_id = candidate_permit_id self.sequence = 6 return self._host_frame("HOST_DISPATCH_PERMIT", 6, {"permit": permit}) if kind == "EXTENSION_ACTION_RESULT" and frame["sequence"] == 7: if self.permit_id is None or frame.get("permit_id") != self.permit_id: raise ProtocolError("E_PROTOCOL", "action result permit differs") self.action_result = frame.get("result") self.sequence = 8 return self._host_frame("HOST_OBSERVATION_REQUEST", 8, {"action_id": self.action["action_id"]}) if kind in {"EXTENSION_OBSERVATION", "EXTENSION_TERMINAL_ERROR"} and frame["sequence"] == 9: if self.action_result is None: raise ProtocolError("E_ACTION_RESULT_REQUIRED", "observation preceded action result") observation = frame.get("observation") if kind == "EXTENSION_OBSERVATION" else {"terminal_error": frame.get("error_code")} result = dict(self.commit_observation({ "run_id": self.run_id, "request_id": self.request_id, "deadline_at": self.deadline_at, "action": self.action, "action_result": self.action_result, "observation": observation, "trusted_identity": { "extension_id": self.identity.extension_id, "source_manifest_sha256": self.identity.source_manifest_sha256, "payload_tree_sha256": self.identity.payload_tree_sha256, "host_install_receipt_sha256": self.identity.host_install_receipt_sha256, }, })) self.sequence = 10 return self._host_frame("HOST_COMMIT_RESULT", 10, {"result": result}) raise ProtocolError("E_SEQUENCE", "message type or sequence differs") def _verify_extension(self, frame: Mapping[str, Any]) -> None: required = {"schema_version", "type", "run_id", "request_id", "sequence", "hmac"} if not isinstance(frame, dict) or not required.issubset(frame): raise ProtocolError("E_PROTOCOL", "extension frame is incomplete") if frame["run_id"] != self.run_id or frame["request_id"] != self.request_id: raise ProtocolError("E_PROTOCOL", "run or request identity differs") unsigned = dict(frame) supplied = unsigned.pop("hmac") if not isinstance(supplied, str) or not hmac.compare_digest(supplied, _sign(self.secret, unsigned)): raise ProtocolError("E_HMAC", "extension frame HMAC differs") def _host_frame(self, kind: str, sequence: int, fields: Mapping[str, Any]) -> dict[str, Any]: frame = { "schema_version": 1, "type": kind, "run_id": self.run_id, "request_id": self.request_id, "sequence": sequence, **dict(fields), } frame["hmac"] = _sign(self.secret, frame) return frame def sign_extension_frame(secret_hex: str, frame: Mapping[str, Any]) -> dict[str, Any]: """Test/extension parity helper; production main never exposes the session secret.""" unsigned = dict(frame) unsigned["hmac"] = _sign(bytes.fromhex(secret_hex), unsigned) return unsigned