q3alique-codeflow/codeflow/analysis/js_extractor.py

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


JS_SINK_PATTERNS = {
    "database": {
        "query": {"db.query", "connection.query", "pool.query", "client.query"},
        "execute": {".execute"},
    },
    "command_execution": {
        "execSync": {"execSync", "exec"},
        "spawn": {"spawn"},
        "fork": {"fork"},
    },
    "file_write": {
        "writeFile": {"fs.writeFile", "fs.writeFileSync"},
        "appendFile": {"fs.appendFile", "fs.appendFileSync"},
    },
    "file_read": {
        "readFile": {"fs.readFile", "fs.readFileSync"},
    },
    "http_request": {
        "fetch": {"fetch"},
        "get": {"axios.get", "request.get"},
        "post": {"axios.post", "request.post"},
    },
    "code_execution": {
        "eval": {"eval"},
        "Function": {"Function"},
    },
    "deserialization": {
        "parse": {"JSON.parse"},
    },
}


def _detect_js_framework(source_text: str) -> str:
    text_lower = source_text.lower()
    if "require('express')" in text_lower or 'require("express")' in text_lower:
        return "express"
    if "from 'express'" in text_lower or 'from "express"' in text_lower:
        return "express"
    if "require('fastify')" in text_lower or 'require("fastify")' in text_lower:
        return "fastify"
    return "unknown"


class JSExtractor:
    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_js_framework(source_text)

        self._extract_module_node(result)
        self._extract_functions(result, tree)
        self._extract_classes(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._extract_calls_and_imports(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 _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 <= line_num <= n.line_end:
                    return n.id
        return None

    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": self.framework or "javascript", "imports": []},
        )
        result.nodes.append(node)

    def _extract_functions(self, result: ExtractionResult, tree):
        query = """
        (function_declaration
          name: (identifier) @func_name) @func_decl

        (variable_declarator
          name: (identifier) @func_name
          value: [(arrow_function) (function)]) @var_decl
        """
        captures = query_tree(tree, self._get_lang(), query, capture_filter={"func_decl", "var_decl"})
        for func_node, _ in captures:
            func_name = None
            name_child = func_node.child_by_field_name("name")
            if name_child:
                func_name = get_node_text(name_child, self.source_bytes)
            if not func_name:
                continue
            node_id = make_node_id(self.file_stem, func_name)
            try:
                start = func_node.start_point[0] + 1
                end = func_node.end_point[0] + 1
            except Exception:
                start, end = 0, 0
            node = Node(
                id=node_id,
                node_type=NodeType.FUNCTION,
                label=func_name[:40],
                file_path=self.file_path,
                line_start=start,
                line_end=end,
            )
            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_declaration
          name: (identifier) @class_name) @class_decl
        """
        captures = query_tree(tree, self._get_lang(), query, capture_filter="class_decl")
        for cls_node, _ in captures:
            cls_name = None
            for child in cls_node.children:
                if child.type == "name" and child.children:
                    cls_name = get_node_text(child.children[0], self.source_bytes)
                    break
            if not cls_name:
                continue
            node_id = make_node_id(self.file_stem, cls_name)
            try:
                start = cls_node.start_point[0] + 1
                end = cls_node.end_point[0] + 1
            except Exception:
                start, end = 0, 0
            node = Node(
                id=node_id,
                node_type=NodeType.CLASS,
                label=cls_name,
                file_path=self.file_path,
                line_start=start,
                line_end=end,
            )
            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 _get_lang(self) -> str:
        return "javascript"

    def _extract_entry_points(self, result: ExtractionResult, tree):
        if self.framework != "express":
            return

        query = """
        (call_expression
          function: (member_expression
            object: (identifier) @app_obj
            property: (property_identifier) @http_method)
          arguments: (arguments
            (string) @route_path
            [(identifier) (arrow_function) (function)] @handler)) @route_call
        """
        captures = query_tree(tree, self._get_lang(), query, capture_filter="route_call")
        route_methods = {"get", "post", "put", "delete", "patch", "options", "use"}

        for call_node, _ in captures:
            func_child = call_node.child_by_field_name("function")
            args_child = call_node.child_by_field_name("arguments")
            if not func_child or not args_child:
                continue

            http_method = None
            if func_child.type == "member_expression":
                prop = func_child.child_by_field_name("property")
                if prop:
                    http_method = get_node_text(prop, self.source_bytes).lower()

            if http_method not in route_methods:
                continue

            args = [c for c in args_child.children if c.type not in ("(", ")", ",")]
            route_path = None
            handler = None
            for arg in args:
                if arg.type == "string" and route_path is None:
                    route_path = get_node_text(arg, self.source_bytes).strip("'\"")
                elif arg.type in ("identifier", "arrow_function", "function_expression"):
                    handler = arg

            if not route_path:
                continue

            http_method_upper = http_method.upper()

            if handler and handler.type == "identifier":
                handler_name = get_node_text(handler, self.source_bytes)
                func_id = make_node_id(self.file_stem, handler_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_upper} {route_path}"
                        n.metadata["http_method"] = http_method_upper
                        n.metadata["route_path"] = route_path
                        n.metadata["framework"] = "express"
                        break

            entry_id = make_node_id(self.file_stem, f"route_{http_method}_{route_path.replace('/', '_')}")
            if not any(n.id == entry_id for n in result.nodes):
                entry_node = Node(
                    id=entry_id,
                    node_type=NodeType.ENTRY_POINT,
                    label=f"{http_method_upper} {route_path}",
                    file_path=self.file_path,
                    line_start=call_node.start_point[0] + 1,
                    line_end=call_node.end_point[0] + 1,
                    is_entry_point=True,
                    metadata={
                        "http_method": http_method_upper,
                        "route_path": route_path,
                        "framework": "express",
                    },
                )
                result.nodes.append(entry_node)
                result.edges.append(Edge(
                    source_id=make_module_node_id(self.file_stem),
                    target_id=entry_id,
                    edge_type=EdgeType.CONTAINS,
                ))

    def _extract_sources(self, result: ExtractionResult, tree):
        source_text = self.source_bytes.decode("utf-8", errors="replace")

        req_patterns = [
            (r"req\.body\.(\w+)", "json_body"),
            (r"req\.params\.(\w+)", "url_param"),
            (r"req\.query\.(\w+)", "query_string"),
            (r"req\.headers\.(\w+)", "header"),
            (r"req\.cookies\.(\w+)", "cookie"),
            (r"req\.files\.(\w+)", "file_upload"),
            (r"req\.body\b", "json_body"),
            (r"req\.params\b", "url_param"),
            (r"req\.query\b", "query_string"),
        ]

        seen_source_ids: set[str] = set()

        for pattern, input_type in req_patterns:
            for match in re.finditer(pattern, source_text):
                param_name = match.lastgroup if match.lastgroup else match.group(1) if match.lastindex and match.lastindex >= 1 else f"*{input_type}*"
                if match.lastindex and match.lastindex >= 1:
                    param_name = match.group(1)
                else:
                    param_name = f"*{input_type}*"
                line_num = source_text[:match.start()].count("\n") + 1
                containing_func = self._find_containing_function(line_num, result)
                if not containing_func:
                    continue
                func_part = containing_func.split("::")[-1]
                source_id = make_node_id(self.file_stem, f"{func_part}::{param_name}")
                if source_id in seen_source_ids:
                    continue
                seen_source_ids.add(source_id)
                src = 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": "express"},
                )
                result.nodes.append(src)
                result.edges.append(Edge(
                    source_id=source_id,
                    target_id=containing_func,
                    edge_type=EdgeType.DATA_FLOW,
                ))

        if "process.argv" in source_text:
            for match in re.finditer(r"process\.argv", source_text):
                line_num = source_text[:match.start()].count("\n") + 1
                containing_func = self._find_containing_function(line_num, result)
                if not containing_func:
                    continue
                func_part = containing_func.split("::")[-1]
                source_id = make_node_id(self.file_stem, f"{func_part}::argv")
                if source_id not in seen_source_ids:
                    seen_source_ids.add(source_id)
                    result.nodes.append(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": "process.argv"},
                    ))
                    result.edges.append(Edge(
                        source_id=source_id,
                        target_id=containing_func,
                        edge_type=EdgeType.DATA_FLOW,
                    ))

    def _extract_sinks(self, result: ExtractionResult, tree):
        source_text = self.source_bytes.decode("utf-8", errors="replace")

        sink_by_name = {}
        for sink_type, funcs in JS_SINK_PATTERNS.items():
            for func_name, patterns in funcs.items():
                for pat in patterns:
                    sink_by_name[pat.split(".")[-1]] = (sink_type, func_name)

        seen_sink_ids: set[str] = set()
        for i, line in enumerate(source_text.splitlines()):
            line_num = i + 1
            for method_name, (sink_type, sink_func) in sink_by_name.items():
                if method_name in line:
                    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_sink_ids:
                        continue
                    seen_sink_ids.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,
                    ))

    def _extract_env_vars(self, result: ExtractionResult, tree):
        source_text = self.source_bytes.decode("utf-8", errors="replace")
        env_pattern = re.compile(r"process\.env\.(\w+)")
        for match in env_pattern.finditer(source_text):
            var_name = match.group(1)
            line_num = source_text[:match.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,
                ))

    def _extract_calls_and_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

        source_text = self.source_bytes.decode("utf-8", errors="replace")
        imports_list = []

        require_pattern = re.compile(r"require\(['\"](\S+?)['\"]\)")
        for match in require_pattern.finditer(source_text):
            module_name = match.group(1)
            imports_list.append(module_name)

        es6_import_pattern = re.compile(r"(?:import\s+.*\s+from\s+['\"]|import\s+['\"'])(\S+?)['\"]")
        for match in es6_import_pattern.finditer(source_text):
            module_name = match.group(1)
            imports_list.append(module_name)

        if module_node:
            module_node.metadata["imports"] = imports_list

        func_node_map = {}
        for n in result.nodes:
            if n.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
                func_node_map[n.label] = n.id

        unresolved: list[dict] = []
        call_pattern = re.compile(r"(\w+)\.(\w+)\s*\(")
        for match in call_pattern.finditer(source_text):
            callee = match.group(2)
            line_num = source_text[:match.start()].count("\n") + 1
            containing_func = self._find_containing_function(line_num, result)
            if not containing_func:
                continue
            found = False
            for fname, fid in func_node_map.items():
                if fname == callee:
                    result.edges.append(Edge(
                        source_id=containing_func,
                        target_id=fid,
                        edge_type=EdgeType.CALLS,
                    ))
                    found = True
                    break
            if not found:
                unresolved.append({
                    "caller_id": containing_func,
                    "callee_name": callee,
                    "line": line_num,
                })

        result.metadata["unresolved_calls"] = unresolved


def extract(file_info: FileInfo) -> ExtractionResult:
    extractor = JSExtractor()
    return extractor.extract(file_info)