veritensor 1.9.1__tar.gz → 1.9.2__tar.gz

{veritensor-1.9.1/src/veritensor.egg-info → veritensor-1.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veritensor
-Version: 1.9.1
+Version: 1.9.2
 Summary: Antivirus for the AI Supply Chain. Scans models, datasets, notebooks, and RAG documents for threats.
 Author-email: Arsenii Brazhnyk <veritensor@gmail.com>
 License: Apache-2.0
@@ -35,6 +35,8 @@ Requires-Dist: gitpython>=3.1.49
 Requires-Dist: pyarrow>=14.0.0
 Requires-Dist: pandas>=2.0.0
 Requires-Dist: openpyxl>=3.1.0
+Requires-Dist: defusedxml>=0.7.1
+Requires-Dist: tomli>=2.0.0; python_version < "3.11"
 Provides-Extra: test
 Requires-Dist: pytest>=7.0.0; extra == "test"
 Requires-Dist: pytest-mock>=3.10.0; extra == "test"

{veritensor-1.9.1 → veritensor-1.9.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "veritensor"
-version = "1.9.1"
+version = "1.9.2"
 description = "Antivirus for the AI Supply Chain. Scans models, datasets, notebooks, and RAG documents for threats."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -36,11 +36,12 @@ dependencies = [
     "pillow>=12.2.0",
     "lxml>=6.1.0",
     "urllib3 >= 2.7.0",
-    # Fickling removed from here to avoid installing PyTorch (800MB+) for regular users
     "gitpython>=3.1.49",
     "pyarrow>=14.0.0",
     "pandas>=2.0.0",
     "openpyxl>=3.1.0",
+    "defusedxml>=0.7.1",
+    "tomli>=2.0.0; python_version < '3.11'",
 ]
 
 

veritensor-1.9.2/src/veritensor/__init__.py

@@ -0,0 +1,2 @@
+__version__ = "1.9.2"
+

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/cli/main.py

@@ -11,6 +11,7 @@ import os
 import datetime
 import time
 import requests
+import copy
 import fnmatch
 import concurrent.futures
 import multiprocessing
@@ -204,7 +205,8 @@ def is_suppressed(file_path: str, threat_msg: str, suppressions: List[dict]) ->
         if not supp_path or not supp_type:
             continue
         file_name = Path(file_path).name
-        if file_name == supp_path and supp_type in threat_msg:
+        supp_name = Path(supp_path).name
+        if file_name == supp_name and supp_type in threat_msg:
             return True
     return False
 
@@ -239,6 +241,16 @@ def scan_worker(args: Tuple[str, VeritensorConfig, Optional[str], bool, bool, bo
     scan_res = ScanResult(file_path=file_path_str)
     scan_res.repo_id = repo 
 
+    # Check MCP config files early, before extension-based routing swallows .json files
+    if file_path and ext == ".json" and not is_s3:
+        try:
+            if is_mcp_config_file(file_path):
+                perm_result = audit_mcp_config(file_path)
+                for t in perm_result.to_threat_strings():
+                    scan_res.add_threat(t)
+        except Exception as e:
+            logger.debug(f"MCP Audit failed for {file_name}: {e}")
+
     # --- HYBRID ROUTING (Sending heavy files to the server) ---
     # If the user has an Enterprise server connected
     already_remote_scanned = False
@@ -360,12 +372,7 @@ def scan_worker(args: Tuple[str, VeritensorConfig, Optional[str], bool, bool, bo
                 if mcp_result.mcp_tools_found:
                     for t in mcp_result.to_threat_strings():
                         scan_res.add_threat(t)
-            elif ext == ".json" and file_path:
-                if is_mcp_config_file(file_path):
-                    from veritensor.engines.static.mcp_permission_auditor import audit_mcp_config
-                    perm_result = audit_mcp_config(file_path)
-                    for t in perm_result.to_threat_strings():
-                        scan_res.add_threat(t)
+          
         else:
             # If the format is unknown to any engine at all.
             # We add INFO, but DO NOT call add_threat(),
@@ -379,11 +386,6 @@ def scan_worker(args: Tuple[str, VeritensorConfig, Optional[str], bool, bool, bo
     except Exception as e:
         scan_res.add_threat(f"CRITICAL: Engine Error: {str(e)}")
 
-    # --- MERGE REMOTE THREATS ---
-    for t in remote_threats:
-        if t not in scan_res.threats:
-            scan_res.add_threat(t)
-
     # --- C. License Check ---
     if not is_s3 and file_path:
         reader = get_reader_for_file(file_path)
@@ -414,6 +416,11 @@ def scan_worker(args: Tuple[str, VeritensorConfig, Optional[str], bool, bool, bo
 
     return scan_res
 
+def _worker_initializer():
+    """Pre-load spaCy model once per worker process to avoid RAM spikes during scans."""
+    from veritensor.engines.content.pii import PIIScanner
+    PIIScanner.get_engine()
+    
 # --- SHARED SCAN LOGIC ---
 def _run_scan_process(
     paths: List[str], repo: Optional[str], jobs: Optional[int], 
@@ -527,7 +534,10 @@ def _run_scan_process(
         ) as progress:
             
             main_task = progress.add_task("Scanning...", total=len(tasks))
-            executor = concurrent.futures.ProcessPoolExecutor(max_workers=jobs)
+            executor = concurrent.futures.ProcessPoolExecutor(
+                max_workers=jobs,
+                initializer=_worker_initializer
+            )
             
             future_to_file = {
                 executor.submit(scan_worker, task_args): task_args[0] 
@@ -583,14 +593,14 @@ def scan(
     watch: bool = typer.Option(False, "--watch", "-w", help="Watch files for changes and scan automatically"),
 
 ):
-    """
-    Scans models, data, and code for security threats.
-    """
-    config = ConfigLoader.load()
+    """Scans models, data, and code for security threats."""
+    # Work on a shallow copy to prevent mutating the singleton between runs
+    base_config = ConfigLoader.load()
+    config = copy.copy(base_config)
+    
     if verbose: logger.setLevel(logging.DEBUG)
     is_machine_output = (json_output or sarif_output or sbom_output) and not output_file
     
-    # Passing the flags from the console to the global config object
     config.report_url = report_to or config.report_url
     config.api_key = api_key or config.api_key
     
@@ -637,7 +647,6 @@ def scan(
         console.print(Panel.fit(f"🛡️  [bold cyan]Veritensor Security Scanner[/bold cyan] v{__version__}", border_style="cyan"))
 
     if watch:
-        import time
         console.print("[bold cyan]👀 Veritensor Watcher started. Monitoring files for changes... (Press Ctrl+C to stop)[/bold cyan]")
         
         # Collecting the initial file modification dates
@@ -921,7 +930,10 @@ def _print_table(results: List[ScanResult]):
         else:
             unique_threats = list(dict.fromkeys(res.threats))
             display_threats = "\n".join(unique_threats)
-        table.add_row(res.file_path.split("/")[-1], f"[{status_style}]{res.status}[/{status_style}]", display_threats)
+            
+        # Use Path.name for cross-platform compatibility
+        file_name = Path(res.file_path).name if res.file_path else "unknown"
+        table.add_row(file_name, f"[{status_style}]{res.status}[/{status_style}]", display_threats)
     console.print(table)
 
 def _perform_signing(image: str, status: str, config, timestamp: str, results: List[ScanResult]):

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/core/cache.py

@@ -20,8 +20,11 @@ class HashCache:
         if not self.conn:
             return
         try:
-            cutoff_mtime = time.time() - (days * 86400)
-            self.cursor.execute("DELETE FROM file_cache WHERE mtime < ?", (cutoff_mtime,))
+            cutoff = time.time() - (days * 86400)
+            self.cursor.execute(
+                "DELETE FROM file_cache WHERE COALESCE(last_scanned, mtime) < ?", 
+                (cutoff,)
+            )
             self.conn.commit()
         except Exception as e:
             logger.debug(f"Cache cleanup failed: {e}")
@@ -35,11 +38,9 @@ class HashCache:
             # check_same_thread=False allows using connection across threads.
             # In our multiprocessing architecture, only the MAIN process interacts with this class.
             self.conn = sqlite3.connect(str(CACHE_FILE), check_same_thread=False)
-            
             # Enable WAL mode for better performance and concurrency safety
             self.conn.execute("PRAGMA journal_mode=WAL;")
             self.conn.execute("PRAGMA synchronous=NORMAL;")
-            
             self.cursor = self.conn.cursor()
             
             # Create table
@@ -48,14 +49,24 @@ class HashCache:
                     path TEXT PRIMARY KEY,
                     hash TEXT,
                     size INTEGER,
-                    mtime REAL
+                    mtime REAL,
+                    last_scanned REAL
                 )
             """)
+            
+            # Handle migration for existing DBs
+            try:
+                self.cursor.execute("ALTER TABLE file_cache ADD COLUMN last_scanned REAL")
+                self.conn.commit()
+            except Exception:
+                pass  # Column already exists
+                
             self.conn.commit()
         except Exception as e:
             logger.warning(f"Failed to initialize cache DB: {e}")
             self.conn = None
 
+
     def get(self, file_path: Path) -> Optional[str]:
         """Returns the hash if the file has not been changed."""
         if not self.conn: return None
@@ -71,16 +82,21 @@ class HashCache:
             row = self.cursor.fetchone()
             
             if row:
-                cached_hash, cached_size, cached_mtime = row
-                # Compare size and mtime (float comparison usually safe for exact system mtime)
+                cached_hash, cached_size, cached_mtime = row[:3]
                 if cached_size == stats.st_size and cached_mtime == stats.st_mtime:
+                    # Update last_scanned on cache hit
+                    self.cursor.execute(
+                        "UPDATE file_cache SET last_scanned = ? WHERE path = ?",
+                        (time.time(), abs_path)
+                    )
+                    self.conn.commit()
                     return cached_hash
             
             return None
         except Exception:
-            # If file not found or permission error during stat
             return None
 
+
     def set(self, file_path: Path, file_hash: str):
         """Saves or updates the hash in the cache."""
         if not self.conn: return
@@ -90,9 +106,9 @@ class HashCache:
             stats = file_path.stat()
             
             self.cursor.execute("""
-                INSERT OR REPLACE INTO file_cache (path, hash, size, mtime)
-                VALUES (?, ?, ?, ?)
-            """, (abs_path, file_hash, stats.st_size, stats.st_mtime))
+                INSERT OR REPLACE INTO file_cache (path, hash, size, mtime, last_scanned)
+                VALUES (?, ?, ?, ?, ?)
+            """, (abs_path, file_hash, stats.st_size, stats.st_mtime, time.time()))
             
             self.conn.commit()
         except Exception as e:

veritensor-1.9.2/src/veritensor/core/networking.py

@@ -0,0 +1,67 @@
+# Copyright 2026 Veritensor Security Apache 2.0
+import socket
+import ipaddress
+import logging
+import urllib3.util.connection as urllib3_cn
+import requests
+
+logger = logging.getLogger(__name__)
+
+# Save the original connection function
+_orig_create_connection = urllib3_cn.create_connection
+
+def validate_ip_ssrf(ip_addr: str) -> bool:
+    """Checks if an IP address is public and safe to connect to."""
+    try:
+        ip_obj = ipaddress.ip_address(ip_addr)
+        if (ip_obj.is_private or ip_obj.is_loopback or 
+            ip_obj.is_link_local or ip_obj.is_multicast or 
+            ip_obj.is_reserved or ip_obj.is_unspecified): 
+            return False
+        return True
+    except ValueError:
+        return False
+
+def safe_create_connection(address, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None, socket_options=None):
+    """
+    Patched connection creator. Resolves DNS, validates against SSRF, 
+    and forces the socket to connect to the safe IP, preventing TOCTOU DNS Rebinding.
+    TLS/SNI remains intact because urllib3 wraps the socket with the original hostname later.
+    """
+    host, port = address
+    
+    try:
+        ip_list = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM)
+    except socket.gaierror:
+        raise ValueError(f"Could not resolve hostname: {host}")
+
+    safe_ip = None
+    for item in ip_list:
+        ip_addr = item[4][0]
+        if validate_ip_ssrf(ip_addr):
+            safe_ip = ip_addr
+            break
+
+    if not safe_ip:
+        raise ValueError(f"SSRF Protection: No safe public IP found for {host}. Access forbidden.")
+        
+    # Connect to the resolved and validated IP
+    return _orig_create_connection((safe_ip, port), timeout, source_address, socket_options)
+
+# Apply the patch globally for the CLI runtime
+urllib3_cn.create_connection = safe_create_connection
+
+def get_safe_session(url: str = None) -> requests.Session:
+    """
+    Returns an SSRF-protected session. 
+    The global urllib3 patch ensures all requests made by this session are safe.
+    """
+    if url:
+        # Trigger a dry-run resolution check immediately if a URL is provided
+        from urllib.parse import urlparse
+        parsed = urlparse(url)
+        hostname = parsed.hostname
+        if hostname:
+            safe_create_connection((hostname, parsed.port or 80))
+            
+    return requests.Session()

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/core/safe_zip.py

@@ -1,4 +1,3 @@
-
 import zipfile
 import logging
 
@@ -32,9 +31,9 @@ class SafeZipReader:
             
             total_size += info.file_size
         
-        # Checking the total size
-        if total_size > SafeZipReader.MAX_UNZIPPED_SIZE:
-            raise ZipBombError(f"Zip Bomb detected! Total unzipped size {total_size} bytes exceeds limit.")
+            # Checking the total size
+            if total_size > SafeZipReader.MAX_UNZIPPED_SIZE:
+                raise ZipBombError(f"Zip Bomb detected! Total unzipped size {total_size} bytes exceeds limit.")
 
     def read(self, zfile: zipfile.ZipFile, name: str) -> bytes:
         info = zfile.getinfo(name)

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/core/streaming.py

@@ -29,7 +29,7 @@ class RemoteStream(io.IOBase):
     def __init__(self, url: str, session: Optional[requests.Session] = None):
         self._validate_url(url)
         self.url = url
-        self.session = session or requests.Session()
+        self._session = session or get_safe_session(self.url)
         self.pos = 0
         self.size = self._fetch_size()
         self._closed = False
@@ -60,10 +60,9 @@ class RemoteStream(io.IOBase):
     def _fetch_size(self) -> int:
         try:
             headers = {"Range": "bytes=0-0"}
-            with get_safe_session(self.url):
-                resp = self.session.get(self.url, headers=headers, stream=True, timeout=10)
-                resp.raise_for_status()
-            
+            resp = self._session.get(self.url, headers=headers, stream=True, timeout=10)
+            resp.raise_for_status()
+                
             content_range = resp.headers.get("Content-Range")
             if content_range and "/" in content_range:
                 return int(content_range.split("/")[-1])
@@ -76,6 +75,7 @@ class RemoteStream(io.IOBase):
             logger.error(f"Failed to fetch size for {self.url}: {e}")
             raise
 
+
     def read(self, size: int = -1) -> bytes:
         if self._closed: raise ValueError("I/O operation on closed file.")
         if self.pos >= self.size: return b""
@@ -89,9 +89,9 @@ class RemoteStream(io.IOBase):
 
         headers = {"Range": f"bytes={self.pos}-{end}"}
         try:
-            with get_safe_session(self.url):
-                resp = self.session.get(self.url, headers=headers, timeout=30)
-                resp.raise_for_status()
+            #  Use safe session
+            resp = self._session.get(self.url, headers=headers, timeout=30)
+            resp.raise_for_status()
             data = resp.content
             self.pos += len(data)
             return data
@@ -111,7 +111,12 @@ class RemoteStream(io.IOBase):
     def tell(self) -> int: return self.pos
     def seekable(self) -> bool: return True
     def readable(self) -> bool: return True
-    def close(self): self._closed = True
+    def close(self):
+        # Properly close the session and call super().close()
+        if not self._closed:
+            self._closed = True
+            self._session.close()
+            super().close()
     def __enter__(self): return self
     def __exit__(self, exc_type, exc, tb): self.close()
 
@@ -189,7 +194,11 @@ class S3Stream(io.IOBase):
     def tell(self) -> int: return self.pos
     def seekable(self) -> bool: return True
     def readable(self) -> bool: return True
-    def close(self): self._closed = True
+    def close(self):
+        # Call super().close() to properly update io.IOBase state
+        if not self._closed:
+            self._closed = True
+            super().close()
     def __enter__(self): return self
     def __exit__(self, exc_type, exc, tb): self.close()
 

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/core/types.py

@@ -1,12 +1,5 @@
 from dataclasses import dataclass, field
 from typing import List, Optional
-from enum import Enum
-
-class Severity(str, Enum):
-    CRITICAL = "CRITICAL"
-    HIGH = "HIGH"
-    MEDIUM = "MEDIUM"
-    LOW = "LOW"
 
 @dataclass
 class ScanResult:

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/engines/content/pii.py

@@ -19,7 +19,7 @@ class PIIScanner:
         
         if PRESIDIO_AVAILABLE is None:
             try:
-                # CRITICAL FIX: Lazy import inside the method to prevent DLL/OSError crashes on startup
+                # Lazy import inside the method to prevent DLL/OSError crashes on startup
                 from presidio_analyzer import AnalyzerEngine
                 from presidio_analyzer.nlp_engine import NlpEngineProvider
                 import spacy

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/engines/data/dataset_engine.py

@@ -1,6 +1,6 @@
 # Copyright 2026 Veritensor Security Apache 2.0
 # Dataset Scanner (Parquet, CSV, JSONL) for Data Poisoning, Malicious URLs, and Toxic Columns
-
+import threading
 import logging
 import json
 from pathlib import Path
@@ -31,23 +31,28 @@ FALLBACK_SUSPICIOUS = [
     "regex:(?i)secret"
 ]
 
-# --- CRITICAL FIX: Robust Singleton for GLiNER ---
+# Thread-safe lazy loading for GLiNER
 _GLINER_UNAVAILABLE = object()
 _gliner_model = None
+_gliner_lock = threading.Lock()
 
 def _get_gliner_model():
-    """Lazy loads GLiNER to avoid 30-90s delay on every function call."""
     global _gliner_model
-    if _gliner_model is None:
-        try:
-            from gliner import GLiNER
-            logger.info("Loading GLiNER model for column analysis...")
-            _gliner_model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
-        except ImportError:
-            _gliner_model = _GLINER_UNAVAILABLE
-        except Exception as e:
-            logger.warning(f"GLiNER failed to load: {e}. Falling back to regex.")
-            _gliner_model = _GLINER_UNAVAILABLE
+    if _gliner_model is not None:
+        return _gliner_model
+        
+    with _gliner_lock:
+        if _gliner_model is None:  # Double-checked locking
+            try:
+                from gliner import GLiNER
+                logger.info("Loading GLiNER model for column analysis...")
+                _gliner_model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
+            except ImportError:
+                _gliner_model = _GLINER_UNAVAILABLE
+            except Exception as e:
+                logger.warning(f"GLiNER failed to load: {e}. Falling back to regex.")
+                _gliner_model = _GLINER_UNAVAILABLE
+                
     return _gliner_model
 
 def _check_toxic_columns(columns: List[str]) -> List[str]:
@@ -60,7 +65,8 @@ def _check_toxic_columns(columns: List[str]) -> List[str]:
         
     toxic_keywords = [
         "social_credit", "political_affiliation", "sexual_orientation", 
-        "race", "religion", "trade_union", "biometric_categorization", "criminal_prediction"
+        "race", "religion", "trade_union", "biometric_categorization", "criminal_prediction",
+        "emotion_state", "sentiment_analysis_workplace", "facial_features", "cctv_face_id", "subliminal_metrics"
     ]
     
     model = _get_gliner_model()
@@ -95,15 +101,19 @@ def scan_dataset(file_path: Path, full_scan: bool = False, bias_profile: Optiona
 
     row_limit = None if full_scan else MAX_ROWS_DEFAULT
     
-    # Initialize Bias Aggregator if profile is provided
-    aggregator = BiasAggregator(bias_profile) if bias_profile else None
+    aggregator = None
+    if bias_profile:
+        try:
+            aggregator = BiasAggregator(bias_profile)
+        except ValueError as e:
+            logger.warning(f"Invalid bias profile for {file_path.name}: {e}")
     
     try:
         text_stream = None
         
         if ext == ".parquet":
             if not PYARROW_AVAILABLE:
-                return ["WARNING: pyarrow not installed. Run 'pip install veritensor[data]'"], None
+                return ["WARNING: pyarrow and pandas are required for local parquet scanning. Run 'pip install pyarrow pandas' or use the Enterprise Server."], None
             
             try:
                 parquet_file = pq.ParquetFile(file_path)
@@ -268,9 +278,16 @@ def _extract_strings_from_json(data: Any) -> Generator[str, None, None]:
 class BiasAggregator:
     """Aggregates counts for fairness metrics using O(1) memory. Safe for multiprocessing (no lambdas)."""
     def __init__(self, profile: dict):
-        self.protected_attrs = profile.get("protected_attributes", [])
+        # Validate required fields to prevent AttributeError during processing
         self.target_var = profile.get("target_variable")
-        self.favorable_label = str(profile.get("favorable_label"))
+        if not self.target_var:
+            raise ValueError("BiasAggregator requires 'target_variable' in bias profile")
+            
+        self.protected_attrs = profile.get("protected_attributes", [])
+        if not self.protected_attrs:
+            raise ValueError("BiasAggregator requires 'protected_attributes' in bias profile")
+            
+        self.favorable_label = str(profile.get("favorable_label", "1"))
         
         self.group_target_counts = {}
         self.proxy_counts = {}

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/engines/data/excel_engine.py

@@ -24,7 +24,7 @@ FORMULA_PREFIXES = ('=', '+', '-', '@')
 def scan_excel(file_path: Path) -> List[str]:
     threats = []
     if not EXCEL_AVAILABLE:
-        return ["WARNING: openpyxl not installed. Run 'pip install veritensor[data]'"]
+        return ["WARNING: openpyxl is required for local Excel scanning. Run 'pip install openpyxl' or use the Enterprise Server."]
 
     try:
         # read_only=True is faster and uses less memory

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/engines/static/dependency_engine.py

@@ -20,7 +20,11 @@ try:
     # Python 3.11+ native TOML support
     import tomllib
 except ImportError:
-    tomllib = None
+    try:
+        # Backport for Python 3.10
+        import tomli as tomllib
+    except ImportError:
+        tomllib = None
 
 # --- Known Malware DB (MVP) ---
 KNOWN_MALICIOUS = {
@@ -214,9 +218,15 @@ def _parse_requirements(path: Path) -> Dict[str, Optional[str]]:
             if not line or line.startswith(("#", "-")): 
                 continue
             
+            # Strip inline comments before any further parsing
+            comment_idx = line.find(' #')
+            if comment_idx != -1:
+                line = line[:comment_idx].strip()
+            if not line:
+                continue
+            
             # Processing direct links (git+, http) with #egg=
             if "egg=" in line:
-                # Пример: git+https://github.com/user/repo.git#egg=evil_pkg
                 name_part = line.split("egg=")[-1].split("&")[0].strip()
                 deps[name_part] = None
                 continue
@@ -236,6 +246,7 @@ def _parse_requirements(path: Path) -> Dict[str, Optional[str]]:
     return deps
 
 
+
 def _parse_pyproject(path: Path) -> Dict[str, Optional[str]]:
     deps = {}
     try:

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/engines/static/mcp_scanner.py

@@ -131,7 +131,6 @@ _SEVERITY_MAP = {
 # Human-in-the-loop parameter names that indicate the tool has a confirmation gate
 _HITL_PARAM_NAMES = frozenset({"confirm", "approved", "dry_run", "dryrun", "force"})
 
-_STDLIB_ENTRYPOINTS = frozenset({"main", "__main__", "cli", "run", "start", "app"})
 # ---------------------------------------------------------------------------
 # Result types
 # ---------------------------------------------------------------------------
@@ -216,13 +215,22 @@ class _MCPToolVisitor(ast.NodeVisitor):
 
     
     def _is_mcp_tool(self, node: ast.FunctionDef) -> bool:
-        if node.name.lower() in _STDLIB_ENTRYPOINTS:
-            return False
+         # Check decorator FIRST. If @tool is present, scan regardless of name.
+        has_mcp_decorator = False
         for dec in node.decorator_list:
             name = _decorator_base_name(dec)
             if name and name.lower() in _MCP_TOOL_DECORATORS:
-                return True
-        return False
+                has_mcp_decorator = True
+                break
+        
+        if not has_mcp_decorator:
+            return False
+        
+        # Only skip dunder names (e.g., __init__) — they are never MCP tools
+        if node.name.startswith("__") and node.name.endswith("__"):
+            return False
+        
+        return True
 
     def _has_hitl_param(self, node: ast.FunctionDef) -> bool:
         params = {arg.arg.lower() for arg in node.args.args}

{veritensor-1.9.1 → veritensor-1.9.2}/src/veritensor/engines/static/notebook_engine.py

@@ -10,6 +10,8 @@ import re
 from veritensor.core.entropy import is_high_entropy
 from veritensor.engines.static.rules import get_severity, SignatureLoader, is_match
 from veritensor.engines.content.pii import PIIScanner  
+import tokenize
+import io as _io
 
 logger = logging.getLogger(__name__)
 
@@ -140,15 +142,21 @@ def _clean_magics(source: str) -> str:
     return "\n".join(lines)
 
 def _check_nesting_level(code: str, max_depth: int = 100) -> bool:
-    # Protection against AST Bomb (Stack Overflow)
+    """Check bracket depth, ignoring content inside string literals."""
     depth = 0
-    for char in code:
-        if char in "([{":
-            depth += 1
-            if depth > max_depth:
-                return False
-        elif char in ")]}":
-            depth -= 1
+    try:
+        tokens = tokenize.generate_tokens(_io.StringIO(code).readline)
+        for tok_type, tok_string, _, _, _ in tokens:
+            if tok_type == tokenize.OP:
+                if tok_string in ('(', '[', '{'):
+                    depth += 1
+                    if depth > max_depth:
+                        return False
+                elif tok_string in (')', ']', '}'):
+                    depth = max(0, depth - 1)
+    except tokenize.TokenError:
+        # Incomplete code (common in notebooks) — allow it to proceed to AST
+        pass
     return True