q3alique-codeflow/codeflow/analysis/graph_builder.py

from __future__ import annotations
from pathlib import Path
from codeflow.discovery.file_scanner import FileInfo
from codeflow.analysis.base_extractor import ExtractionResult
from codeflow.analysis.python_extractor import PythonExtractor
from codeflow.analysis.js_extractor import JSExtractor
from codeflow.analysis.java_extractor import JavaExtractor
from codeflow.analysis.go_extractor import GoExtractor
from codeflow.analysis.csharp_extractor import CSharpExtractor
from codeflow.analysis.generic_extractor import GenericExtractor
from codeflow.models.node import Node, NodeType, make_node_id, make_module_node_id
from codeflow.models.edge import Edge, EdgeType
from codeflow.models.graph import CodeGraph
# Only create BOUNDARY nodes for calls that may become dangerous sinks.
# This prevents builtins (print, len, str, append …) from flooding the graph.
from codeflow.analysis.boundary_marker import (
    DANGEROUS_CALLS as _DANGEROUS_CALL_NAMES,
    SANITIZER_CALLS as _SANITIZER_CALLEES,
)

_TRACKED_CALLEES: frozenset[str] = frozenset(_DANGEROUS_CALL_NAMES.keys())


def _select_extractor(language: str):
    if language in ("python",):
        return PythonExtractor()
    if language in ("javascript", "typescript"):
        return JSExtractor()
    if language == "java":
        return JavaExtractor()
    if language == "go":
        return GoExtractor()
    if language == "c_sharp":
        return CSharpExtractor()
    return GenericExtractor()


def _extract_file(file_info: FileInfo) -> ExtractionResult:
    extractor = _select_extractor(file_info.language)
    return extractor.extract(file_info)


def build(file_infos: list[FileInfo]) -> CodeGraph:
    graph = CodeGraph()
    all_results: list[ExtractionResult] = []

    for fi in file_infos:
        result = _extract_file(fi)
        all_results.append(result)

    for result in all_results:
        for node in result.nodes:
            if not graph.has_node(node.id):
                graph.add_node(node)

    for result in all_results:
        for edge in result.edges:
            if graph.has_node(edge.source_id) and graph.has_node(edge.target_id):
                graph.add_edge(edge)

    # Build a lookup: function label → node id (for cross-file call resolution)
    all_callee_names: dict[str, str] = {}
    for result in all_results:
        for node in result.nodes:
            if node.node_type in (NodeType.FUNCTION, NodeType.ENTRY_POINT):
                all_callee_names[node.label] = node.id

    # Single pass: resolve unresolved calls OR create BOUNDARY nodes only for
    # callees that are in _TRACKED_CALLEES (i.e. potentially dangerous sinks).
    for result in all_results:
        unresolved = result.metadata.get("unresolved_calls", [])
        for call_info in unresolved:
            caller_id = call_info["caller_id"]
            callee_name = call_info["callee_name"]

            # Try cross-file resolution first
            resolved = False
            for fname, fid in all_callee_names.items():
                if fname == callee_name or fid.endswith(f"::{callee_name}"):
                    if graph.has_node(caller_id) and graph.has_node(fid):
                        graph.add_edge(Edge(
                            source_id=caller_id,
                            target_id=fid,
                            edge_type=EdgeType.CALLS,
                        ))
                    resolved = True
                    break

            if resolved:
                continue

            # Sanitizer call: create a node with is_sanitizer=True.
            # Taint propagation will stop at these nodes (see taint_tracker.py).
            if callee_name in _SANITIZER_CALLEES:
                san_id = make_node_id(
                    Path(result.file_path).stem,
                    f"SANITIZER::{callee_name}",
                )
                if not graph.has_node(san_id):
                    graph.add_node(Node(
                        id=san_id,
                        node_type=NodeType.BOUNDARY,
                        label=f"✓ {callee_name}",
                        file_path=result.file_path,
                        metadata={"callee_name": callee_name, "is_sanitizer": True},
                    ))
                if graph.has_node(caller_id):
                    graph.add_edge(Edge(
                        source_id=caller_id,
                        target_id=san_id,
                        edge_type=EdgeType.CALLS,
                    ))
                continue  # don't also create a dangerous-sink boundary

            # Only track external calls that could become dangerous sinks.
            # Skip stdlib noise (print, len, str, format, append, …).
            if callee_name not in _TRACKED_CALLEES:
                continue

            boundary_id = make_node_id(
                Path(result.file_path).stem,
                f"BOUNDARY::{callee_name}",
            )
            if not graph.has_node(boundary_id):
                graph.add_node(Node(
                    id=boundary_id,
                    node_type=NodeType.BOUNDARY,
                    label=f"boundary: {callee_name}",
                    file_path=result.file_path,
                    metadata={"callee_name": callee_name},
                ))
            if graph.has_node(caller_id):
                graph.add_edge(Edge(
                    source_id=caller_id,
                    target_id=boundary_id,
                    edge_type=EdgeType.CALLS,
                ))

    # Wire up IMPORTS edges between modules in the same repo
    for result in all_results:
        for i, node_i in enumerate(result.nodes):
            if node_i.node_type != NodeType.MODULE:
                continue
            for j, node_j in enumerate(result.nodes):
                if i >= j or node_j.node_type != NodeType.MODULE:
                    continue
                imports_i = set(node_i.metadata.get("imports", []))
                stem_j = Path(node_j.file_path).stem

                for imp in imports_i:
                    if stem_j.lower() in imp.lower():
                        graph.add_edge(Edge(
                            source_id=node_i.id,
                            target_id=node_j.id,
                            edge_type=EdgeType.IMPORTS,
                        ))

    return graph