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>
350 lines
12 KiB
Python
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,
|
|
)
|
|
)
|