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)