Users/mmaitre/include exclude (#63)

Author: mmaitre314

README.md

@@ -40,11 +40,37 @@ picklescan --url https://huggingface.co/sshleifer/tiny-distilbert-base-cased-dis
 
 To scan Numpy's `.npy` files, pip install the `numpy` package first.
 
+## Usage
+
+### Exit codes
+
 The scanner exit status codes are (a-la [ClamAV](https://www.clamav.net/)):
 - `0`: scan did not find malware
 - `1`: scan found malware
 - `2`: scan failed
 
+### Filtering files and directories
+
+When scanning directories, files and subdirectories can be filtered using regular expressions (again modeled after ClamAV). Each option can be specified multiple times:
+
+| Option | Description |
+|---|---|
+| `--exclude=REGEX` | Don't scan files whose path matches the regex |
+| `--include=REGEX` | Only scan files whose path matches the regex |
+| `--exclude-dir=REGEX` | Don't descend into directories whose path matches the regex |
+| `--include-dir=REGEX` | Only descend into directories whose path matches the regex |
+
+Key behaviors:
+- **Excludes always win over includes.** A file or directory matching both an exclude and an include pattern is skipped.
+- **Multiple patterns OR together.** A file is included if it matches *any* `--include` pattern.
+- **No includes = everything eligible.** Include patterns only narrow the scan when specified.
+- **`--exclude-dir` prunes traversal.** The directory and all of its contents are skipped entirely.
+
+```bash
+# Only scan .pkl files, skip the cache/ subdirectory
+picklescan --path models/ --include='\.pkl$' --exclude-dir='cache'
+```
+
 ## Develop
 
 Create and activate the conda environment ([miniconda](https://docs.conda.io/en/latest/miniconda.html) is sufficient):

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = picklescan
-version = 1.0.1
+version = 1.0.2
 author = Matthieu Maitre
 author_email = mmaitre314@users.noreply.github.com
 description = Security scanner detecting Python Pickle files performing suspicious actions

src/picklescan/cli.py

@@ -1,9 +1,10 @@
 import argparse
 import logging
 import os
+import re
 import sys
 
-from .scanner import ScanResult, scan_directory_path
+from .scanner import ScanFilter, ScanResult, scan_directory_path
 from .scanner import scan_file_path
 from .scanner import scan_url
 from .scanner import scan_huggingface_model
@@ -40,6 +41,36 @@ def main():
     )
     parser.add_argument("-g", "--globals", help="list all globals found", action="store_true")
     parser.set_defaults(globals=False)
+    parser.add_argument(
+        "--exclude",
+        action="append",
+        default=[],
+        metavar="REGEX",
+        help="Don't scan file names matching regular expression. Can be used multiple times.",
+    )
+    parser.add_argument(
+        "--include",
+        action="append",
+        default=[],
+        metavar="REGEX",
+        help="Only scan file names matching regular expression. Can be used multiple times.",
+    )
+    parser.add_argument(
+        "--exclude-dir",
+        action="append",
+        default=[],
+        metavar="REGEX",
+        dest="exclude_dir",
+        help="Don't scan directory names matching regular expression. Can be used multiple times.",
+    )
+    parser.add_argument(
+        "--include-dir",
+        action="append",
+        default=[],
+        metavar="REGEX",
+        dest="include_dir",
+        help="Only scan directory names matching regular expression. Can be used multiple times.",
+    )
     parser.add_argument(
         "-l",
         "--log",
@@ -53,13 +84,20 @@ def main():
     if "log_level" in args and args.log_level is not None:
         _log.setLevel(getattr(logging, args.log_level))
 
+    scan_filter = ScanFilter(
+        exclude=[re.compile(p) for p in args.exclude],
+        include=[re.compile(p) for p in args.include],
+        exclude_dir=[re.compile(p) for p in args.exclude_dir],
+        include_dir=[re.compile(p) for p in args.include_dir],
+    )
+
     try:
         if args.path is not None:
             path = os.path.abspath(args.path)
             if not os.path.exists(path):
                 raise FileNotFoundError(f"Path {path} does not exist")
             if os.path.isdir(path):
-                scan_result = scan_directory_path(path)
+                scan_result = scan_directory_path(path, scan_filter=scan_filter)
             else:
                 scan_result = scan_file_path(path)
         elif args.url is not None:

src/picklescan/scanner.py

@@ -1,11 +1,12 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 import http.client
 import io
 import json
 import logging
 import os
 import pickletools
+import re
 from tarfile import TarError
 from tempfile import TemporaryDirectory
 from typing import IO, List, Optional, Set, Tuple
@@ -51,6 +52,25 @@ def merge(self, sr: "ScanResult"):
         self.scan_err = self.scan_err or sr.scan_err
 
 
+@dataclass
+class ScanFilter:
+    """Filtering options for directory scans, modeled after ClamAV clamscan.
+
+    - ``exclude``: regexes matched against the full file path; matching files are skipped.
+    - ``include``: regexes matched against the full file path; when set, only matching files are scanned.
+    - ``exclude_dir``: regexes matched against directory paths; matching directories are not traversed.
+    - ``include_dir``: regexes matched against directory paths; when set, only matching directories are traversed.
+
+    Excludes always take precedence over includes.
+    Multiple patterns of the same kind are combined with logical OR.
+    """
+
+    exclude: List[re.Pattern] = field(default_factory=list)
+    include: List[re.Pattern] = field(default_factory=list)
+    exclude_dir: List[re.Pattern] = field(default_factory=list)
+    include_dir: List[re.Pattern] = field(default_factory=list)
+
+
 class GenOpsError(Exception):
     def __init__(self, msg: str, globals: Optional[Set[Tuple[str, str]]]):
         self.msg = msg
@@ -593,12 +613,31 @@ def scan_huggingface_model(repo_id):
     return scan_result
 
 
-def scan_directory_path(path) -> ScanResult:
+def _matches_any(patterns: List[re.Pattern], text: str) -> bool:
+    """Return True if *text* matches at least one compiled regex in *patterns*."""
+    return any(p.search(text) for p in patterns)
+
+
+def scan_directory_path(path, scan_filter: Optional[ScanFilter] = None) -> ScanResult:
     _log.debug(f"scan_directory_path({path})")
 
     scan_result = ScanResult([])
 
-    for base_path, _, file_names in os.walk(path):
+    for base_path, dir_names, file_names in os.walk(path):
+        # --- directory filtering (prune in-place so os.walk skips them) ---
+        if scan_filter is not None:
+            filtered_dirs = []
+            for d in dir_names:
+                dir_path = os.path.join(base_path, d)
+                if _matches_any(scan_filter.exclude_dir, dir_path):
+                    _log.debug("Excluding directory %s (matched --exclude-dir)", dir_path)
+                    continue
+                if scan_filter.include_dir and not _matches_any(scan_filter.include_dir, dir_path):
+                    _log.debug("Skipping directory %s (no --include-dir match)", dir_path)
+                    continue
+                filtered_dirs.append(d)
+            dir_names[:] = filtered_dirs
+
         for file_name in file_names:
             file_ext = os.path.splitext(file_name)[1]
             if (
@@ -608,6 +647,16 @@ def scan_directory_path(path) -> ScanResult:
             ):
                 continue
             file_path = os.path.join(base_path, file_name)
+
+            # --- file filtering ---
+            if scan_filter is not None:
+                if _matches_any(scan_filter.exclude, file_path):
+                    _log.debug("Excluding file %s (matched --exclude)", file_path)
+                    continue
+                if scan_filter.include and not _matches_any(scan_filter.include, file_path):
+                    _log.debug("Skipping file %s (no --include match)", file_path)
+                    continue
+
             _log.debug("Scanning file %s", file_path)
             with open(file_path, "rb") as file:
                 scan_result.merge(scan_bytes(file, file_path, file_ext))

tests/test_scanner.py

@@ -3,6 +3,7 @@
 import io
 import os
 import pickle
+import re
 import sys
 from typing import Union
 from unittest import TestCase
@@ -14,6 +15,7 @@
 from picklescan.scanner import (
     Global,
     SafetyLevel,
+    ScanFilter,
     ScanResult,
     _http_get,
     _list_globals,
@@ -586,3 +588,95 @@ def test_not_a_pickle_file():
     # File is not a valid pickle, but scanner should not error - just report no threats
     not_a_pickle = ScanResult([], scanned_files=1, issues_count=0, infected_files=0, scan_err=False)
     compare_scan_results(scan_file_path(f"{_root_path}/data/not_a_pickle.bin"), not_a_pickle)
+
+
+# ---------------------------------------------------------------------------
+# Tests for scan_directory_path with ScanFilter (--include/--exclude support)
+# ---------------------------------------------------------------------------
+
+
+def test_scan_directory_exclude_file():
+    """--exclude skips files whose full path matches the regex."""
+    # Exclude all .zip files – only .pkl/.pickle/.pt/.bin/.7z remain
+    sf = ScanFilter(exclude=[re.compile(r"\.zip$")])
+    sr = scan_directory_path(f"{_root_path}/data/", scan_filter=sf)
+    # No .zip file should have been scanned
+    assert sr.scanned_files > 0
+    # The unfiltered scan has 44 scanned files (from test_scan_directory_path);
+    # we just verify that some files were dropped.
+    unfiltered = scan_directory_path(f"{_root_path}/data/")
+    assert sr.scanned_files < unfiltered.scanned_files
+
+
+def test_scan_directory_include_file():
+    """--include restricts scans to files whose path matches the regex."""
+    # Only scan benign .pkl files
+    sf = ScanFilter(include=[re.compile(r"benign0_v3\.pkl$")])
+    sr = scan_directory_path(f"{_root_path}/data/", scan_filter=sf)
+    assert sr.scanned_files == 1
+    assert sr.issues_count == 0
+
+
+def test_scan_directory_exclude_wins_over_include():
+    """Excludes always take precedence over includes (ClamAV semantics)."""
+    sf = ScanFilter(
+        include=[re.compile(r"benign0_v3\.pkl$")],
+        exclude=[re.compile(r"benign")],
+    )
+    sr = scan_directory_path(f"{_root_path}/data/", scan_filter=sf)
+    assert sr.scanned_files == 0
+
+
+def test_scan_directory_exclude_dir():
+    """--exclude-dir prevents traversal into matching directories."""
+    # Scanning the parent tests/ directory but excluding 'data2'
+    sf = ScanFilter(exclude_dir=[re.compile(r"data2")])
+    sr = scan_directory_path(f"{_root_path}/", scan_filter=sf)
+    # Should still find files in data/ but none from data2/
+    assert sr.scanned_files > 0
+    # Compare with an include_dir that only allows data/
+    sf2 = ScanFilter(include_dir=[re.compile(r"/data$")])
+    sr2 = scan_directory_path(f"{_root_path}/", scan_filter=sf2)
+    assert sr2.scanned_files > 0
+    # Both should give the same set of scanned files (only data/)
+    assert sr.scanned_files == sr2.scanned_files
+
+
+def test_scan_directory_include_dir():
+    """--include-dir restricts which directories are traversed."""
+    # Only descend into data2/
+    sf = ScanFilter(include_dir=[re.compile(r"data2")])
+    sr = scan_directory_path(f"{_root_path}/", scan_filter=sf)
+    assert sr.scanned_files > 0
+
+    # Verify data/ files are NOT included by scanning only data/ and comparing
+    sf_data_only = ScanFilter(include_dir=[re.compile(r"/data$")])
+    sr_data = scan_directory_path(f"{_root_path}/", scan_filter=sf_data_only)
+    # data2 results should differ from data-only results
+    assert sr.scanned_files != sr_data.scanned_files
+
+
+def test_scan_directory_multiple_patterns():
+    """Multiple patterns of the same kind are combined with logical OR."""
+    sf = ScanFilter(
+        include=[re.compile(r"benign0_v3\.pkl$"), re.compile(r"benign0_v4\.pkl$")],
+    )
+    sr = scan_directory_path(f"{_root_path}/data/", scan_filter=sf)
+    assert sr.scanned_files == 2
+    assert sr.issues_count == 0
+
+
+def test_scan_directory_no_filter():
+    """Passing no filter (None) gives the same result as default behaviour."""
+    sr_none = scan_directory_path(f"{_root_path}/data/", scan_filter=None)
+    sr_default = scan_directory_path(f"{_root_path}/data/")
+    assert sr_none.scanned_files == sr_default.scanned_files
+    assert sr_none.issues_count == sr_default.issues_count
+
+
+def test_scan_directory_empty_filter():
+    """An empty ScanFilter (no patterns) behaves like no filter at all."""
+    sf = ScanFilter()
+    sr = scan_directory_path(f"{_root_path}/data/", scan_filter=sf)
+    sr_default = scan_directory_path(f"{_root_path}/data/")
+    assert sr.scanned_files == sr_default.scanned_files