main
py 135 lines 4.7 KB
Raw
1 """Tests for canary token generation and leak detection."""
2
3 from __future__ import annotations
4
5 import logging
6
7 from scripts.analyze_fallback import validate_output_safety
8 from scripts.canary_token import (
9 CANARY_PREFIX,
10 check_output_for_any_canary,
11 check_output_for_leak,
12 generate_canary,
13 inject_canary,
14 )
15
16
17 def test_generate_canary_format() -> None:
18 canary = generate_canary()
19 assert canary.startswith(f"{CANARY_PREFIX}-")
20 # prefix + dash + 16 hex chars
21 hex_part = canary.split("-", 2)[-1]
22 assert len(hex_part) == 16
23 assert all(c in "0123456789abcdef" for c in hex_part)
24
25
26 def test_generate_canary_uniqueness() -> None:
27 canaries = {generate_canary() for _ in range(100)}
28 assert len(canaries) == 100
29
30
31 def test_inject_canary_into_prompt_with_heading() -> None:
32 prompt = "# Weekly Analysis\n\nAnalyze repos...\n"
33 canary = generate_canary()
34 result = inject_canary(prompt, canary)
35 assert canary in result
36 assert "INTERNAL VERIFICATION TOKEN" in result
37 assert result.startswith("# Weekly Analysis\n")
38
39
40 def test_inject_canary_into_prompt_without_heading() -> None:
41 prompt = "Analyze repos...\n"
42 canary = generate_canary()
43 result = inject_canary(prompt, canary)
44 assert canary in result
45
46
47 def test_check_output_no_leak() -> None:
48 canary = generate_canary()
49 output = "## This Week's Trends\n\nRust tools gained momentum..."
50 result = check_output_for_leak(output, canary)
51 assert not result.leaked
52 assert result.match_position is None
53
54
55 def test_check_output_exact_leak() -> None:
56 canary = generate_canary()
57 output = f"## Analysis\n\nThe system uses {canary} internally.\n"
58 result = check_output_for_leak(output, canary)
59 assert result.leaked
60 assert result.match_position is not None
61
62
63 def test_check_output_case_insensitive_leak() -> None:
64 canary = generate_canary()
65 output = f"Found token: {canary.upper()}\n"
66 result = check_output_for_leak(output, canary)
67 assert result.leaked
68
69
70 def test_check_output_partial_leak(caplog) -> None:
71 canary = generate_canary()
72 # Include just prefix + 8 hex chars
73 partial = canary[: len(CANARY_PREFIX) + 1 + 8]
74 output = f"Extracted: {partial}\n"
75 with caplog.at_level(logging.WARNING):
76 result = check_output_for_leak(output, canary)
77 assert result.leaked
78
79
80 def test_check_output_for_any_canary_finds_pattern() -> None:
81 output = f"Leaked: {CANARY_PREFIX}-abcdef1234567890 in text"
82 result = check_output_for_any_canary(output)
83 assert result.leaked
84 assert result.canary == f"{CANARY_PREFIX}-abcdef1234567890"
85
86
87 def test_check_output_for_any_canary_clean() -> None:
88 output = "Normal editorial content about repos and trends."
89 result = check_output_for_any_canary(output)
90 assert not result.leaked
91
92
93 def test_check_empty_output() -> None:
94 canary = generate_canary()
95 result = check_output_for_leak("", canary)
96 assert not result.leaked
97
98
99 # ═══════════════════════════════════════════════════════════════════════════════
100 # validate_output_safety integration tests
101 # ═══════════════════════════════════════════════════════════════════════════════
102
103
104 def test_validate_output_safety_clean() -> None:
105 output = "## This Week's Trends\n\nRust and Go dominate this week.\n"
106 violations = validate_output_safety(output)
107 assert violations == []
108
109
110 def test_validate_output_safety_canary_leak() -> None:
111 canary = generate_canary()
112 output = f"## Analysis\n\nThe internal token is {canary}.\n"
113 violations = validate_output_safety(output, canary)
114 assert len(violations) >= 1
115 assert "Canary token leaked" in violations[0]
116
117
118 def test_validate_output_safety_boundary_marker_leak() -> None:
119 output = "## Analysis\n\n<untrusted-content>some data</untrusted-content>\n"
120 violations = validate_output_safety(output)
121 assert len(violations) >= 1
122 assert "boundary marker" in violations[0].lower() or "boundary marker" in violations[1].lower()
123
124
125 def test_validate_output_safety_unknown_canary_pattern() -> None:
126 output = f"Found: {CANARY_PREFIX}-deadbeef12345678 in output\n"
127 violations = validate_output_safety(output)
128 assert len(violations) >= 1
129 assert "Unknown canary pattern" in violations[0]
130
131
132 def test_validate_output_safety_no_false_positives_on_normal_hex() -> None:
133 output = "Commit hash: abcdef1234567890abcdef\nSHA: deadbeef12345678\n"
134 violations = validate_output_safety(output)
135 assert violations == []