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
|