Merge pull request #1 from Ceaseless04/feat/tests-and-detectors

Additional tests and Detectors

Author: Ceaseless04

README.md

@@ -46,6 +46,25 @@ python -m pytest -q
 You should see five tests covering core functionality.
 
 
+## Project Structure
+
+### Extending detection
+
+The :class:`agentshield.secret_scanner.SecretScanner` class ships with a
+set of common regexes (API keys, tokens, AWS formats, JWTs, etc.).  If
+you need to recognise additional secrets, simply:
+
+```python
+from agentshield.secret_scanner import SecretScanner
+import re
+
+scanner = SecretScanner()
+scanner.register_pattern("MY_SECRET", re.compile(r"mysecret=\S+"))
+```
+
+Patterns are applied in the order they are registered, and you can also
+provide a custom list during initialization.
+
 ## Project Structure
 
 ```
@@ -67,4 +86,3 @@ LICENSE
 ## License
 
 This project is open source under the Apache license.
-🛡️ AgentShield is a zero-trust security layer that lets AI agents safely access repositories, tools, and APIs without exposing secrets or proprietary data.

agentshield/secret_scanner.py

@@ -34,13 +34,31 @@ class SecretScanner:
         ("ENV_VAR", re.compile(r"\b[A-Z0-9_]+=[^\n]+")),
         ("TOKEN", re.compile(r"(?i)token\s*=\s*\S+")),
         ("PASSWORD", re.compile(r"(?i)password\s*=\s*[^\s]+")),
+        # common provider keys
+        ("AWS_ACCESS_KEY", re.compile(r"AKIA[0-9A-Z]{16}")),
+        ("AWS_SECRET_KEY", re.compile(r"(?i)aws_secret_access_key\s*=\s*\S+")),
+        ("JWT", re.compile(r"[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+")),
     ]
     # entropy threshold for high-entropy strings
     ENTROPY_THRESHOLD: float = 4.5
 
     def __init__(self, patterns: List[tuple[str, Pattern]] | None = None) -> None:
+        """Initialize scanner with optional custom patterns list.
+
+        ``patterns`` should be a list of ``(name, regex)`` tuples.  If not
+        provided, the built-in defaults are used.  Callers may also call
+        :func:`register_pattern` on the instance to add detectors lazily.
+        """
         self.patterns = patterns or list(self.DEFAULT_PATTERNS)
 
+    def register_pattern(self, name: str, pattern: Pattern) -> None:
+        """Add a new named regex pattern to this scanner instance.
+
+        Useful for plugins or application-specific secrets.  Patterns are
+        evaluated in insertion order during ``scan``.
+        """
+        self.patterns.append((name, pattern))
+
     def scan(self, text: str) -> List[SecretMatch]:
         """Return a list of secrets found in ``text``.
 

tests/test_scanner.py

@@ -1,3 +1,5 @@
+import re
+
 import pytest
 
 from agentshield.secret_scanner import SecretScanner
@@ -15,3 +17,20 @@ def test_entropy_detector():
     scanner = SecretScanner()
     results = scanner.scan(high)
     assert any(r.secret_type == "HIGH_ENTROPY" for r in results)
+
+
+def test_additional_patterns():
+    scanner = SecretScanner()
+    # AWS access key format (16 characters after prefix)
+    assert any(r.secret_type == "AWS_ACCESS_KEY" for r in scanner.scan("AKIA1234567890ABCDEF"))
+    # JWT-like string
+    jwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.e30.TJVA95OrM7E2cBab30RMHrHDcEfhlp"
+    assert any(r.secret_type == "JWT" for r in scanner.scan(jwt))
+
+
+def test_register_pattern():
+    scanner = SecretScanner()
+    scanner.register_pattern("FOO", re.compile(r"foo=\d+"))
+    results = scanner.scan("foo=1234")
+    assert any(r.secret_type == "FOO" for r in results)
+    assert any(r.secret_type == "FOO" for r in results)