@orionscanner/cli 0.1.0

package/src/orion/engine/taint/engine.py

@@ -0,0 +1,464 @@
+"""AST-based taint tracking engine for detecting user input flowing to dangerous sinks."""
+import ast
+import asyncio
+from dataclasses import dataclass, field
+from typing import Dict, List, Set, Optional, Any
+from orion.engine.config import EngineConfig, ResourceLimits
+from orion.engine.ir import OrionIR
+from orion.engine.results import EngineResult, Finding
+
+
+@dataclass
+class TaintState:
+    tainted_vars: Set[str] = field(default_factory=set)
+    tainted_attrs: Dict[str, Set[str]] = field(default_factory=dict)
+    tainted_indices: Dict[str, Set[Any]] = field(default_factory=dict)
+
+
+SOURCE_PATTERNS = {
+    "request": ["request.args", "request.form", "request.json", "request.files",
+                "request.values", "request.data", "request.get_json",
+                "request.GET", "request.POST", "request.query_params", "request.body"],
+    "input": ["input(", "sys.stdin", "sys.argv", "os.environ"],
+    "flask": ["request.args", "request.form", "request.json", "request.files",
+              "request.values", "request.data", "request.get_json"],
+    "fastapi": ["request.query_params", "request.body", "request.form", "request.json"],
+    "django": ["request.GET", "request.POST", "request.body", "request.FILES",
+               "request.COOKIES", "request.META"],
+    "tornado": ["self.get_argument", "self.get_arguments", "self.request.arguments",
+                "self.request.body"],
+    "bottle": ["request.query", "request.forms", "request.json", "request.files"],
+    "cherrypy": ["cherrypy.request.params", "cherrypy.request.body"],
+    "web": ["self.args", "self.form", "self.json", "self.request.query",
+            "self.request.form", "self.request.json"],
+}
+
+SINK_PATTERNS = {
+    "sql": ["execute", "executemany", "raw", "cursor", "executescript"],
+    "command": ["system", "popen", "subprocess.call", "subprocess.run",
+                "subprocess.Popen", "os.system", "os.popen", "commands.getstatusoutput",
+                "commands.getoutput"],
+    "code_exec": ["eval", "exec", "compile", "__import__", "execfile"],
+    "template": ["render_template_string", "render_template", "Markup", "MarkupSafe",
+                 "Template", "Environment", "jinja2.Environment"],
+    "file": ["open", "file", "io.open", "pathlib.Path.open"],
+    "deserialize": ["pickle.loads", "pickle.load", "yaml.load", "yaml.unsafe_load",
+                    "marshmallow.load", "json.loads"],
+    "ldap": ["ldap.initialize", "ldap.search", "ldap.bind"],
+    "xpath": ["etree.XPath", "lxml.etree.XPath"],
+    "regex": ["re.compile", "re.search", "re.match", "re.findall"],
+}
+
+SANITIZER_PATTERNS = {
+    "sql": ["escape_string", "quote", "quote_identifier", "mogrify",
+            "parameterize", "cursor.execute", "text(", "bindparam"],
+    "html": ["escape", "html_escape", "escapehtml", "markupsafe.escape",
+             "cgi.escape", "html.escape", "bleach.clean"],
+    "shell": ["shlex.quote", "pipes.quote", "subprocess.list2cmdline"],
+    "path": ["os.path.normpath", "os.path.realpath", "pathlib.Path.resolve"],
+    "url": ["urllib.parse.quote", "urllib.parse.quote_plus", "quote", "quote_plus"],
+}
+
+TAINT_PROPAGATION_OPS = {
+    ast.Add, ast.Mult, ast.Sub, ast.Div, ast.FloorDiv, ast.Mod, ast.Pow,
+    ast.LShift, ast.RShift, ast.BitOr, ast.BitXor, ast.BitAnd,
+    ast.And, ast.Or, ast.Not, ast.UAdd, ast.USub, ast.Invert,
+    ast.Eq, ast.NotEq, ast.Lt, ast.LtE, ast.Gt, ast.GtE,
+    ast.Is, ast.IsNot, ast.In, ast.NotIn,
+}
+
+
+class TaintVisitor(ast.NodeVisitor):
+    def __init__(self, filename: str, source_lines: List[str]):
+        self.filename = filename
+        self.source_lines = source_lines
+        self.findings: List[Finding] = []
+        self.taint_state = TaintState()
+        self.function_stack: List[str] = []
+        self.class_stack: List[str] = []
+        self.imports: Dict[str, str] = {}
+        self.current_function: Optional[str] = None
+        self.assignment_targets: Dict[str, Set[str]] = {}
+        self.call_depth = 0
+        self.max_call_depth = 50
+
+    def get_snippet(self, node: ast.AST, context: int = 2) -> str:
+        start = max(0, node.lineno - 1 - context)
+        end = min(len(self.source_lines), node.end_lineno + context if hasattr(node, 'end_lineno') and node.end_lineno else node.lineno + context)
+        return "\n".join(self.source_lines[start:end])
+
+    def is_tainted(self, node: ast.AST) -> bool:
+        if isinstance(node, ast.Name):
+            return node.id in self.taint_state.tainted_vars
+        elif isinstance(node, ast.Attribute):
+            if isinstance(node.value, ast.Name):
+                base = node.value.id
+                if base in self.taint_state.tainted_attrs:
+                    return node.attr in self.taint_state.tainted_attrs[base]
+            return self.is_tainted(node.value)
+        elif isinstance(node, ast.Subscript):
+            if isinstance(node.value, ast.Name):
+                base = node.value.id
+                if base in self.taint_state.tainted_indices:
+                    return True
+            return self.is_tainted(node.value)
+        elif isinstance(node, ast.BinOp):
+            return self.is_tainted(node.left) or self.is_tainted(node.right)
+        elif isinstance(node, ast.Call):
+            return self.is_tainted(node.func) or any(self.is_tainted(arg) for arg in node.args)
+        elif isinstance(node, ast.JoinedStr):
+            return any(self.is_tainted(v) for v in node.values if isinstance(v, ast.FormattedValue))
+        elif isinstance(node, ast.List):
+            return any(self.is_tainted(e) for e in node.elts)
+        elif isinstance(node, ast.Dict):
+            return any(self.is_tainted(v) for v in node.values)
+        elif isinstance(node, ast.Tuple):
+            return any(self.is_tainted(e) for e in node.elts)
+        return False
+
+    def mark_tainted(self, target: ast.AST, source_node: Optional[ast.AST] = None):
+        if isinstance(target, ast.Name):
+            self.taint_state.tainted_vars.add(target.id)
+        elif isinstance(target, ast.Attribute):
+            if isinstance(target.value, ast.Name):
+                base = target.value.id
+                if base not in self.taint_state.tainted_attrs:
+                    self.taint_state.tainted_attrs[base] = set()
+                self.taint_state.tainted_attrs[base].add(target.attr)
+        elif isinstance(target, ast.Subscript):
+            if isinstance(target.value, ast.Name):
+                base = target.value.id
+                if base not in self.taint_state.tainted_indices:
+                    self.taint_state.tainted_indices[base] = set()
+                self.taint_state.tainted_indices[base].add("all")
+        elif isinstance(target, ast.Tuple):
+            for elt in target.elts:
+                self.mark_tainted(elt, source_node)
+        elif isinstance(target, ast.List):
+            for elt in target.elts:
+                self.mark_tainted(elt, source_node)
+
+    def untaint(self, target: ast.AST):
+        if isinstance(target, ast.Name):
+            self.taint_state.tainted_vars.discard(target.id)
+        elif isinstance(target, ast.Attribute):
+            if isinstance(target.value, ast.Name):
+                base = target.value.id
+                if base in self.taint_state.tainted_attrs:
+                    self.taint_state.tainted_attrs[base].discard(target.attr)
+        elif isinstance(target, ast.Subscript):
+            if isinstance(target.value, ast.Name):
+                base = target.value.id
+                if base in self.taint_state.tainted_indices:
+                    self.taint_state.tainted_indices[base].discard("all")
+
+    def check_source(self, node: ast.AST) -> Optional[str]:
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name):
+                if node.func.id in ["input", "eval", "exec"]:
+                    return f"builtin_{node.func.id}"
+            elif isinstance(node.func, ast.Attribute):
+                attr = node.func.attr
+                if isinstance(node.func.value, ast.Name):
+                    base = node.func.value.id
+                    full = f"{base}.{attr}"
+                    for cat, patterns in SOURCE_PATTERNS.items():
+                        if any(p in full for p in patterns):
+                            return f"source_{cat}"
+        elif isinstance(node, ast.Attribute):
+            if isinstance(node.value, ast.Name):
+                base = node.value.id
+                full = f"{base}.{node.attr}"
+                for cat, patterns in SOURCE_PATTERNS.items():
+                    if any(p in full for p in patterns):
+                        return f"source_{cat}"
+        elif isinstance(node, ast.Name):
+            if node.id in ["request", "sys.argv", "os.environ"]:
+                return "source_global"
+        elif isinstance(node, ast.Subscript):
+            if isinstance(node.value, ast.Name):
+                base = node.value.id
+                if base in ["request", "os", "sys"]:
+                    return "source_subscript"
+        return None
+
+    def check_sink(self, node: ast.Call) -> Optional[tuple]:
+        if isinstance(node.func, ast.Name):
+            func_name = node.func.id
+            for cat, patterns in SINK_PATTERNS.items():
+                if func_name in patterns:
+                    return (cat, func_name)
+        elif isinstance(node.func, ast.Attribute):
+            attr = node.func.attr
+            for cat, patterns in SINK_PATTERNS.items():
+                if attr in patterns:
+                    if isinstance(node.func.value, ast.Name):
+                        return (cat, f"{node.func.value.id}.{attr}")
+                    return (cat, attr)
+        return None
+
+    def check_sanitizer(self, node: ast.Call) -> Optional[str]:
+        if isinstance(node.func, ast.Attribute):
+            attr = node.func.attr
+            for cat, patterns in SANITIZER_PATTERNS.items():
+                if attr in patterns:
+                    return cat
+        elif isinstance(node.func, ast.Name):
+            func_name = node.func.id
+            for cat, patterns in SANITIZER_PATTERNS.items():
+                if func_name in patterns:
+                    return cat
+        return None
+
+    def visit_Import(self, node: ast.Import):
+        for alias in node.names:
+            self.imports[alias.asname or alias.name] = alias.name
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node: ast.ImportFrom):
+        if node.module:
+            for alias in node.names:
+                self.imports[alias.asname or alias.name] = f"{node.module}.{alias.name}"
+        self.generic_visit(node)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        self.function_stack.append(node.name)
+        old_function = self.current_function
+        self.current_function = node.name
+        old_taint = TaintState()
+        old_taint.tainted_vars = self.taint_state.tainted_vars.copy()
+        old_taint.tainted_attrs = {k: v.copy() for k, v in self.taint_state.tainted_attrs.items()}
+        old_taint.tainted_indices = {k: v.copy() for k, v in self.taint_state.tainted_indices.items()}
+
+        for arg in node.args.args:
+            self.taint_state.tainted_vars.discard(arg.arg)
+            self.taint_state.tainted_vars.discard(arg.arg)
+
+        self.generic_visit(node)
+
+        self.taint_state = old_taint
+        self.current_function = old_function
+        self.function_stack.pop()
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
+        self.visit_FunctionDef(node)
+
+    def visit_ClassDef(self, node: ast.ClassDef):
+        self.class_stack.append(node.name)
+        self.generic_visit(node)
+        self.class_stack.pop()
+
+    def visit_Assign(self, node: ast.Assign):
+        for target in node.targets:
+            source_tainted = self.is_tainted(node.value)
+            source_type = self.check_source(node.value)
+            
+            if source_tainted or source_type:
+                self.mark_tainted(target, node.value)
+            else:
+                self.untaint(target)
+
+            if isinstance(target, ast.Name) and source_tainted:
+                self.assignment_targets.setdefault(target.id, set()).add(self.get_snippet(node))
+
+        self.generic_visit(node)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign):
+        if node.value:
+            self.visit_Assign(ast.Assign(targets=[node.target], value=node.value, lineno=node.lineno))
+        self.generic_visit(node)
+
+    def visit_AugAssign(self, node: ast.AugAssign):
+        if self.is_tainted(node.value):
+            self.mark_tainted(node.target, node.value)
+        self.generic_visit(node)
+
+    def visit_Call(self, node: ast.Call):
+        sink_info = self.check_sink(node)
+        sanitizer = self.check_sanitizer(node)
+        
+        args_tainted = any(self.is_tainted(arg) for arg in node.args)
+        kwargs_tainted = any(self.is_tainted(kw.value) for kw in node.keywords)
+        any_tainted = args_tainted or kwargs_tainted or self.is_tainted(node.func)
+
+        if sink_info and any_tainted:
+            cat, sink_name = sink_info
+            confidence = 0.85
+            evidence = []
+            
+            if args_tainted:
+                evidence.append("tainted_positional_args")
+            if kwargs_tainted:
+                evidence.append("tainted_keyword_args")
+            if self.is_tainted(node.func):
+                evidence.append("tainted_callable")
+            
+            if sanitizer:
+                confidence = 0.35
+                evidence.append(f"sanitized_by_{sanitizer}")
+            
+            if len(evidence) >= 2 or (sanitizer and confidence < 0.7):
+                confidence = max(confidence, 0.7)
+
+            cwe_map = {
+                "sql": "CWE-89",
+                "command": "CWE-78",
+                "code_exec": "CWE-94",
+                "template": "CWE-79",
+                "file": "CWE-22",
+                "deserialize": "CWE-502",
+                "ldap": "CWE-90",
+                "xpath": "CWE-643",
+                "regex": "CWE-1333",
+            }
+            
+            severity_map = {
+                "sql": "high",
+                "command": "critical",
+                "code_exec": "critical",
+                "template": "high",
+                "file": "medium",
+                "deserialize": "high",
+                "ldap": "high",
+                "xpath": "medium",
+                "regex": "medium",
+            }
+
+            self.findings.append(Finding(
+                engine="taint",
+                severity=severity_map.get(cat, "medium"),
+                title=f"Taint flow: user input to {cat} sink ({sink_name})",
+                description=f"Untrusted data flows into {sink_name} without proper sanitization. "
+                           f"Category: {cat}. Evidence: {', '.join(evidence)}",
+                location=f"{self.filename}:{node.lineno}",
+                metadata={
+                    "snippet": self.get_snippet(node),
+                    "cwe": cwe_map.get(cat, "CWE-20"),
+                    "confidence": confidence,
+                    "sink_category": cat,
+                    "sink_name": sink_name,
+                    "evidence": evidence,
+                    "sanitizer": sanitizer,
+                }
+            ))
+
+        if sanitizer and any_tainted:
+            for arg in node.args:
+                if self.is_tainted(arg):
+                    self.untaint(arg)
+            for kw in node.keywords:
+                if self.is_tainted(kw.value):
+                    self.untaint(kw.value)
+
+        for arg in node.args:
+            self.visit(arg)
+        for kw in node.keywords:
+            self.visit(kw.value)
+        self.visit(node.func)
+
+    def visit_BinOp(self, node: ast.BinOp):
+        left_tainted = self.is_tainted(node.left)
+        right_tainted = self.is_tainted(node.right)
+        if left_tainted or right_tainted:
+            pass
+        self.generic_visit(node)
+
+    def visit_JoinedStr(self, node: ast.JoinedStr):
+        tainted = any(self.is_tainted(v) for v in node.values if isinstance(v, ast.FormattedValue))
+        if tainted:
+            pass
+        self.generic_visit(node)
+
+    def visit_List(self, node: ast.List):
+        tainted = any(self.is_tainted(e) for e in node.elts)
+        if tainted:
+            pass
+        self.generic_visit(node)
+
+    def visit_Dict(self, node: ast.Dict):
+        tainted = any(self.is_tainted(v) for v in node.values)
+        if tainted:
+            pass
+        self.generic_visit(node)
+
+    def visit_Return(self, node: ast.Return):
+        if node.value and self.is_tainted(node.value):
+            pass
+        self.generic_visit(node)
+
+    def visit_Yield(self, node: ast.Yield):
+        if node.value and self.is_tainted(node.value):
+            pass
+        self.generic_visit(node)
+
+    def visit_YieldFrom(self, node: ast.YieldFrom):
+        if self.is_tainted(node.value):
+            pass
+        self.generic_visit(node)
+
+    def visit_Await(self, node: ast.Await):
+        if self.is_tainted(node.value):
+            pass
+        self.generic_visit(node)
+
+    def generic_visit(self, node: ast.AST):
+        if self.call_depth > self.max_call_depth:
+            return
+        self.call_depth += 1
+        try:
+            super().generic_visit(node)
+        finally:
+            self.call_depth -= 1
+
+
+class Engine:
+    name = "taint"
+    version = "1.0.0"
+
+    async def check_availability(self) -> bool:
+        return True
+
+    def resource_limits(self) -> ResourceLimits:
+        return ResourceLimits(max_cpu_seconds=60, max_memory_mb=512, max_files=50000)
+
+    async def run(self, ctx, config: EngineConfig) -> EngineResult:
+        # Handle both single IR and list of IRs
+        if not isinstance(ctx, list):
+            ctx = [ctx]
+        all_findings = []
+        errors = []
+        files_processed = 0
+
+        exclude_patterns = config.toggles.get("exclude", []) if config.toggles else []
+
+        for ir in ctx:
+            try:
+                if not ir.metadata.get("language") == "python":
+                    continue
+
+                rel_path = ir.metadata.get("path", "unknown.py")
+                if any(pattern in rel_path for pattern in exclude_patterns):
+                    continue
+
+                content = ir.content.decode("utf-8", errors="replace")
+                source_lines = content.splitlines()
+                tree = ast.parse(content, filename=rel_path)
+
+                visitor = TaintVisitor(rel_path, source_lines)
+                visitor.visit(tree)
+
+                all_findings.extend(visitor.findings)
+                files_processed += 1
+
+            except SyntaxError as e:
+                errors.append(f"Syntax error in {ir.metadata.get('path', 'unknown')}: {e}")
+            except Exception as e:
+                errors.append(f"Error processing {ir.metadata.get('path', 'unknown')}: {e}")
+
+        return EngineResult(
+            engine=self.name,
+            findings=all_findings,
+            errors=errors,
+            metrics={"files_processed": files_processed, "findings_count": len(all_findings)},
+            ok=len(errors) == 0
+        )