mirror of
https://github.com/q3alique/codeflow
synced 2026-06-21 14:06:02 +00:00
q3alique
acc1b4f3e7
Static taint-analysis and visualization tool for source-code security review. Supports deep analysis for Python, JavaScript/TypeScript, Java, Go, and C#, with structural support for all other languages via the generic extractor. Outputs: interactive HTML report, LLM-ready Markdown review document, and optional Burp Suite JSON export. Self-bootstrapping launcher (run.py) requires no virtual environment.
1029 lines
42 KiB
Python
1029 lines
42 KiB
Python
from __future__ import annotations
|
|
import re
|
|
from pathlib import Path
|
|
from codeflow.discovery.file_scanner import FileInfo
|
|
from codeflow.analysis.base_extractor import ExtractionResult
|
|
from codeflow.analysis.ast_parser import parse_file, query_tree, get_node_text
|
|
from codeflow.models.node import Node, NodeType, make_node_id, make_module_node_id
|
|
from codeflow.models.edge import Edge, EdgeType
|
|
|
|
|
|
SINK_PATTERNS = {
|
|
"database": {
|
|
"execute": {".execute"},
|
|
"query": {".query"},
|
|
"executemany": {".executemany"},
|
|
},
|
|
"command_execution": {
|
|
"run": {"subprocess.run", "subprocess.call", "subprocess.Popen",
|
|
"subprocess.check_output", "os.system", "os.popen"},
|
|
"call": {"subprocess.call"},
|
|
"Popen": {"subprocess.Popen"},
|
|
"check_output": {"subprocess.check_output"},
|
|
"system": {"os.system"},
|
|
"popen": {"os.popen"},
|
|
},
|
|
"file_write": {
|
|
"open": {"open"},
|
|
},
|
|
"file_read": {
|
|
"open": {"open"},
|
|
"send_file": {"send_file"},
|
|
},
|
|
"http_request": {
|
|
"get": {"requests.get", "httpx.get"},
|
|
"post": {"requests.post", "httpx.post"},
|
|
"put": {"requests.put", "httpx.put"},
|
|
"delete": {"requests.delete", "httpx.delete"},
|
|
},
|
|
"code_execution": {
|
|
"eval": {"eval"},
|
|
"exec": {"exec"},
|
|
"compile": {"compile"},
|
|
},
|
|
"deserialization": {
|
|
"loads": {"pickle.loads"},
|
|
"load": {"pickle.load", "yaml.load"},
|
|
},
|
|
"template_injection": {
|
|
"render_template_string": {"render_template_string"},
|
|
},
|
|
}
|
|
|
|
_PRIORITY_ORDER = [
|
|
"code_execution", "command_execution", "database",
|
|
"template_injection", "deserialization",
|
|
"file_write", "file_read", "http_request",
|
|
]
|
|
|
|
|
|
def _detect_framework(source_text: str) -> str:
|
|
text_lower = source_text.lower()
|
|
if "from flask" in text_lower or "import flask" in text_lower:
|
|
return "flask"
|
|
if "from django" in text_lower or "import django" in text_lower:
|
|
return "django"
|
|
if "from fastapi" in text_lower or "import fastapi" in text_lower:
|
|
return "fastapi"
|
|
return "unknown"
|
|
|
|
|
|
class PythonExtractor:
|
|
def __init__(self):
|
|
self.file_stem: str = ""
|
|
self.file_path: str = ""
|
|
self.framework: str = "unknown"
|
|
self.source_bytes: bytes = b""
|
|
self.file_lines: list[str] = []
|
|
|
|
def extract(self, file_info: FileInfo) -> ExtractionResult:
|
|
result = ExtractionResult(file_path=file_info.path)
|
|
self.file_path = file_info.path
|
|
self.file_stem = Path(file_info.path).stem
|
|
|
|
try:
|
|
self.source_bytes = Path(file_info.path).read_bytes()
|
|
self.file_lines = Path(file_info.path).read_text(encoding="utf-8", errors="replace").splitlines()
|
|
except (OSError, IOError):
|
|
result.errors.append(f"Cannot read file: {file_info.path}")
|
|
return result
|
|
|
|
tree = parse_file(file_info.path, file_info.language)
|
|
if tree is None:
|
|
result.errors.append(f"Could not parse: {file_info.path}")
|
|
return result
|
|
|
|
source_text = self.source_bytes.decode("utf-8", errors="replace")
|
|
self.framework = _detect_framework(source_text)
|
|
|
|
self._extract_module_node(result)
|
|
self._extract_classes(result, tree)
|
|
self._extract_functions(result, tree)
|
|
self._extract_calls(result, tree)
|
|
self._extract_imports(result, tree)
|
|
self._extract_entry_points(result, tree)
|
|
self._extract_sources(result, tree)
|
|
self._extract_sinks(result, tree)
|
|
self._extract_env_vars(result, tree)
|
|
|
|
self._ensure_unique_ids(result)
|
|
return result
|
|
|
|
def _ensure_unique_ids(self, result: ExtractionResult):
|
|
seen: set[str] = set()
|
|
deduped: list[Node] = []
|
|
for node in result.nodes:
|
|
if node.id not in seen:
|
|
seen.add(node.id)
|
|
deduped.append(node)
|
|
result.nodes = deduped
|
|
|
|
def _get_func_line_range(self, func_node) -> tuple[int, int]:
|
|
try:
|
|
start = func_node.start_point[0] + 1
|
|
end = func_node.end_point[0] + 1
|
|
return start, end
|
|
except Exception:
|
|
return 0, 0
|
|
|
|
def _is_method(self, func_node, tree) -> tuple[bool, str]:
|
|
try:
|
|
parent = func_node.parent
|
|
while parent:
|
|
if parent.type == "class_definition":
|
|
for child in parent.children:
|
|
if child.type == "identifier":
|
|
return True, self.source_bytes[child.start_byte:child.end_byte].decode("utf-8", errors="replace")
|
|
parent = parent.parent
|
|
except Exception:
|
|
pass
|
|
return False, ""
|
|
|
|
def _extract_module_node(self, result: ExtractionResult):
|
|
node_id = make_module_node_id(self.file_stem)
|
|
node = Node(
|
|
id=node_id,
|
|
node_type=NodeType.MODULE,
|
|
label=Path(self.file_path).name,
|
|
file_path=self.file_path,
|
|
line_start=1,
|
|
line_end=len(self.file_lines),
|
|
metadata={
|
|
"language": "python",
|
|
"imports": [],
|
|
},
|
|
)
|
|
result.nodes.append(node)
|
|
|
|
def _extract_functions(self, result: ExtractionResult, tree):
|
|
query = """
|
|
(function_definition
|
|
name: (identifier) @func_name) @func_def
|
|
"""
|
|
captures = query_tree(tree, "python", query, capture_filter="func_def")
|
|
for func_node, _ in captures:
|
|
func_name = get_node_text(func_node.child_by_field_name("name"), self.source_bytes)
|
|
if not func_name:
|
|
continue
|
|
is_method, class_name = self._is_method(func_node, tree)
|
|
qualified = f"{class_name}.{func_name}" if is_method else func_name
|
|
node_id = make_node_id(self.file_stem, qualified)
|
|
start, end = self._get_func_line_range(func_node)
|
|
|
|
params = []
|
|
params_node = func_node.child_by_field_name("parameters")
|
|
if params_node:
|
|
for child in params_node.children:
|
|
if child.type in ("identifier", "typed_parameter"):
|
|
try:
|
|
pname = self.source_bytes[child.start_byte:child.end_byte].decode("utf-8", errors="replace").split(":")[0].strip()
|
|
if pname:
|
|
params.append(pname)
|
|
except Exception:
|
|
pass
|
|
|
|
node = Node(
|
|
id=node_id,
|
|
node_type=NodeType.FUNCTION,
|
|
label=func_name[:40],
|
|
file_path=self.file_path,
|
|
line_start=start,
|
|
line_end=end,
|
|
metadata={
|
|
"qualified_name": qualified,
|
|
"params": params,
|
|
"is_method": is_method,
|
|
"class_name": class_name if is_method else None,
|
|
},
|
|
)
|
|
result.nodes.append(node)
|
|
result.edges.append(Edge(
|
|
source_id=make_module_node_id(self.file_stem),
|
|
target_id=node_id,
|
|
edge_type=EdgeType.CONTAINS,
|
|
))
|
|
|
|
def _extract_classes(self, result: ExtractionResult, tree):
|
|
query = """
|
|
(class_definition
|
|
name: (identifier) @class_name) @class_def
|
|
"""
|
|
captures = query_tree(tree, "python", query, capture_filter="class_def")
|
|
for class_node, _ in captures:
|
|
class_name_node = class_node.child_by_field_name("name")
|
|
if not class_name_node:
|
|
continue
|
|
class_name = get_node_text(class_name_node, self.source_bytes)
|
|
if not class_name:
|
|
continue
|
|
node_id = make_node_id(self.file_stem, class_name)
|
|
start = class_node.start_point[0] + 1
|
|
end = class_node.end_point[0] + 1
|
|
node = Node(
|
|
id=node_id,
|
|
node_type=NodeType.CLASS,
|
|
label=class_name,
|
|
file_path=self.file_path,
|
|
line_start=start,
|
|
line_end=end,
|
|
metadata={},
|
|
)
|
|
result.nodes.append(node)
|
|
result.edges.append(Edge(
|
|
source_id=make_module_node_id(self.file_stem),
|
|
target_id=node_id,
|
|
edge_type=EdgeType.CONTAINS,
|
|
))
|
|
|
|
def _extract_imports(self, result: ExtractionResult, tree):
|
|
module_id = make_module_node_id(self.file_stem)
|
|
module_node = None
|
|
for n in result.nodes:
|
|
if n.id == module_id:
|
|
module_node = n
|
|
break
|
|
|
|
query = """
|
|
(import_statement
|
|
name: (dotted_name) @module_name) @import_stmt
|
|
|
|
(import_from_statement
|
|
module_name: (dotted_name) @module_name) @import_stmt
|
|
"""
|
|
captures = query_tree(tree, "python", query, capture_filter="import_stmt")
|
|
imports_list = []
|
|
for node, _ in captures:
|
|
try:
|
|
text = self.source_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace").strip()
|
|
if text:
|
|
imports_list.append(text)
|
|
except Exception:
|
|
pass
|
|
|
|
if module_node:
|
|
module_node.metadata["imports"] = imports_list
|
|
|
|
def _extract_calls(self, result: ExtractionResult, tree):
|
|
func_nodes = {}
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
func_nodes[n.id] = n
|
|
|
|
call_query = """
|
|
(call
|
|
function: (identifier) @callee) @call_expr
|
|
"""
|
|
captures = query_tree(tree, "python", call_query)
|
|
unresolved: list[dict] = []
|
|
|
|
for call_node, capture_name in captures:
|
|
if capture_name != "callee":
|
|
continue
|
|
callee_text = get_node_text(call_node, self.source_bytes)
|
|
if not callee_text:
|
|
continue
|
|
caller_id = self._find_containing_function(call_node, result)
|
|
if caller_id is None:
|
|
continue
|
|
callee_name = callee_text
|
|
found = False
|
|
for nid, n in func_nodes.items():
|
|
if nid.endswith(f"::{callee_name}") or n.metadata.get("qualified_name") == callee_name:
|
|
result.edges.append(Edge(
|
|
source_id=caller_id,
|
|
target_id=nid,
|
|
edge_type=EdgeType.CALLS,
|
|
))
|
|
found = True
|
|
break
|
|
if not found:
|
|
try:
|
|
line = call_node.start_point[0] + 1
|
|
except Exception:
|
|
line = 0
|
|
unresolved.append({
|
|
"caller_id": caller_id,
|
|
"callee_name": callee_name,
|
|
"line": line,
|
|
})
|
|
|
|
method_call_query = """
|
|
(call
|
|
function: (attribute
|
|
attribute: (identifier) @callee)) @call_expr
|
|
"""
|
|
method_captures = query_tree(tree, "python", method_call_query)
|
|
for call_node, capture_name in method_captures:
|
|
if capture_name != "callee":
|
|
continue
|
|
callee_text = get_node_text(call_node, self.source_bytes)
|
|
if not callee_text:
|
|
continue
|
|
caller_id = self._find_containing_function(call_node, result)
|
|
if caller_id is None:
|
|
continue
|
|
callee_name = callee_text
|
|
found = False
|
|
for nid, n in func_nodes.items():
|
|
if nid.endswith(f"::{callee_name}") or n.metadata.get("qualified_name") == callee_name:
|
|
result.edges.append(Edge(
|
|
source_id=caller_id,
|
|
target_id=nid,
|
|
edge_type=EdgeType.CALLS,
|
|
))
|
|
found = True
|
|
break
|
|
if not found:
|
|
try:
|
|
line = call_node.start_point[0] + 1
|
|
except Exception:
|
|
line = 0
|
|
unresolved.append({
|
|
"caller_id": caller_id,
|
|
"callee_name": callee_name,
|
|
"line": line,
|
|
})
|
|
|
|
result.metadata["unresolved_calls"] = unresolved
|
|
|
|
def _find_containing_function(self, node, result: ExtractionResult) -> str | None:
|
|
try:
|
|
node_line = node.start_point[0] + 1
|
|
except Exception:
|
|
return None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= node_line <= n.line_end:
|
|
return n.id
|
|
return None
|
|
|
|
def _extract_entry_points(self, result: ExtractionResult, tree):
|
|
if self.framework not in ("flask", "fastapi", "django"):
|
|
return
|
|
|
|
query = """
|
|
(decorated_definition
|
|
(decorator
|
|
(call
|
|
function: (attribute
|
|
attribute: (identifier) @dec_method)
|
|
arguments: (argument_list
|
|
(string) @route_path))) @decorator
|
|
(function_definition
|
|
name: (identifier) @func_name)) @decorated
|
|
"""
|
|
captures = query_tree(tree, "python", query, capture_filter="decorated")
|
|
|
|
path_params_re = re.compile(r"<(?:\w+:)?(\w+)>")
|
|
|
|
for decorated_node, _ in captures:
|
|
dec_method = None
|
|
route_path = None
|
|
func_name = None
|
|
dec_args = None
|
|
|
|
for child in decorated_node.children:
|
|
if child.type == "decorator":
|
|
for child2 in child.children:
|
|
if child2.type == "call":
|
|
func_part = None
|
|
for cc in child2.children:
|
|
if cc.type == "attribute":
|
|
for ccc in cc.children:
|
|
if ccc.type == "identifier":
|
|
dec_method = get_node_text(ccc, self.source_bytes)
|
|
if ccc.type == "attribute" and ccc.children:
|
|
pass
|
|
elif cc.type == "argument_list":
|
|
dec_args = cc
|
|
for arg in cc.children:
|
|
if arg.type == "string":
|
|
route_path_clean = get_node_text(arg, self.source_bytes).strip("'\"")
|
|
if route_path_clean:
|
|
route_path = route_path_clean
|
|
break
|
|
if func_part is None:
|
|
pass
|
|
elif child.type == "function_definition":
|
|
name_node = child.child_by_field_name("name")
|
|
if name_node:
|
|
func_name = get_node_text(name_node, self.source_bytes)
|
|
|
|
if not all([dec_method, route_path, func_name]):
|
|
continue
|
|
|
|
is_flask = dec_method in ("route", "get", "post", "put", "delete", "patch", "options")
|
|
is_fastapi = dec_method in ("get", "post", "put", "delete", "patch", "options", "head") and self.framework == "fastapi"
|
|
|
|
if not (is_flask or is_fastapi):
|
|
continue
|
|
|
|
if dec_method == "route":
|
|
http_method = "GET"
|
|
if dec_args:
|
|
for arg in dec_args.children:
|
|
if arg.type == "keyword_argument":
|
|
kw_text = get_node_text(arg, self.source_bytes)
|
|
if "methods" in kw_text:
|
|
match = re.search(r"\[(.*?)\]", kw_text)
|
|
if match:
|
|
methods = [m.strip().strip("'\"").upper() for m in match.group(1).split(",")]
|
|
http_method = ",".join(methods)
|
|
else:
|
|
http_method = dec_method.upper()
|
|
|
|
func_id = make_node_id(self.file_stem, func_name)
|
|
for n in result.nodes:
|
|
if n.id == func_id and n.node_type == NodeType.FUNCTION:
|
|
n.node_type = NodeType.ENTRY_POINT
|
|
n.is_entry_point = True
|
|
n.label = f"{http_method} {route_path}"
|
|
n.metadata["http_method"] = http_method
|
|
n.metadata["route_path"] = route_path
|
|
n.metadata["framework"] = self.framework
|
|
break
|
|
|
|
route_params = path_params_re.findall(route_path)
|
|
for rp in route_params:
|
|
source_id = make_node_id(self.file_stem, f"{func_name}::{rp}")
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label=f"param: {rp} [url_param]",
|
|
file_path=self.file_path,
|
|
line_start=0,
|
|
line_end=0,
|
|
is_tainted=False,
|
|
metadata={
|
|
"param_name": rp,
|
|
"input_type": "url_param",
|
|
"framework": self.framework,
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=func_id,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
def _extract_sources(self, result: ExtractionResult, tree):
|
|
source_text = self.source_bytes.decode("utf-8", errors="replace")
|
|
|
|
if "sys.argv" in source_text:
|
|
lines = source_text.splitlines()
|
|
for i, line in enumerate(lines):
|
|
if "sys.argv" in line:
|
|
line_num = i + 1
|
|
containing_func = None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line_num <= n.line_end:
|
|
containing_func = n.id
|
|
break
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::argv")
|
|
if not any(n.id == source_id for n in result.nodes):
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label="param: argv [cli_arg]",
|
|
file_path=self.file_path,
|
|
line_start=line_num,
|
|
line_end=line_num,
|
|
metadata={
|
|
"param_name": "argv",
|
|
"input_type": "cli_arg",
|
|
"framework": "sys.argv",
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
if "argparse" in source_text or "from argparse" in source_text:
|
|
for i, line in enumerate(source_text.splitlines()):
|
|
if "add_argument" in line:
|
|
match = re.search(r'["\']([^"\']+)["\']', line)
|
|
if match:
|
|
param_name = match.group(1).lstrip("-")
|
|
line_num = i + 1
|
|
containing_func = None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line_num <= n.line_end:
|
|
containing_func = n.id
|
|
break
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::{param_name}")
|
|
if not any(n.id == source_id for n in result.nodes):
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label=f"param: {param_name} [cli_arg]",
|
|
file_path=self.file_path,
|
|
line_start=line_num,
|
|
line_end=line_num,
|
|
metadata={
|
|
"param_name": param_name,
|
|
"input_type": "cli_arg",
|
|
"framework": "argparse",
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
click_pattern = re.compile(r"@click\.(argument|option)\b")
|
|
for match in click_pattern.finditer(source_text):
|
|
line_num = source_text[:match.start()].count("\n") + 1
|
|
containing_func = None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line_num <= n.line_end:
|
|
containing_func = n.id
|
|
break
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::click_arg")
|
|
if not any(n.id == source_id for n in result.nodes):
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label="param: click_arg [cli_arg]",
|
|
file_path=self.file_path,
|
|
line_start=line_num,
|
|
line_end=line_num,
|
|
metadata={
|
|
"param_name": "click_arg",
|
|
"input_type": "cli_arg",
|
|
"framework": "click",
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
if self.framework == "unknown":
|
|
return
|
|
|
|
func_id_map = {}
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
func_id_map[n.label] = n.id
|
|
if n.label.startswith(("GET ", "POST ", "PUT ", "DELETE ")):
|
|
func_id_map[n.metadata.get("route_path", n.label)] = n.id
|
|
|
|
request_sources = {
|
|
"args": ("query_string", "get"),
|
|
"form": ("form_field", "get"),
|
|
"files": ("file_upload", "get"),
|
|
"cookies": ("cookie", "get"),
|
|
"headers": ("header", "get"),
|
|
}
|
|
|
|
access_query = """
|
|
(call
|
|
function: (attribute
|
|
object: (attribute
|
|
object: (identifier) @obj
|
|
attribute: (identifier) @attr)
|
|
attribute: (identifier) @method)
|
|
arguments: (argument_list
|
|
[(string) (integer)] @arg)) @call_expr
|
|
"""
|
|
captures = query_tree(tree, "python", access_query, capture_filter="call_expr")
|
|
for node, _ in captures:
|
|
obj = None
|
|
attr = None
|
|
arg = None
|
|
text = get_node_text(node, self.source_bytes)
|
|
func_child = node.child_by_field_name("function")
|
|
if func_child:
|
|
for c2 in func_child.children:
|
|
if c2.type == "attribute":
|
|
parts = []
|
|
for c3 in c2.children:
|
|
if c3.type == "identifier":
|
|
parts.append(c3)
|
|
elif c3.type == "attribute":
|
|
for c4 in c3.children:
|
|
if c4.type == "identifier":
|
|
parts.append(c4)
|
|
if len(parts) >= 2:
|
|
obj = get_node_text(parts[0], self.source_bytes)
|
|
attr = get_node_text(parts[-1], self.source_bytes)
|
|
args_child = node.child_by_field_name("arguments")
|
|
if args_child:
|
|
for arg_child in args_child.children:
|
|
if arg_child.type == "string" and arg is None:
|
|
arg = get_node_text(arg_child, self.source_bytes).strip("'\"")
|
|
|
|
if obj == "request" and attr in request_sources:
|
|
input_type, _ = request_sources[attr]
|
|
param_name = arg or f"*{attr}*"
|
|
containing_func = self._find_containing_function(node, result)
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::{param_name}")
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label=f"param: {param_name} [{input_type}]"[:40],
|
|
file_path=self.file_path,
|
|
line_start=node.start_point[0] + 1,
|
|
line_end=node.end_point[0] + 1,
|
|
is_tainted=False,
|
|
metadata={
|
|
"param_name": param_name,
|
|
"input_type": input_type,
|
|
"framework": self.framework,
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
subscript_query = """
|
|
(subscript
|
|
value: (attribute
|
|
object: (identifier) @obj
|
|
attribute: (identifier) @attr)
|
|
(string) @key) @subscript_expr
|
|
"""
|
|
sub_captures = query_tree(tree, "python", subscript_query, capture_filter="subscript_expr")
|
|
for node, _ in sub_captures:
|
|
obj = None
|
|
attr = None
|
|
key = None
|
|
for child in node.children:
|
|
if child.type == "attribute":
|
|
for c2 in child.children:
|
|
if c2.type == "identifier":
|
|
if obj is None:
|
|
obj = get_node_text(c2, self.source_bytes)
|
|
else:
|
|
attr = get_node_text(c2, self.source_bytes)
|
|
elif child.type == "string":
|
|
key = get_node_text(child, self.source_bytes).strip("'\"")
|
|
|
|
if obj != "request":
|
|
continue
|
|
|
|
input_type_map = {
|
|
"args": "query_string",
|
|
"form": "form_field",
|
|
"files": "file_upload",
|
|
"data": "raw_body",
|
|
"cookies": "cookie",
|
|
"headers": "header",
|
|
"json": "json_body",
|
|
}
|
|
|
|
if attr in input_type_map:
|
|
param_name = key or f"*{attr}*"
|
|
input_type = input_type_map[attr]
|
|
containing_func = self._find_containing_function(node, result)
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::{param_name}")
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label=f"param: {param_name} [{input_type}]"[:40],
|
|
file_path=self.file_path,
|
|
line_start=node.start_point[0] + 1,
|
|
line_end=node.end_point[0] + 1,
|
|
metadata={
|
|
"param_name": param_name,
|
|
"input_type": input_type,
|
|
"framework": self.framework,
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
if self.framework == "flask":
|
|
json_pattern = re.compile(r"request\.(get_json\(\)|json)\b")
|
|
for match in json_pattern.finditer(self.source_bytes.decode("utf-8", errors="replace")):
|
|
line = self.source_bytes[:match.start()].count(b"\n") + 1
|
|
containing_func = None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line <= n.line_end:
|
|
containing_func = n.id
|
|
break
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::*json_body*")
|
|
if not any(n.id == source_id for n in result.nodes):
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label="param: *json_body* [json_body]",
|
|
file_path=self.file_path,
|
|
line_start=line,
|
|
line_end=line,
|
|
metadata={
|
|
"param_name": "*json_body*",
|
|
"input_type": "json_body",
|
|
"framework": "flask",
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
data_pattern = re.compile(r"request\.data\b")
|
|
for match in data_pattern.finditer(self.source_bytes.decode("utf-8", errors="replace")):
|
|
line = self.source_bytes[:match.start()].count(b"\n") + 1
|
|
containing_func = None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line <= n.line_end:
|
|
containing_func = n.id
|
|
break
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::*raw_body*")
|
|
if not any(n.id == source_id for n in result.nodes):
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label="param: *raw_body* [raw_body]",
|
|
file_path=self.file_path,
|
|
line_start=line,
|
|
line_end=line,
|
|
metadata={
|
|
"param_name": "*raw_body*",
|
|
"input_type": "raw_body",
|
|
"framework": "flask",
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
input_pattern = re.compile(r"\binput\(\)")
|
|
for match in input_pattern.finditer(source_text):
|
|
line_num = source_text[:match.start()].count("\n") + 1
|
|
containing_func = None
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line_num <= n.line_end:
|
|
containing_func = n.id
|
|
break
|
|
if containing_func:
|
|
func_name_part = containing_func.split("::")[-1]
|
|
source_id = make_node_id(self.file_stem, f"{func_name_part}::stdin_input")
|
|
if not any(n.id == source_id for n in result.nodes):
|
|
source_node = Node(
|
|
id=source_id,
|
|
node_type=NodeType.SOURCE,
|
|
label="param: stdin_input [stdin]",
|
|
file_path=self.file_path,
|
|
line_start=line_num,
|
|
line_end=line_num,
|
|
metadata={
|
|
"param_name": "stdin_input",
|
|
"input_type": "stdin",
|
|
"framework": "builtin",
|
|
},
|
|
)
|
|
result.nodes.append(source_node)
|
|
result.edges.append(Edge(
|
|
source_id=source_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
def _extract_env_vars(self, result: ExtractionResult, tree):
|
|
source_text = self.source_bytes.decode("utf-8", errors="replace")
|
|
env_pattern = re.compile(r"""
|
|
os\.environ(?:\.get)?\s*\(\s*['"](\w+)['"]
|
|
|os\.environ\s*\[\s*['"](\w+)['"]\s*\]
|
|
|os\.getenv\s*\(\s*['"](\w+)['"]
|
|
|environ\.get\s*\(\s*['"](\w+)['"]
|
|
""", re.VERBOSE)
|
|
|
|
for match in env_pattern.finditer(source_text):
|
|
var_name = next(g for g in match.groups() if g is not None)
|
|
line_num = source_text[:match.start()].count("\n") + 1
|
|
containing_func = self._find_containing_function_from_line(line_num, result)
|
|
env_id = make_node_id(self.file_stem, f"ENV::{var_name}")
|
|
if any(n.id == env_id for n in result.nodes):
|
|
continue
|
|
env_node = Node(
|
|
id=env_id,
|
|
node_type=NodeType.ENV_VAR,
|
|
label=f"env: {var_name}",
|
|
file_path=self.file_path,
|
|
line_start=line_num,
|
|
line_end=line_num,
|
|
metadata={"var_name": var_name},
|
|
)
|
|
result.nodes.append(env_node)
|
|
if containing_func:
|
|
result.edges.append(Edge(
|
|
source_id=env_id,
|
|
target_id=containing_func,
|
|
edge_type=EdgeType.DATA_FLOW,
|
|
))
|
|
|
|
def _find_containing_function_from_line(self, line_num: int, result: ExtractionResult) -> str | None:
|
|
for n in result.nodes:
|
|
if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
|
|
if n.line_start <= line_num <= n.line_end:
|
|
return n.id
|
|
return None
|
|
|
|
def _extract_sinks(self, result: ExtractionResult, tree):
|
|
source_text = self.source_bytes.decode("utf-8", errors="replace")
|
|
|
|
call_query = """
|
|
(call
|
|
function: (identifier) @func_name) @call_expr
|
|
"""
|
|
captures = query_tree(tree, "python", call_query, capture_filter="call_expr")
|
|
|
|
seen_sink_ids: set[str] = set()
|
|
|
|
for call_node, _ in captures:
|
|
func_name = None
|
|
for child in call_node.children:
|
|
if child.type == "identifier":
|
|
func_name = get_node_text(child, self.source_bytes)
|
|
break
|
|
if not func_name:
|
|
continue
|
|
|
|
sink_info = self._classify_sink(func_name)
|
|
if sink_info is None:
|
|
continue
|
|
|
|
sink_type, sink_func = sink_info
|
|
|
|
containing_func = self._find_containing_function(call_node, result)
|
|
if containing_func is None:
|
|
continue
|
|
|
|
try:
|
|
line = call_node.start_point[0] + 1
|
|
except Exception:
|
|
line = 0
|
|
|
|
sink_id = make_node_id(self.file_stem, f"SINK::{sink_type}::{sink_func}::{line}")
|
|
if sink_id in seen_sink_ids:
|
|
continue
|
|
seen_sink_ids.add(sink_id)
|
|
|
|
sink_node = Node(
|
|
id=sink_id,
|
|
node_type=NodeType.SINK,
|
|
label=f"sink: {sink_func} [{sink_type}]",
|
|
file_path=self.file_path,
|
|
line_start=line,
|
|
line_end=line,
|
|
metadata={
|
|
"sink_type": sink_type,
|
|
"function_name": sink_func,
|
|
},
|
|
)
|
|
result.nodes.append(sink_node)
|
|
result.edges.append(Edge(
|
|
source_id=containing_func,
|
|
target_id=sink_id,
|
|
edge_type=EdgeType.CALLS,
|
|
))
|
|
|
|
method_call_query = """
|
|
(call
|
|
function: (attribute
|
|
attribute: (identifier) @method_name)) @call_expr
|
|
"""
|
|
method_captures = query_tree(tree, "python", method_call_query, capture_filter="call_expr")
|
|
for call_node, _ in method_captures:
|
|
method_name = None
|
|
func_child = call_node.child_by_field_name("function")
|
|
if func_child:
|
|
method_ids = [c for c in func_child.children if c.type == "identifier"]
|
|
if method_ids:
|
|
method_name = get_node_text(method_ids[-1], self.source_bytes)
|
|
if not method_name:
|
|
continue
|
|
|
|
sink_info = self._classify_sink(method_name)
|
|
if sink_info is None:
|
|
continue
|
|
|
|
sink_type, sink_func = sink_info
|
|
|
|
containing_func = self._find_containing_function(call_node, result)
|
|
if containing_func is None:
|
|
continue
|
|
|
|
try:
|
|
line = call_node.start_point[0] + 1
|
|
except Exception:
|
|
line = 0
|
|
|
|
sink_id = make_node_id(self.file_stem, f"SINK::{sink_type}::{sink_func}::{line}")
|
|
if sink_id in seen_sink_ids:
|
|
continue
|
|
seen_sink_ids.add(sink_id)
|
|
|
|
sink_node = Node(
|
|
id=sink_id,
|
|
node_type=NodeType.SINK,
|
|
label=f"sink: {sink_func} [{sink_type}]",
|
|
file_path=self.file_path,
|
|
line_start=line,
|
|
line_end=line,
|
|
metadata={
|
|
"sink_type": sink_type,
|
|
"function_name": sink_func,
|
|
},
|
|
)
|
|
result.nodes.append(sink_node)
|
|
result.edges.append(Edge(
|
|
source_id=containing_func,
|
|
target_id=sink_id,
|
|
edge_type=EdgeType.CALLS,
|
|
))
|
|
|
|
open_pattern = re.compile(r"\bopen\s*\([^)]*\)")
|
|
for match in open_pattern.finditer(source_text):
|
|
line_num = source_text[:match.start()].count("\n") + 1
|
|
mode_match = re.search(r'["\']([rwa+b]*)["\']', match.group())
|
|
if mode_match:
|
|
mode = mode_match.group(1)
|
|
else:
|
|
mode = "r"
|
|
|
|
sink_type = "file_write" if ("w" in mode or "a" in mode) else "file_read"
|
|
|
|
containing_func = self._find_containing_function_from_line(line_num, result)
|
|
if containing_func is None:
|
|
continue
|
|
|
|
sink_id = make_node_id(self.file_stem, f"SINK::{sink_type}::open::{line_num}")
|
|
if sink_id in seen_sink_ids:
|
|
continue
|
|
seen_sink_ids.add(sink_id)
|
|
|
|
sink_node = Node(
|
|
id=sink_id,
|
|
node_type=NodeType.SINK,
|
|
label=f"sink: open [{sink_type}]",
|
|
file_path=self.file_path,
|
|
line_start=line_num,
|
|
line_end=line_num,
|
|
metadata={
|
|
"sink_type": sink_type,
|
|
"function_name": "open",
|
|
},
|
|
)
|
|
result.nodes.append(sink_node)
|
|
result.edges.append(Edge(
|
|
source_id=containing_func,
|
|
target_id=sink_id,
|
|
edge_type=EdgeType.CALLS,
|
|
))
|
|
|
|
def _classify_sink(self, func_name: str) -> tuple[str, str] | None:
|
|
for sink_type in _PRIORITY_ORDER:
|
|
for sink_func, patterns in SINK_PATTERNS[sink_type].items():
|
|
for pat in patterns:
|
|
if func_name == pat.split(".")[-1] or func_name == pat:
|
|
return sink_type, sink_func
|
|
return None
|
|
|
|
|
|
def extract(file_info: FileInfo) -> ExtractionResult:
|
|
extractor = PythonExtractor()
|
|
return extractor.extract(file_info)
|