tests/test_signer.py
3.8 KB · 103 lines · python Raw
1 """Tests for MessageSigner — signing, verification, and canonicalization."""
2
3 from __future__ import annotations
4
5 import copy
6
7 import pytest
8 from quantumshield.identity.agent import AgentIdentity
9
10 from pqc_mcp_transport.signer import MessageSigner, VerificationResult
11
12
13 class TestCanonicalize:
14 def test_canonicalize_deterministic(self, sample_tool_call: dict) -> None:
15 """Canonicalization of the same message always yields the same bytes."""
16 a = MessageSigner.canonicalize(sample_tool_call)
17 b = MessageSigner.canonicalize(sample_tool_call)
18 assert a == b
19
20 # Insertion order shouldn't matter
21 reordered = dict(reversed(list(sample_tool_call.items())))
22 c = MessageSigner.canonicalize(reordered)
23 assert a == c
24
25 def test_canonicalize_strips_pqc(self, sample_tool_call: dict) -> None:
26 """The _pqc field is excluded from the canonical form."""
27 with_pqc = dict(sample_tool_call)
28 with_pqc["_pqc"] = {"signature": "deadbeef"}
29 canon_with = MessageSigner.canonicalize(with_pqc)
30 canon_without = MessageSigner.canonicalize(sample_tool_call)
31 assert canon_with == canon_without
32
33
34 class TestSignAndVerify:
35 def test_sign_message_adds_pqc_envelope(
36 self, message_signer: MessageSigner, sample_tool_call: dict
37 ) -> None:
38 signed = message_signer.sign_message(sample_tool_call)
39 assert "_pqc" in signed
40 pqc = signed["_pqc"]
41 assert "signer_did" in pqc
42 assert "algorithm" in pqc
43 assert "timestamp" in pqc
44 assert "nonce" in pqc
45 assert "signature" in pqc
46 assert "public_key" in pqc
47
48 def test_verify_valid_signature(
49 self, message_signer: MessageSigner, sample_tool_call: dict
50 ) -> None:
51 signed = message_signer.sign_message(sample_tool_call)
52 result = MessageSigner.verify_message(signed)
53 assert result.valid is True
54 assert result.signer_did == message_signer.identity.did
55
56 def test_verify_tampered_message_fails(
57 self, message_signer: MessageSigner, sample_tool_call: dict
58 ) -> None:
59 signed = message_signer.sign_message(sample_tool_call)
60 # Tamper with the payload
61 signed["id"] = "tampered"
62 result = MessageSigner.verify_message(signed)
63 assert result.valid is False
64
65 def test_verify_wrong_key_fails(
66 self,
67 message_signer: MessageSigner,
68 server_identity: AgentIdentity,
69 sample_tool_call: dict,
70 ) -> None:
71 signed = message_signer.sign_message(sample_tool_call)
72 # Replace the public key with the server's key (mismatch)
73 signed["_pqc"]["public_key"] = server_identity.signing_keypair.public_key.hex()
74 result = MessageSigner.verify_message(signed)
75 assert result.valid is False
76
77 def test_verify_no_pqc_envelope(self, sample_tool_call: dict) -> None:
78 result = MessageSigner.verify_message(sample_tool_call)
79 assert result.valid is False
80 assert result.error == "No _pqc envelope"
81
82
83 class TestStripPQC:
84 def test_strip_pqc_removes_envelope(
85 self, message_signer: MessageSigner, sample_tool_call: dict
86 ) -> None:
87 signed = message_signer.sign_message(sample_tool_call)
88 stripped = MessageSigner.strip_pqc(signed)
89 assert "_pqc" not in stripped
90 # Original fields should still be there
91 assert stripped["jsonrpc"] == "2.0"
92 assert stripped["method"] == "tools/call"
93
94
95 class TestNonce:
96 def test_nonce_uniqueness(
97 self, message_signer: MessageSigner, sample_tool_call: dict
98 ) -> None:
99 """Each signed message gets a unique nonce."""
100 signed1 = message_signer.sign_message(sample_tool_call)
101 signed2 = message_signer.sign_message(sample_tool_call)
102 assert signed1["_pqc"]["nonce"] != signed2["_pqc"]["nonce"]
103