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 catalogue ──────────────────────────────────────────────────────────── JAVA_SINK_PATTERNS: dict[str, dict[str, str]] = { "database": { "execute": r"\.execute\s*\(", "executeQuery": r"\.executeQuery\s*\(", "executeUpdate": r"\.executeUpdate\s*\(", "executeBatch": r"\.executeBatch\s*\(", "prepareStatement": r"\.prepareStatement\s*\(", "prepareCall": r"\.prepareCall\s*\(", }, "command_execution": { "exec": r"\.exec\s*\(", "ProcessBuilder": r"new\s+ProcessBuilder\s*\(", "start": r"\.start\s*\(\s*\)", # ProcessBuilder.start() }, "deserialization": { "readObject": r"\.readObject\s*\(", "readValue": r"\.readValue\s*\(", "fromJson": r"\.fromJson\s*\(", "deserialize": r"\.deserialize\s*\(", }, "file_write": { "Files.write": r"Files\.write\s*\(", "FileWriter": r"new\s+FileWriter\s*\(", "PrintWriter": r"new\s+PrintWriter\s*\(", "FileOutputStream": r"new\s+FileOutputStream\s*\(", "write": r"\.write\s*\(", }, "file_read": { "Files.readAllBytes": r"Files\.readAllBytes\s*\(", "Files.readAllLines": r"Files\.readAllLines\s*\(", "FileReader": r"new\s+FileReader\s*\(", "FileInputStream": r"new\s+FileInputStream\s*\(", }, "code_execution": { "eval": r"\.eval\s*\(", "newInstance": r"\.newInstance\s*\(", "Class.forName": r"Class\.forName\s*\(", }, "template_injection": { "process": r"\.process\s*\(", "evaluate": r"\.evaluate\s*\(", "merge": r"\.merge\s*\(", }, "http_request": { "send": r"\.send\s*\(", "execute": r"httpClient\.execute\s*\(", }, } _PRIORITY_ORDER = [ "code_execution", "command_execution", "database", "template_injection", "deserialization", "file_write", "file_read", "http_request", ] # Spring MVC / Spring Boot HTTP mapping annotations _HTTP_ANNOTATIONS: dict[str, str] = { "GetMapping": "GET", "PostMapping": "POST", "PutMapping": "PUT", "DeleteMapping": "DELETE", "PatchMapping": "PATCH", "RequestMapping": "REQUEST", } # Spring parameter binding annotations → input_type _PARAM_INPUT_TYPES: dict[str, str] = { "PathVariable": "url_param", "RequestParam": "query_string", "RequestBody": "json_body", "RequestHeader": "header", "CookieValue": "cookie", "MatrixVariable":"url_param", } def _detect_java_framework(source_text: str) -> str: if "@RestController" in source_text or "@Controller" in source_text: return "spring" if "HttpServletRequest" in source_text or "extends HttpServlet" in source_text: return "servlet" return "unknown" class JavaExtractor: def __init__(self): self.file_stem: str = "" self.file_path: str = "" self.framework: str = "unknown" self.source_bytes: bytes = b"" self.source_lines: list[str] = [] # ── Public entry ────────────────────────────────────────────────────────── 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.source_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, "java") 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_java_framework(source_text) self._extract_module_node(result) self._extract_classes(result, tree) self._extract_methods(result, tree) self._extract_entry_points(result, source_text) self._extract_sources(result, source_text) self._extract_sinks(result, source_text) self._extract_env_vars(result, source_text) self._extract_imports(result, source_text) self._extract_calls(result, tree) self._ensure_unique_ids(result) return result # ── Helpers ─────────────────────────────────────────────────────────────── 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 _find_containing_function(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 and n.line_end and n.line_start <= line_num <= n.line_end: return n.id return None def _get_containing_class(self, ts_node) -> str: parent = ts_node.parent while parent: if parent.type == "class_declaration": name_child = parent.child_by_field_name("name") if name_child: return get_node_text(name_child, self.source_bytes) parent = parent.parent return "" # ── Module node ─────────────────────────────────────────────────────────── def _extract_module_node(self, result: ExtractionResult): node_id = make_module_node_id(self.file_stem) result.nodes.append(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.source_lines), metadata={"language": "java", "framework": self.framework, "imports": []}, )) # ── Classes ─────────────────────────────────────────────────────────────── def _extract_classes(self, result: ExtractionResult, tree): query = """ (class_declaration name: (identifier) @class_name) @class_decl """ for cls_node, _ in query_tree(tree, "java", query, capture_filter="class_decl"): name_child = cls_node.child_by_field_name("name") if not name_child: continue cls_name = get_node_text(name_child, self.source_bytes) if not cls_name: continue node_id = make_node_id(self.file_stem, cls_name) result.nodes.append(Node( id=node_id, node_type=NodeType.CLASS, label=cls_name, file_path=self.file_path, line_start=cls_node.start_point[0] + 1, line_end=cls_node.end_point[0] + 1, metadata={"language": "java"}, )) result.edges.append(Edge( source_id=make_module_node_id(self.file_stem), target_id=node_id, edge_type=EdgeType.CONTAINS, )) # ── Methods ─────────────────────────────────────────────────────────────── def _extract_methods(self, result: ExtractionResult, tree): query = """ (method_declaration name: (identifier) @method_name) @method_decl """ for method_node, _ in query_tree(tree, "java", query, capture_filter="method_decl"): name_child = method_node.child_by_field_name("name") if not name_child: continue method_name = get_node_text(name_child, self.source_bytes) if not method_name: continue class_name = self._get_containing_class(method_node) qualified = f"{class_name}.{method_name}" if class_name else method_name node_id = make_node_id(self.file_stem, qualified) # Collect formal parameter names params: list[str] = [] params_node = method_node.child_by_field_name("parameters") if params_node: for child in params_node.children: if child.type == "formal_parameter": pname_node = child.child_by_field_name("name") if pname_node: pname = get_node_text(pname_node, self.source_bytes) if pname: params.append(pname) result.nodes.append(Node( id=node_id, node_type=NodeType.FUNCTION, label=method_name[:40], file_path=self.file_path, line_start=method_node.start_point[0] + 1, line_end=method_node.end_point[0] + 1, metadata={ "qualified_name": qualified, "params": params, "class_name": class_name, "language": "java", }, )) result.edges.append(Edge( source_id=make_module_node_id(self.file_stem), target_id=node_id, edge_type=EdgeType.CONTAINS, )) # ── Entry points (Spring MVC) ───────────────────────────────────────────── def _extract_entry_points(self, result: ExtractionResult, source_text: str): if self.framework != "spring": return # Match @GetMapping("/path"), @RequestMapping(value="/path"), etc. ann_re = re.compile( r"@(?P" + "|".join(_HTTP_ANNOTATIONS.keys()) + r")" r"(?:\s*\(\s*" r'(?:value\s*=\s*)?' r'"(?P[^"]*)"' r"[^)]*\))?", re.MULTILINE, ) for m in ann_re.finditer(source_text): ann_name = m.group("ann") http_method = _HTTP_ANNOTATIONS[ann_name] route_path = m.group("path") or "/" ann_line = source_text[: m.start()].count("\n") + 1 # Find the method node whose line_start is closest to (and >= ann_line). # In tree-sitter-java the method_declaration node starts at its first # modifier line, which is the annotation itself. best: Node | None = None best_dist = 9999 for n in result.nodes: if n.node_type != NodeType.FUNCTION: continue dist = n.line_start - ann_line if 0 <= dist <= 8 and dist < best_dist: best_dist = dist best = n if best is None: continue best.node_type = NodeType.ENTRY_POINT best.is_entry_point = True best.label = f"{http_method} {route_path}" best.metadata["http_method"] = http_method best.metadata["route_path"] = route_path best.metadata["framework"] = self.framework # ── Sources ─────────────────────────────────────────────────────────────── def _extract_sources(self, result: ExtractionResult, source_text: str): seen: set[str] = set() # Spring: @PathVariable / @RequestParam / @RequestBody on parameters # Pattern handles: @PathVariable Long id AND @RequestParam("name") String s param_re = re.compile( r"@(?P" + "|".join(_PARAM_INPUT_TYPES.keys()) + r")(?:\s*\(\s*(?:value\s*=\s*)?\"(?P[^\"]+)\"[^)]*\))?" r"\s+(?:(?:(?:@\w+\s+)*)" # optional further annotations r"[\w.<>\[\]]+\s+)+" # type (incl. generics / arrays) r"(?P\w+)", re.MULTILINE, ) for m in param_re.finditer(source_text): ann_name = m.group("ann") input_type = _PARAM_INPUT_TYPES[ann_name] param_name = m.group("alias") or m.group("pname") line_num = source_text[: m.start()].count("\n") + 1 self._add_source(result, seen, source_text, param_name, input_type, line_num) # HttpServletRequest-style (raw Servlet API) servlet_sources = [ (r'\.getParameter\s*\(\s*"(?P

[^"]+)"\s*\)', "query_string"), (r'\.getHeader\s*\(\s*"(?P

[^"]+)"\s*\)', "header"), (r'\.getQueryString\s*\(\s*\)', "query_string"), (r'\.getInputStream\s*\(\s*\)', "raw_body"), (r'\.getCookies\s*\(\s*\)', "cookie"), ] for pattern, input_type in servlet_sources: for m in re.finditer(pattern, source_text): gd = m.groupdict() param_name = gd.get("p") or f"*{input_type}*" line_num = source_text[: m.start()].count("\n") + 1 self._add_source(result, seen, source_text, param_name, input_type, line_num) def _add_source( self, result: ExtractionResult, seen: set[str], source_text: str, param_name: str, input_type: str, line_num: int, ): containing_func = self._find_containing_function(line_num, result) if not containing_func: return func_part = containing_func.split("::")[-1] source_id = make_node_id(self.file_stem, f"{func_part}::{param_name}") if source_id in seen: return seen.add(source_id) result.nodes.append(Node( id=source_id, node_type=NodeType.SOURCE, label=f"param: {param_name} [{input_type}]"[:40], file_path=self.file_path, line_start=line_num, line_end=line_num, metadata={ "param_name": param_name, "input_type": input_type, "framework": self.framework, }, )) result.edges.append(Edge( source_id=source_id, target_id=containing_func, edge_type=EdgeType.DATA_FLOW, )) # ── Sinks ───────────────────────────────────────────────────────────────── def _extract_sinks(self, result: ExtractionResult, source_text: str): seen: set[str] = set() flat: list[tuple[str, str, str]] = [] for sink_type, funcs in JAVA_SINK_PATTERNS.items(): for func_name, pattern in funcs.items(): flat.append((pattern, sink_type, func_name)) flat.sort(key=lambda x: _PRIORITY_ORDER.index(x[1]) if x[1] in _PRIORITY_ORDER else 99) for pattern, sink_type, sink_func in flat: for m in re.finditer(pattern, source_text): line_num = source_text[: m.start()].count("\n") + 1 containing_func = self._find_containing_function(line_num, result) if not containing_func: continue sink_id = make_node_id( self.file_stem, f"SINK::{sink_type}::{sink_func}::{line_num}" ) if sink_id in seen: continue seen.add(sink_id) result.nodes.append(Node( id=sink_id, node_type=NodeType.SINK, label=f"sink: {sink_func} [{sink_type}]", file_path=self.file_path, line_start=line_num, line_end=line_num, metadata={"sink_type": sink_type, "function_name": sink_func}, )) result.edges.append(Edge( source_id=containing_func, target_id=sink_id, edge_type=EdgeType.CALLS, )) # ── Env vars ────────────────────────────────────────────────────────────── def _extract_env_vars(self, result: ExtractionResult, source_text: str): env_re = re.compile(r'System\.getenv\s*\(\s*"(?P\w+)"\s*\)') for m in env_re.finditer(source_text): var_name = m.group("var") line_num = source_text[: m.start()].count("\n") + 1 env_id = make_node_id(self.file_stem, f"ENV::{var_name}") if any(n.id == env_id for n in result.nodes): continue result.nodes.append(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}, )) containing_func = self._find_containing_function(line_num, result) if containing_func: result.edges.append(Edge( source_id=env_id, target_id=containing_func, edge_type=EdgeType.DATA_FLOW, )) # ── Imports ─────────────────────────────────────────────────────────────── def _extract_imports(self, result: ExtractionResult, source_text: str): module_id = make_module_node_id(self.file_stem) import_re = re.compile( r"^import\s+(?:static\s+)?(?P[\w.]+)\s*;", re.MULTILINE ) imports_list = [m.group("imp") for m in import_re.finditer(source_text)] for n in result.nodes: if n.id == module_id: n.metadata["imports"] = imports_list break # ── Call graph ──────────────────────────────────────────────────────────── def _extract_calls(self, result: ExtractionResult, tree): query = """ (method_invocation name: (identifier) @method_name) @call """ func_map: dict[str, str] = {} for n in result.nodes: if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT): func_map[n.label] = n.id qn = n.metadata.get("qualified_name", "") if qn: func_map[qn] = n.id unresolved: list[dict] = [] for call_node, _ in query_tree(tree, "java", query, capture_filter="call"): name_child = call_node.child_by_field_name("name") if not name_child: continue callee = get_node_text(name_child, self.source_bytes) if not callee: continue line_num = call_node.start_point[0] + 1 caller_id = self._find_containing_function(line_num, result) if not caller_id: continue resolved = False for fname, fid in func_map.items(): if fname == callee or fid.endswith(f"::{callee}"): result.edges.append(Edge( source_id=caller_id, target_id=fid, edge_type=EdgeType.CALLS, )) resolved = True break if not resolved: unresolved.append({"caller_id": caller_id, "callee_name": callee, "line": line_num}) result.metadata["unresolved_calls"] = unresolved def extract(file_info: FileInfo) -> ExtractionResult: return JavaExtractor().extract(file_info)