tests/test_cli.py
2.8 KB · 96 lines · python Raw
1 """Tests for the pqc-bpf CLI."""
2
3 from __future__ import annotations
4
5 import json
6
7 from click.testing import CliRunner
8
9 from pqc_ebpf_attestation.cli import main
10
11
12 def test_sign_creates_sig_json(tmp_path) -> None:
13 runner = CliRunner()
14 bpf_path = tmp_path / "trace.bpf.o"
15 bpf_path.write_bytes(b"\x7fELF" + b"\x00" * 120)
16
17 result = runner.invoke(
18 main,
19 [
20 "sign",
21 str(bpf_path),
22 "--name",
23 "trace-read",
24 "--type",
25 "kprobe",
26 "--author",
27 "ops-team",
28 ],
29 )
30 assert result.exit_code == 0, result.output
31 sig_path = tmp_path / "trace.bpf.o.sig.json"
32 assert sig_path.exists()
33 data = json.loads(sig_path.read_text())
34 assert data["program"]["metadata"]["name"] == "trace-read"
35 assert data["program"]["metadata"]["program_type"] == "kprobe"
36 assert data["signer_did"].startswith("did:pqaid:")
37
38
39 def test_verify_ok_exit_code(tmp_path) -> None:
40 runner = CliRunner()
41 bpf_path = tmp_path / "program.bpf.o"
42 bpf_path.write_bytes(b"\x7fELFpayload")
43
44 sign_res = runner.invoke(
45 main, ["sign", str(bpf_path), "--name", "p", "--type", "xdp"]
46 )
47 assert sign_res.exit_code == 0, sign_res.output
48
49 sig_path = tmp_path / "program.bpf.o.sig.json"
50 verify_res = runner.invoke(main, ["verify", str(sig_path)])
51 assert verify_res.exit_code == 0, verify_res.output
52 assert "signature VALID" in verify_res.output
53
54
55 def test_verify_fails_on_tampered(tmp_path) -> None:
56 runner = CliRunner()
57 bpf_path = tmp_path / "program.bpf.o"
58 bpf_path.write_bytes(b"original-bytes")
59
60 runner.invoke(main, ["sign", str(bpf_path), "--name", "p", "--type", "kprobe"])
61 sig_path = tmp_path / "program.bpf.o.sig.json"
62 data = json.loads(sig_path.read_text())
63 # Swap the bytecode_hash to something wrong, keeping everything else.
64 data["program"]["bytecode_hash"] = "00" * 32
65 sig_path.write_text(json.dumps(data))
66
67 verify_res = runner.invoke(main, ["verify", str(sig_path)])
68 assert verify_res.exit_code == 1
69 assert "FAIL" in verify_res.output
70
71
72 def test_info_shows_metadata(tmp_path) -> None:
73 runner = CliRunner()
74 bpf_path = tmp_path / "program.bpf.o"
75 bpf_path.write_bytes(b"bytecode-here")
76
77 runner.invoke(
78 main,
79 [
80 "sign",
81 str(bpf_path),
82 "--name",
83 "read-latency",
84 "--type",
85 "tracepoint",
86 "--author",
87 "sre",
88 ],
89 )
90 sig_path = tmp_path / "program.bpf.o.sig.json"
91 info_res = runner.invoke(main, ["info", str(sig_path)])
92 assert info_res.exit_code == 0, info_res.output
93 assert "read-latency" in info_res.output
94 assert "tracepoint" in info_res.output
95 assert "signer_did" in info_res.output
96