Files
realmarauderandClaude Opus 4.6 2abace936b Initial release: ComfyUI-TensorTrap security scanner
Three ComfyUI nodes for security scanning:

- Scan Model: Scans model files before loading, blocks execution
  if threats found (pickle exploits, polyglot attacks, CVE bypasses)
- Audit Installed Nodes: Static analysis of all custom node packages
  for eval/exec, subprocess, network access, obfuscation patterns
- Analyze Workflow: Graph analysis of node connections for dangerous
  data flows, known CVE node types, URL injection, and suspicious inputs

Includes:
- Node auditor with regex + AST-based detection
- Workflow graph analyzer tracing data flows through node connections
- Known dangerous node database (8+ CVEs cataloged)
- Detection of code injection, command execution, network access,
  file access, dynamic imports, obfuscation, and supply chain risks

Built on TensorTrap security research documenting 8 CVEs, 4 major
real-world attacks, and 700+ confirmed server compromises in the
ComfyUI ecosystem.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-28 00:19:21 -05:00

350 lines
12 KiB
Python

"""Static analysis scanner for ComfyUI custom node source code.
Scans Python files in custom_nodes/ for dangerous patterns:
- eval(), exec(), compile() — code injection
- subprocess, os.system, os.popen — command execution
- importlib — dynamic module loading
- requests, urllib — network access (data exfiltration, SSRF)
- open() with user-controlled paths — arbitrary file access
- pickle.loads — deserialization attacks
- base64/hex decoding — obfuscation
"""
import ast
import re
from dataclasses import dataclass, field
from pathlib import Path
@dataclass
class NodeFinding:
"""A security finding in a custom node's source code."""
severity: str # CRITICAL, HIGH, MEDIUM, LOW, INFO
category: str # code_injection, command_exec, network, file_access, obfuscation, supply_chain
message: str
file_path: str
line_number: int
code_snippet: str = ""
node_package: str = ""
@dataclass
class NodeAuditResult:
"""Audit results for a single custom node package."""
package_name: str
package_path: str
files_scanned: int = 0
findings: list[NodeFinding] = field(default_factory=list)
@property
def is_safe(self) -> bool:
return not any(f.severity in ("CRITICAL", "HIGH") for f in self.findings)
@property
def max_severity(self) -> str:
severity_order = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "INFO": 4}
if not self.findings:
return "SAFE"
return min(self.findings, key=lambda f: severity_order.get(f.severity, 5)).severity
def to_dict(self) -> dict:
return {
"package_name": self.package_name,
"package_path": self.package_path,
"files_scanned": self.files_scanned,
"is_safe": self.is_safe,
"max_severity": self.max_severity,
"findings": [
{
"severity": f.severity,
"category": f.category,
"message": f.message,
"file_path": f.file_path,
"line_number": f.line_number,
"code_snippet": f.code_snippet,
}
for f in self.findings
],
}
# --- Dangerous patterns ---
# Direct code execution
CODE_INJECTION_PATTERNS = [
(r"\beval\s*\(", "CRITICAL", "eval() call — arbitrary code execution"),
(r"\bexec\s*\(", "CRITICAL", "exec() call — arbitrary code execution"),
(r"\bcompile\s*\(", "HIGH", "compile() call — may bypass eval/exec detection"),
]
# Command execution
COMMAND_EXEC_PATTERNS = [
(r"\bos\.system\s*\(", "CRITICAL", "os.system() — shell command execution"),
(r"\bos\.popen\s*\(", "CRITICAL", "os.popen() — shell command execution"),
(r"\bsubprocess\.\w+\s*\(", "HIGH", "subprocess call — command execution"),
(r"\bPopen\s*\(", "HIGH", "Popen() — process execution"),
]
# Network access
NETWORK_PATTERNS = [
(r"\brequests\.(get|post|put|delete|head|patch)\s*\(", "MEDIUM", "HTTP request — potential data exfiltration or SSRF"),
(r"\burllib\.request\.\w+\s*\(", "MEDIUM", "urllib request — potential data exfiltration or SSRF"),
(r"\bhttpx\.\w+\s*\(", "MEDIUM", "httpx request — potential data exfiltration"),
(r"\bsocket\.\w+\s*\(", "HIGH", "Raw socket access — potential reverse shell or data exfiltration"),
(r"\bwebsocket\b", "MEDIUM", "WebSocket usage — potential C2 communication"),
]
# File access
FILE_ACCESS_PATTERNS = [
(r"\bpickle\.loads?\s*\(", "HIGH", "pickle deserialization — arbitrary code execution risk"),
(r"\bshelve\.open\s*\(", "HIGH", "shelve.open — uses pickle internally"),
(r"\bmarshal\.loads?\s*\(", "HIGH", "marshal deserialization — code execution risk"),
]
# Dynamic imports
DYNAMIC_IMPORT_PATTERNS = [
(r"\bimportlib\.import_module\s*\(", "HIGH", "Dynamic module import — can load arbitrary code"),
(r"\b__import__\s*\(", "HIGH", "Dynamic import — can load arbitrary modules"),
]
# Obfuscation
OBFUSCATION_PATTERNS = [
(r"\bbase64\.b64decode\s*\(", "MEDIUM", "Base64 decoding — may hide malicious payload"),
(r"\bcodecs\.decode\s*\(.*rot_13", "MEDIUM", "ROT13 decoding — obfuscation technique"),
(r"\bbytes\.fromhex\s*\(", "MEDIUM", "Hex decoding — may hide malicious payload"),
(r"\bzlib\.decompress\s*\(", "LOW", "zlib decompression — may hide obfuscated code"),
]
# Supply chain
SUPPLY_CHAIN_PATTERNS = [
(r"\bpip\s+install\b", "HIGH", "Runtime pip install — supply chain risk"),
(r"\bsubprocess.*pip\b", "HIGH", "subprocess pip call — supply chain risk"),
(r"requirements\.txt", "INFO", "requirements.txt reference — check dependencies"),
]
ALL_PATTERN_GROUPS = [
("code_injection", CODE_INJECTION_PATTERNS),
("command_exec", COMMAND_EXEC_PATTERNS),
("network", NETWORK_PATTERNS),
("file_access", FILE_ACCESS_PATTERNS),
("dynamic_import", DYNAMIC_IMPORT_PATTERNS),
("obfuscation", OBFUSCATION_PATTERNS),
("supply_chain", SUPPLY_CHAIN_PATTERNS),
]
# Known safe patterns that trigger false positives
FALSE_POSITIVE_CONTEXTS = [
# Comments
r"^\s*#",
# String definitions / docstrings
r'^\s*["\']',
r'^\s*"""',
r"^\s*'''",
]
def scan_node_package(package_path: Path) -> NodeAuditResult:
"""Scan a single custom node package for dangerous patterns.
Args:
package_path: Path to the custom node package directory
Returns:
NodeAuditResult with findings
"""
package_name = package_path.name
result = NodeAuditResult(
package_name=package_name,
package_path=str(package_path),
)
if not package_path.is_dir():
return result
# Scan all Python files
for py_file in package_path.rglob("*.py"):
result.files_scanned += 1
findings = _scan_python_file(py_file, package_name)
result.findings.extend(findings)
# Check for suspicious non-Python files
_check_suspicious_files(package_path, package_name, result)
return result
def scan_all_nodes(custom_nodes_dir: Path) -> list[NodeAuditResult]:
"""Scan all installed custom node packages.
Args:
custom_nodes_dir: Path to ComfyUI's custom_nodes/ directory
Returns:
List of NodeAuditResult for each package
"""
results = []
if not custom_nodes_dir.is_dir():
return results
for package_dir in sorted(custom_nodes_dir.iterdir()):
if package_dir.is_dir() and not package_dir.name.startswith("."):
result = scan_node_package(package_dir)
results.append(result)
return results
def _scan_python_file(filepath: Path, package_name: str) -> list[NodeFinding]:
"""Scan a single Python file for dangerous patterns."""
findings = []
try:
content = filepath.read_text(encoding="utf-8", errors="ignore")
except OSError:
return findings
lines = content.splitlines()
for line_num, line in enumerate(lines, 1):
# Skip obvious false positives
stripped = line.strip()
if any(re.match(fp, stripped) for fp in FALSE_POSITIVE_CONTEXTS):
continue
for category, patterns in ALL_PATTERN_GROUPS:
for pattern, severity, message in patterns:
if re.search(pattern, line):
# Get a clean code snippet
snippet = stripped[:120]
findings.append(
NodeFinding(
severity=severity,
category=category,
message=message,
file_path=str(filepath),
line_number=line_num,
code_snippet=snippet,
node_package=package_name,
)
)
# AST-based analysis for more accurate detection
findings.extend(_ast_analysis(content, filepath, package_name))
return findings
def _ast_analysis(content: str, filepath: Path, package_name: str) -> list[NodeFinding]:
"""Use AST parsing for more accurate detection of dangerous patterns."""
findings = []
try:
tree = ast.parse(content, filename=str(filepath))
except SyntaxError:
return findings
for node in ast.walk(tree):
# Detect eval/exec as function calls (more accurate than regex)
if isinstance(node, ast.Call):
func_name = _get_call_name(node)
# Check for dangerous builtins used as function calls
if func_name == "eval":
# Check if the argument comes from user input (node inputs)
findings.append(
NodeFinding(
severity="CRITICAL",
category="code_injection",
message="eval() call detected via AST analysis",
file_path=str(filepath),
line_number=node.lineno,
code_snippet=f"eval() at line {node.lineno}",
node_package=package_name,
)
)
elif func_name == "exec":
findings.append(
NodeFinding(
severity="CRITICAL",
category="code_injection",
message="exec() call detected via AST analysis",
file_path=str(filepath),
line_number=node.lineno,
code_snippet=f"exec() at line {node.lineno}",
node_package=package_name,
)
)
# Detect hidden inputs that access PROMPT (workflow inspection)
if isinstance(node, ast.Dict):
for key in node.keys:
if isinstance(key, ast.Constant) and key.value == "hidden":
findings.append(
NodeFinding(
severity="INFO",
category="workflow_access",
message="Node declares hidden inputs (can access full workflow context)",
file_path=str(filepath),
line_number=node.lineno,
code_snippet=f"hidden input declaration at line {node.lineno}",
node_package=package_name,
)
)
return findings
def _get_call_name(node: ast.Call) -> str:
"""Extract the function name from a Call node."""
if isinstance(node.func, ast.Name):
return node.func.id
if isinstance(node.func, ast.Attribute):
return node.func.attr
return ""
def _check_suspicious_files(package_path: Path, package_name: str, result: NodeAuditResult):
"""Check for suspicious non-Python files in a package."""
suspicious_extensions = {
".exe": "CRITICAL",
".dll": "CRITICAL",
".so": "HIGH",
".dylib": "HIGH",
".bin": "MEDIUM",
".dat": "LOW",
".whl": "MEDIUM",
}
for filepath in package_path.rglob("*"):
if filepath.is_file():
ext = filepath.suffix.lower()
if ext in suspicious_extensions:
result.findings.append(
NodeFinding(
severity=suspicious_extensions[ext],
category="suspicious_file",
message=f"Suspicious binary file: {filepath.name}",
file_path=str(filepath),
line_number=0,
node_package=package_name,
)
)
# Check for obfuscated Python (pyarmor, cython compiled)
if ext == ".pyd" or ext == ".pyc":
if "pyarmor" in filepath.name.lower():
result.findings.append(
NodeFinding(
severity="HIGH",
category="obfuscation",
message=f"PyArmor obfuscated file: {filepath.name}",
file_path=str(filepath),
line_number=0,
node_package=package_name,
)
)