207 lines
7.9 KiB
Python
207 lines
7.9 KiB
Python
"""授权凭证签发 / 校验单元测试。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import pytest
|
||
|
||
from app.integrations.sandbox.auth import (
|
||
AuthorizationContext,
|
||
compute_body_hash,
|
||
compute_file_op_hash,
|
||
sign_authorization,
|
||
verify_authorization,
|
||
)
|
||
|
||
SECRET = "python-rust-shared-secret"
|
||
# 固定向量，与 Rust sandbox/src/auth.rs 的 body_hash_known_vector_matches_python 共享
|
||
ARGV = ["sh", "-c", "echo hi"]
|
||
STDIN: str | None = None
|
||
|
||
|
||
def _ctx() -> AuthorizationContext:
|
||
return AuthorizationContext(
|
||
trace_id="trace_1",
|
||
risk_level="L1",
|
||
tool="shell",
|
||
character_id="character_01H",
|
||
run_id="run_01H",
|
||
)
|
||
|
||
|
||
def test_sign_then_verify_roundtrip() -> None:
|
||
token = sign_authorization(SECRET, _ctx(), ARGV, STDIN, ttl_seconds=60)
|
||
payload = verify_authorization(SECRET, token)
|
||
assert payload["trace_id"] == "trace_1"
|
||
assert payload["risk_level"] == "L1"
|
||
assert payload["tool"] == "shell"
|
||
|
||
|
||
def test_verify_rejects_wrong_secret() -> None:
|
||
token = sign_authorization(SECRET, _ctx(), ARGV, STDIN)
|
||
with pytest.raises(ValueError):
|
||
verify_authorization("other-secret", token)
|
||
|
||
|
||
def test_verify_rejects_tampered_signature() -> None:
|
||
token = sign_authorization(SECRET, _ctx(), ARGV, STDIN)
|
||
tampered = token[:-2] + ("AA" if token[-2:] != "AA" else "BB")
|
||
with pytest.raises(ValueError):
|
||
verify_authorization(SECRET, tampered)
|
||
|
||
|
||
def test_verify_rejects_expired() -> None:
|
||
token = sign_authorization(SECRET, _ctx(), ARGV, STDIN, ttl_seconds=-10)
|
||
with pytest.raises(ValueError):
|
||
verify_authorization(SECRET, token)
|
||
|
||
|
||
def test_verify_rejects_malformed() -> None:
|
||
with pytest.raises(ValueError):
|
||
verify_authorization(SECRET, "not-a-valid-token")
|
||
|
||
|
||
def test_sign_requires_secret() -> None:
|
||
with pytest.raises(ValueError):
|
||
sign_authorization("", _ctx(), ARGV, STDIN)
|
||
|
||
|
||
def test_payload_carries_body_hash() -> None:
|
||
token = sign_authorization(SECRET, _ctx(), ARGV, STDIN)
|
||
payload = verify_authorization(SECRET, token)
|
||
assert payload["body_hash"] == compute_body_hash(ARGV, STDIN)
|
||
|
||
|
||
def test_body_hash_differs_per_argv() -> None:
|
||
"""凭证绑死具体执行内容：不同 argv 产生不同 body_hash，杜绝重放任意命令。"""
|
||
t1 = sign_authorization(SECRET, _ctx(), ["sh", "-c", "echo a"], None)
|
||
t2 = sign_authorization(SECRET, _ctx(), ["sh", "-c", "echo b"], None)
|
||
p1 = verify_authorization(SECRET, t1)
|
||
p2 = verify_authorization(SECRET, t2)
|
||
assert p1["body_hash"] != p2["body_hash"]
|
||
|
||
|
||
def test_compute_body_hash_known_vector() -> None:
|
||
"""与 Rust compute_body_hash 共享固定向量，保证跨端字节级一致。"""
|
||
assert (
|
||
compute_body_hash(ARGV, STDIN)
|
||
== "f5fa97ea2ee92c505174672da52c59c952e49c32fa56a1db5bd983c10a829540"
|
||
)
|
||
|
||
|
||
def test_compute_body_hash_none_vs_empty_stdin() -> None:
|
||
assert compute_body_hash(["sh"], None) != compute_body_hash(["sh"], "")
|
||
|
||
|
||
def test_compute_file_op_hash_known_vector() -> None:
|
||
"""与 Rust compute_file_op_hash 共享固定向量，保证跨端字节级一致。"""
|
||
assert (
|
||
compute_file_op_hash("read", "data/note.txt")
|
||
== "717f67252c61cf1681cb3c6018d3801b79a511564664567b57d12d775bb8fa53"
|
||
)
|
||
|
||
|
||
def test_compute_file_op_hash_differs_per_op_and_path() -> None:
|
||
"""凭证绑死具体文件操作：op 或 path 不同则摘要不同。"""
|
||
assert compute_file_op_hash("read", "a.txt") != compute_file_op_hash("write", "a.txt")
|
||
assert compute_file_op_hash("read", "a.txt") != compute_file_op_hash("read", "b.txt")
|
||
|
||
|
||
def test_sign_carries_file_op_hash_only_for_file_ops() -> None:
|
||
"""exec 凭证 file_op_hash 为 None；文件凭证携带 {op, path} 摘要。"""
|
||
# exec 路径：argv 非 None → body_hash 绑定，file_op_hash 为 None
|
||
exec_token = sign_authorization(SECRET, _ctx(), ARGV, STDIN)
|
||
exec_payload = verify_authorization(SECRET, exec_token)
|
||
assert exec_payload["file_op_hash"] is None
|
||
assert exec_payload["body_hash"] == compute_body_hash(ARGV, STDIN)
|
||
# 文件路径：argv=None → body_hash 为 None，file_op_hash 绑定 {op, path}
|
||
file_token = sign_authorization(SECRET, _ctx(), None, None, file_op=("read", "data/note.txt"))
|
||
file_payload = verify_authorization(SECRET, file_token)
|
||
assert file_payload["body_hash"] is None
|
||
assert file_payload["file_op_hash"] == compute_file_op_hash("read", "data/note.txt")
|
||
|
||
|
||
def test_sign_carries_sorted_rlimit_overrides() -> None:
|
||
"""rlimit_overrides 随凭证下发，且按 key 排序（对齐 Rust serde BTreeMap 字母序）。"""
|
||
ctx = AuthorizationContext(
|
||
trace_id="trace_1",
|
||
risk_level="L1",
|
||
tool="python_runner",
|
||
# 反序插入，验证 sign 时被排序
|
||
rlimit_overrides={"fsize": 67_108_864, "as": 536_870_912},
|
||
)
|
||
token = sign_authorization(SECRET, ctx, ARGV, STDIN)
|
||
payload = verify_authorization(SECRET, token)
|
||
# 排序后 as 在 fsize 之前（字母序），与 Rust BTreeMap 序列化一致
|
||
assert payload["rlimit_overrides"] == {"as": 536_870_912, "fsize": 67_108_864}
|
||
# rlimit_overrides 属授权属性，不进 body_hash
|
||
assert payload["body_hash"] == compute_body_hash(ARGV, STDIN)
|
||
|
||
|
||
def test_sign_without_rlimit_overrides_is_null() -> None:
|
||
"""无 rlimit_overrides 时 payload 字段为 null（对齐 Rust Option None → serde null）。"""
|
||
token = sign_authorization(SECRET, _ctx(), ARGV, STDIN)
|
||
payload = verify_authorization(SECRET, token)
|
||
assert payload["rlimit_overrides"] is None
|
||
|
||
|
||
def test_compute_body_hash_code_matches_known_vectors() -> None:
|
||
"""与 Rust compute_body_hash_code 共享固定向量，保证跨端字节级一致。"""
|
||
from app.integrations.sandbox.auth import compute_body_hash_code
|
||
|
||
assert (
|
||
compute_body_hash_code("python", "print(1+1)", None)
|
||
== "e7d633605676c3b19c642ee2437d91f9fb22562aaffad40f217d84c587322b69"
|
||
)
|
||
assert (
|
||
compute_body_hash_code("python", "print(1+1)", "42")
|
||
== "f933764421a59d1578697189fb76374511c5d7e4f1ad03ad37da3241568c35d5"
|
||
)
|
||
assert (
|
||
compute_body_hash_code("node", "print(1+1)", None)
|
||
== "a0440aef4bbcf83d89127ddf89949019b6be939a1c02702d22b35b40b21862aa"
|
||
)
|
||
|
||
|
||
def test_sign_authorization_code_mode_carries_isolation_mode() -> None:
|
||
"""code 模式：body_hash 按 {language,code,stdin} 算；payload 末尾含 isolation_mode。"""
|
||
from app.integrations.sandbox.auth import (
|
||
AuthorizationContext,
|
||
compute_body_hash_code,
|
||
sign_authorization,
|
||
verify_authorization,
|
||
)
|
||
|
||
ctx = AuthorizationContext(
|
||
trace_id="trace_1",
|
||
risk_level="L3",
|
||
tool="shell",
|
||
isolation_mode="container",
|
||
)
|
||
token = sign_authorization(
|
||
"shared-secret",
|
||
ctx,
|
||
code="print(1+1)",
|
||
language="python",
|
||
stdin=None,
|
||
)
|
||
payload = verify_authorization("shared-secret", token)
|
||
assert list(payload)[-1] == "isolation_mode"
|
||
assert payload["isolation_mode"] == "container"
|
||
assert payload["body_hash"] == compute_body_hash_code("python", "print(1+1)", None)
|
||
|
||
|
||
def test_sign_authorization_argv_mode_backward_compatible_positional() -> None:
|
||
"""既有位置参数调用不受影响：argv/stdin 位置传、ttl_seconds 关键字传，仍可验签。"""
|
||
import base64
|
||
|
||
from app.integrations.sandbox.auth import AuthorizationContext, sign_authorization
|
||
|
||
ctx = AuthorizationContext(trace_id="t", risk_level="L1", tool="shell")
|
||
# 既有调用形式：argv/stdin 位置参数（与仓库现有 test_auth.py 一致）
|
||
token = sign_authorization("s", ctx, ["python3", "-c", "print(1)"], None)
|
||
payload_b64 = token.split(".", 1)[0]
|
||
pad = "=" * (-len(payload_b64) % 4)
|
||
payload_json = base64.urlsafe_b64decode(payload_b64 + pad).decode("utf-8")
|
||
# 末尾追加 isolation_mode（旧字段顺序不变），None 时为 null
|
||
assert payload_json.endswith('"isolation_mode":null}')
|