157 lines
7.1 KiB
Python
157 lines
7.1 KiB
Python
"""司命授权凭证：HMAC-SHA256 短期 token，与 Rust 沙箱端 ``auth.rs`` 共用算法与密钥。
|
||
|
||
凭证格式：``<payload_b64url>.<sig_b64url>``
|
||
|
||
- payload：JSON，字段顺序固定为
|
||
``exp / trace_id / risk_level / tool / character_id / run_id / body_hash / rlimit_overrides / file_op_hash / isolation_mode``，
|
||
必须与 Rust ``serde`` 字段顺序字节级一致（否则 HMAC 不匹配）。``isolation_mode`` 末尾追加，旧凭证向后兼容。
|
||
- sig：``HMAC-SHA256(secret, payload_b64url)`` 的 base64url（无填充）。
|
||
- body_hash：``compute_body_hash(argv, stdin)``，把凭证绑死到具体执行内容，防 TTL 内重放执行任意命令。
|
||
|
||
安全语义：客户端不判断风险，只把「经司命审核通过」的上下文连同请求内容摘要签成短期凭证；
|
||
沙箱端验签 + body_hash 比对通过才执行。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import hmac
|
||
import json
|
||
import time
|
||
from dataclasses import dataclass
|
||
from typing import cast
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class AuthorizationContext:
|
||
"""一次授权的上下文：由司命审核通过后随请求下发。"""
|
||
|
||
trace_id: str
|
||
risk_level: str # L0–L4，见契约字典 risk_level
|
||
tool: str
|
||
character_id: str | None = None
|
||
run_id: str | None = None
|
||
# 司命授予的资源配额上限（key: fsize/as/cpu/nproc/nofile → u64）。
|
||
# 随凭证下发，sandbox 验签后信任并 clamp 到全局天花板；None 表示不绑定分级配额。
|
||
# 属授权属性，不进 body_hash（body_hash 只绑 argv + stdin）。
|
||
rlimit_overrides: dict[str, int] | None = None
|
||
# 隔离模式授权（L0-L2 → process，L3-L4 → container）。随凭证下发，沙箱据此路由；
|
||
# None 表示司命未指定。属授权属性，不进 body_hash，与 rlimit_overrides 同链路。
|
||
isolation_mode: str | None = None
|
||
|
||
|
||
def _b64url_no_pad(raw: bytes) -> str:
|
||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")
|
||
|
||
|
||
def _b64url_no_pad_decode(text: str) -> bytes:
|
||
pad = "=" * (-len(text) % 4)
|
||
return base64.urlsafe_b64decode(text + pad)
|
||
|
||
|
||
def compute_body_hash(argv: list[str], stdin: str | None) -> str:
|
||
"""计算请求内容摘要，与 Rust 沙箱端 ``compute_body_hash`` 字节级一致。
|
||
|
||
口径：``SHA256(canonical JSON {"argv": [...], "stdin": null|str})``——
|
||
JSON 紧凑（``separators=(",",":")``）、不转义非 ASCII（``ensure_ascii=False``）、
|
||
字段顺序固定 ``argv, stdin``；返回小写十六进制（64 字符）。
|
||
"""
|
||
payload = {"argv": argv, "stdin": stdin}
|
||
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||
return hashlib.sha256(encoded).hexdigest()
|
||
|
||
|
||
def compute_file_op_hash(op: str, path: str) -> str:
|
||
"""计算文件操作摘要，与 Rust 沙箱端 ``compute_file_op_hash`` 字节级一致。
|
||
|
||
口径：``SHA256(canonical JSON {"op": str, "path": str})``——把凭证绑死到具体文件操作
|
||
（read/write/list/stat），防 TTL 内重放执行任意文件读写。字段顺序固定 ``op, path``。
|
||
"""
|
||
payload = {"op": op, "path": path}
|
||
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||
return hashlib.sha256(encoded).hexdigest()
|
||
|
||
|
||
def compute_body_hash_code(language: str, code: str, stdin: str | None) -> str:
|
||
"""计算 code 模式请求内容摘要，与 Rust 沙箱端 ``compute_body_hash_code`` 字节级一致。
|
||
|
||
口径：``SHA256(canonical JSON {"language": str, "code": str, "stdin": null|str})``——
|
||
字段顺序固定 ``language, code, stdin``、紧凑 JSON、不转义非 ASCII；返回小写十六进制。
|
||
code 模式凭证 body_hash 用此函数（argv 模式仍用 ``compute_body_hash``），把凭证绑死到
|
||
具体代码内容，防 TTL 内重放执行任意代码。
|
||
"""
|
||
payload = {"language": language, "code": code, "stdin": stdin}
|
||
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||
return hashlib.sha256(encoded).hexdigest()
|
||
|
||
|
||
def sign_authorization(
|
||
secret: str,
|
||
ctx: AuthorizationContext,
|
||
argv: list[str] | None = None,
|
||
stdin: str | None = None,
|
||
ttl_seconds: int = 30,
|
||
*,
|
||
code: str | None = None,
|
||
language: str | None = None,
|
||
file_op: tuple[str, str] | None = None,
|
||
) -> str:
|
||
"""签发授权凭证。
|
||
|
||
内容绑定（互斥，由调用方保证不同时传 argv 与 code）：
|
||
- exec 命令模式：传 ``argv``/``stdin``，按 ``compute_body_hash(argv, stdin)`` 绑定。
|
||
- exec 代码模式：传 ``code``/``language``/``stdin``，按 ``compute_body_hash_code`` 绑定。
|
||
- 低风险操作（cancel）：都不传，不绑定 body_hash（file_op 单独走 file_op_hash）。
|
||
|
||
payload 字段顺序与 Rust ``AuthPayload`` 一致（isolation_mode 末尾追加，旧凭证向后兼容）。
|
||
"""
|
||
if not secret:
|
||
raise ValueError("授权密钥不能为空")
|
||
exp = int(time.time()) + ttl_seconds
|
||
if code is not None:
|
||
body_hash = compute_body_hash_code(language or "python", code, stdin)
|
||
elif argv is not None:
|
||
body_hash = compute_body_hash(argv, stdin)
|
||
else:
|
||
body_hash = None
|
||
payload: dict[str, object] = {
|
||
"exp": exp,
|
||
"trace_id": ctx.trace_id,
|
||
"risk_level": ctx.risk_level,
|
||
"tool": ctx.tool,
|
||
"character_id": ctx.character_id,
|
||
"run_id": ctx.run_id,
|
||
"body_hash": body_hash,
|
||
# rlimit_overrides 内部按 key 排序，对齐 Rust serde BTreeMap 的字母序，
|
||
# 保证多 key 时两端 payload JSON 字节一致 → HMAC 一致。
|
||
"rlimit_overrides": (
|
||
dict(sorted(ctx.rlimit_overrides.items())) if ctx.rlimit_overrides else None
|
||
),
|
||
# file_op_hash 绑定文件操作 {op, path}（exec 路径为 None）；字段顺序对齐 Rust AuthPayload。
|
||
"file_op_hash": compute_file_op_hash(file_op[0], file_op[1]) if file_op else None,
|
||
# isolation_mode 末尾追加（字段顺序对齐 Rust AuthPayload；旧凭证无此字段仍可验签）。
|
||
"isolation_mode": ctx.isolation_mode,
|
||
}
|
||
payload_json = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
|
||
payload_b64 = _b64url_no_pad(payload_json.encode("utf-8"))
|
||
digest = hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
|
||
return f"{payload_b64}.{_b64url_no_pad(digest)}"
|
||
|
||
|
||
def verify_authorization(secret: str, token: str) -> dict[str, object]:
|
||
"""校验授权凭证（对称实现，供本地测试与跨端一致性校验）。"""
|
||
payload_b64, sep, sig_b64 = token.rpartition(".")
|
||
if not sep:
|
||
raise ValueError("授权凭证格式不合法")
|
||
expected = _b64url_no_pad(
|
||
hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
|
||
)
|
||
if not hmac.compare_digest(expected, sig_b64):
|
||
raise ValueError("授权凭证签名校验失败")
|
||
payload = cast("dict[str, object]", json.loads(_b64url_no_pad_decode(payload_b64)))
|
||
exp = payload.get("exp")
|
||
if not isinstance(exp, int | float) or exp <= time.time():
|
||
raise ValueError("授权凭证已过期")
|
||
return payload
|