trailofbits-buttercup/program-model/src/buttercup/program_model/codequery.py

"""Codequery based code querying module"""

from __future__ import annotations

import logging
import shutil
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from itertools import groupby
from typing import ClassVar, Optional
import rapidfuzz
import uuid
from buttercup.common.challenge_task import ChallengeTask, ChallengeTaskError
from buttercup.program_model.api.tree_sitter import CodeTS
from buttercup.program_model.api.fuzzy_imports_resolver import (
    FuzzyJavaImportsResolver,
    FuzzyCImportsResolver,
)
from buttercup.program_model.utils.common import (
    Function,
    TypeDefinition,
    TypeUsageInfo,
)
from buttercup.common.project_yaml import ProjectYaml

logger = logging.getLogger(__name__)


CONTAINER_SRC_DIR: str = "container_src_dir"


@dataclass
class CQSearchResult:
    """Result of the cqsearch command."""

    value: str
    file: Path
    line: int
    body: str

    @classmethod
    def from_line(cls, line: str) -> CQSearchResult | None:
        """Parse a line of the cqsearch output into a CQSearchResult."""
        try:
            value, file_line, body = line.split("\t", 2)
            file_str, line = file_line.split(":", 1)
        except ValueError:
            logger.warning("Invalid cqsearch line: %s", line)
            return None

        # Rebase the file path from the challenge task base dir.
        # This is needed because the task-dir part might be different from what
        # was originally used to create the db.
        file: Path = Path(file_str)
        if CONTAINER_SRC_DIR not in file.parts:
            logger.warning("File %s is not in the container source dir", file_str)
            return None

        container_src_dir_idx = file.parts.index(CONTAINER_SRC_DIR)
        assert container_src_dir_idx > 0
        file = Path(*file.parts[container_src_dir_idx:])

        try:
            line_number = int(line)
        except ValueError:
            logger.warning("Invalid line number: %s", line)
            line_number = 0

        return cls(value, file, line_number, body)


@dataclass
class CodeQuery:
    """Class to extract context about a challenge project with CodeQuery.

    This class indexes the codebase as it appears from within the oss-fuzz
    container used for fuzzing. All returned paths are absolute paths in the
    container (e.g. /src/my-source/my-file.c).
    """

    challenge: ChallengeTask
    ts: CodeTS = field(init=False)
    imports_resolver: Optional[FuzzyCImportsResolver | FuzzyJavaImportsResolver] = (
        field(init=False)
    )

    CSCOPE_FILES: ClassVar[str] = "cscope.files"
    CSCOPE_OUT: ClassVar[str] = "cscope.out"
    TAGS: ClassVar[str] = "tags"
    CODEQUERY_DB: ClassVar[str] = "codequery.db"

    def __post_init__(self) -> None:
        """Initialize the CodeQuery object."""
        self._verify_requirements()
        self.ts = CodeTS(self.challenge)
        language = self._get_project_language()
        if language in ["c", "c++"]:
            self.imports_resolver = FuzzyCImportsResolver(self._get_container_src_dir())
        elif language in ["java", "jvm"]:
            self.imports_resolver = FuzzyJavaImportsResolver(self.challenge)
        else:
            self.imports_resolver = None

        if self._is_already_indexed():
            logger.info("CodeQuery DB already exists in %s.", self.challenge.task_dir)
            return

        if self.challenge.local_task_dir is None:
            raise ValueError(
                "Challenge Task is read-only, cannot perform this operation"
            )

        self._create_codequery_db()
        logger.info("CodeQuery DB created successfully.")

    def _verify_requirements(self) -> None:
        """Verify that the required commands are installed."""
        required_commands = ["cscope", "ctags", "cqmakedb", "cqsearch"]
        missing_commands = []

        for command in required_commands:
            if shutil.which(command) is None:
                missing_commands.append(command)

        if missing_commands:
            logger.fatal(
                "Missing commands: %s. Please install the 'codequery' package.",
                ", ".join(missing_commands),
            )
            raise RuntimeError("No code query package")

    def _get_project_language(self) -> str:
        project_yaml = ProjectYaml(
            self.challenge, self.challenge.task_meta.project_name
        )
        return str(project_yaml.language)

    def _is_already_indexed(self) -> bool:
        """Check if the codequery database already exists."""
        return (
            self.challenge.task_dir.joinpath(CONTAINER_SRC_DIR).exists()
            and self._get_container_src_dir().joinpath(self.CSCOPE_FILES).exists()
            and self._get_container_src_dir().joinpath(self.CSCOPE_OUT).exists()
            and self._get_container_src_dir().joinpath(self.CODEQUERY_DB).exists()
            and self._get_container_src_dir().joinpath(self.TAGS).exists()
        )

    def _get_container_src_dir(self) -> Path:
        """Get the container source directory."""
        return Path(self.challenge.task_dir.joinpath(CONTAINER_SRC_DIR))

    def _generated_files(self) -> bool:
        """Check if there are files that will be generated by the build process."""
        src_dir = self.challenge.get_source_path()

        language = self._get_project_language()
        if language in ["c", "c++"]:
            return False
        elif language in ["java", "jvm"]:
            return any(src_dir.glob("**/*.g4"))
        else:
            raise ValueError(f"Unsupported language: {language}")

    def _copy_src_from_container(self) -> None:
        """Build and copy the /src directory from the container to the challenge task directory."""

        # If there are files that will be generated by the build process, we need to
        # copy the /src directory from the container to the challenge task directory
        # AFTER the build process.
        if self._generated_files():
            logger.info("There are files that will be generated by the build process.")

            self.challenge.clean = False
            name = self.challenge.task_meta.task_id + "_" + str(uuid.uuid4())[:16]
            res = self.challenge.build_fuzzers_save_containers(
                container_name=name,
            )
            if not res.success:
                raise RuntimeError("Failed to build image.")

            src_dst = Path(self.challenge.task_dir.joinpath("container_src_dir"))
            src_dst.mkdir(parents=True, exist_ok=True)

            try:
                command = [
                    "docker",
                    "cp",
                    f"{name}:/src",
                    src_dst.resolve().as_posix(),
                ]
                logger.info("Running command: %s", " ".join(command))
                subprocess.run(command, check=True, capture_output=True)
            except subprocess.CalledProcessError as e:
                logger.error("Failed to copy src from container: %s", e)
                raise RuntimeError(f"Failed to copy src from container: {e}")
            finally:
                command = [
                    "docker",
                    "rm",
                    f"{name}",
                ]
                logger.info("Running command: %s", " ".join(command))
                subprocess.run(command, check=True, capture_output=True)

            self.challenge.clean = True

        # Else, can just copy from the built image
        else:
            res = self.challenge.build_image(pull_latest_base_image=True)
            if not res.success:
                raise RuntimeError("Failed to build image.")

            challenge_container_name = self.challenge.container_image()
            src_dst = self._get_container_src_dir()
            src_dst.mkdir(parents=True, exist_ok=True)
            try:
                command = [
                    "docker",
                    "create",
                    "--name",
                    f"codequery-container-{self.challenge.task_meta.task_id}",
                    "-v",
                    f"{self.challenge.get_source_path().as_posix()}:{self.challenge.workdir_from_dockerfile()}",
                    challenge_container_name,
                ]
                subprocess.run(command, check=True, capture_output=True)
                command = [
                    "docker",
                    "cp",
                    f"codequery-container-{self.challenge.task_meta.task_id}:/src",
                    src_dst.resolve().as_posix(),
                ]
                subprocess.run(command, check=True, capture_output=True)
            except subprocess.CalledProcessError as e:
                logger.error("Failed to copy src from container: %s", e)
                raise RuntimeError(f"Failed to copy src from container: {e}")
            finally:
                command = [
                    "docker",
                    "rm",
                    f"codequery-container-{self.challenge.task_meta.task_id}",
                ]
                subprocess.run(command, check=True, capture_output=True)

    def _create_codequery_db(self) -> None:
        """Create the codequery database."""
        self._copy_src_from_container()

        with self._get_container_src_dir().joinpath(self.CSCOPE_FILES).open("w") as f:
            project_yaml = ProjectYaml(
                self.challenge, self.challenge.task_meta.project_name
            )
            if project_yaml.language == "c" or project_yaml.language == "c++":
                extensions = [
                    "*.c",
                    "*.cpp",
                    "*.cxx",
                    "*.cc",
                    "*.h",
                    "*.hpp",
                    "*.hxx",
                    "*.hh",
                ]
            elif project_yaml.language == "jvm":
                extensions = ["*.java"]
            else:
                raise ValueError(f"Unsupported language: {project_yaml.language}")

            # Find all files with the given extensions
            for ext in extensions:
                for file in self._get_container_src_dir().rglob(ext):
                    f.write(str(file) + "\n")

        try:
            subprocess.run(
                ["cscope", "-bkq"],
                cwd=self._get_container_src_dir(),
                capture_output=True,
                timeout=200,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
            raise RuntimeError("Failed to create cscope index.")

        if not self._get_container_src_dir().joinpath(self.CSCOPE_OUT).exists():
            raise RuntimeError("Failed to create cscope out.")

        try:
            subprocess.run(
                ["ctags", "--fields=+i", "-n", "-L", self.CSCOPE_FILES],
                cwd=self._get_container_src_dir(),
                capture_output=True,
                timeout=300,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
            raise RuntimeError("Failed to create ctags index.")

        if not self._get_container_src_dir().joinpath(self.TAGS).exists():
            raise RuntimeError("Failed to create ctags index.")

        try:
            subprocess.run(
                [
                    "cqmakedb",
                    "-s",
                    self.CODEQUERY_DB,
                    "-c",
                    self.CSCOPE_OUT,
                    "-t",
                    self.TAGS,
                    "-p",
                ],
                cwd=self._get_container_src_dir(),
                capture_output=True,
                timeout=2700,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
            raise RuntimeError("Failed to create cquery database.")

        if not self._get_container_src_dir().joinpath(self.CODEQUERY_DB).exists():
            raise RuntimeError("Failed to create cquery database.")

    def __repr__(self) -> str:
        return f"CodeQuery(challenge={self.challenge})"

    def _run_cqsearch(self, *args: str) -> list[CQSearchResult]:
        """Run the cqsearch command and parse the results."""
        try:
            result = subprocess.run(
                ["cqsearch", *args],
                cwd=self._get_container_src_dir(),
                capture_output=True,
                text=True,
                check=True,
            )
            output = result.stdout
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Failed to run cqsearch: {e}")

        results = [CQSearchResult.from_line(line) for line in output.splitlines()]
        return [result for result in results if result is not None]

    def _rebase_path(self, path: Path) -> Path:
        if CONTAINER_SRC_DIR not in path.parts:
            return path
        container_src_dir_idx = path.parts.index(CONTAINER_SRC_DIR)
        return Path("/", *path.parts[container_src_dir_idx + 1 :])

    def _rebase_functions_file_paths(self, functions: list[Function]) -> list[Function]:
        """Rebase the file paths of the functions to the challenge task container structure."""
        return [
            Function(
                name=function.name,
                file_path=self._rebase_path(function.file_path),
                bodies=function.bodies,
            )
            for function in functions
        ]

    def _rebase_types_file_paths(
        self, types: list[TypeDefinition]
    ) -> list[TypeDefinition]:
        """Rebase the file paths of the types to the challenge task container structure."""
        return [
            TypeDefinition(
                name=td.name,
                file_path=self._rebase_path(td.file_path),
                definition=td.definition,
                definition_line=td.definition_line,
                type=td.type,
            )
            for td in types
        ]

    def _rebase_type_usages_file_paths(
        self, type_usages: list[TypeUsageInfo]
    ) -> list[TypeUsageInfo]:
        """Rebase the file paths of the types to the challenge task container structure."""
        return [
            TypeUsageInfo(
                name=tu.name,
                file_path=self._rebase_path(tu.file_path),
                line_number=tu.line_number,
            )
            for tu in type_usages
        ]

    def _get_all_functions(self) -> list[CQSearchResult]:
        """Get all functions in the codebase."""
        return [
            f
            for f in self._run_cqsearch(
                "-s", self.CODEQUERY_DB, "-p", "2", "-t", "*", "-u"
            )
        ]

    def _get_all_types(self) -> list[CQSearchResult]:
        """Get all symbols in the codebase."""
        return [
            t
            for t in self._run_cqsearch(
                "-s", self.CODEQUERY_DB, "-p", "1", "-t", "*", "-u"
            )
        ]

    def get_functions(
        self,
        function_name: str,
        file_path: Path | None = None,
        line_number: int | None = None,
        fuzzy: bool | None = False,
        fuzzy_threshold: int = 80,
    ) -> list[Function]:
        """Get the definition(s) of a function in the codebase or in a specific
        file. File paths are based on the challenge task container structure
        (e.g. /src).

        The order of the results is (1) exact matches and (2) fuzzy matches sorted in descending order of similarity.

        NOTE: Fuzzy search will be disabled if a file path is provided.
        """
        if fuzzy and file_path:
            logger.warning(
                "Fuzzy search will be disabled because file path %s was provided.",
                file_path,
            )

        # FIXME(Evan): Sometimes cscope doesn't identify a function. They can be found by looking for symbols.
        results: list[CQSearchResult] = []
        for search_type in ["1", "2"]:  # 1 for symbols, 2 for functions
            cqsearch_args = [
                "-s",
                self.CODEQUERY_DB,
                "-p",
                search_type,
                "-t",
                function_name,
                "-e",
                "-u",
            ]
            if file_path:
                cqsearch_args += ["-b", file_path.as_posix()]

            results.extend(self._run_cqsearch(*cqsearch_args))

        # Extended fuzzy matching
        if fuzzy and file_path is None:
            # Fuzzy match the function name against all functions in the codebase
            fuzzy_matches: list[tuple[CQSearchResult, float]] = sorted(
                [
                    (f, rapidfuzz.fuzz.ratio(function_name, f.value))
                    for f in self._get_all_functions()
                    if f.value
                    and rapidfuzz.fuzz.ratio(function_name, f.value) > fuzzy_threshold
                ],
                key=lambda x: x[1],
                reverse=True,
            )
            fuzzy_matches = [f for f, _ in fuzzy_matches]
            logger.info(
                "Found %d fuzzy matches for function %s",
                len(fuzzy_matches),
                function_name,
            )
            results.extend(fuzzy_matches)
        logger.info(
            "Found %d instances (fuzzy or exact) of function %s",
            len(results),
            function_name,
        )

        res: set[Function] = set()
        results_by_file = groupby(results, key=lambda x: x.file)
        for file, file_results in results_by_file:
            file_results_list = list(file_results)
            functions_found = list(set(result.value for result in file_results_list))
            if not fuzzy and not all(function_name == f for f in functions_found):
                logger.warning(
                    "Function name mismatch, this should not happen: %s",
                    function_name,
                )
                continue

            for function in functions_found:
                f = self.ts.get_function(function, file)
                logger.info("Function: %s | File: %s", function, file)
                if f is None:
                    logger.warning(
                        "Function not found in tree-sitter: %s/%s", file, function
                    )
                    continue
                if line_number:
                    lines = [
                        (
                            body.start_line,
                            body.end_line,
                        )
                        for body in f.bodies
                    ]
                    # NOTE(boyan): We check whether the supplied line to look up for the function
                    # is contained within at least one of the function bodies found by
                    # tree-sitter
                    if any(
                        True
                        for start_line, end_line in lines
                        if start_line <= line_number <= end_line
                    ):
                        res.add(f)
                    else:
                        logger.warning(
                            "Function (%s) not found using tree-sitter in file %s for line %d",
                            function,
                            file,
                            line_number,
                        )
                else:
                    res.add(f)

        logger.info(
            "Found %d functions for %s after filtering", len(res), function_name
        )

        # Sort in same order as results
        results_value = [r.value for r in results]
        res_sorted: list[Function] = sorted(
            res, key=lambda x: results_value.index(x.name)
        )

        return self._rebase_functions_file_paths(res_sorted)

    def _filter_callees(
        self, caller_function: Function, callees: list[Function]
    ) -> list[Function]:
        # If no resolver available, don't filter anything
        if not self.imports_resolver:
            return callees
        return self.imports_resolver.filter_callees(caller_function, callees)

    def get_callers(self, function: Function | str) -> list[Function]:
        """Get the callers of a function. File paths are based on the challenge
        task container structure (e.g. /src)."""
        if isinstance(function, str):
            function_name = function
        elif isinstance(function, Function):
            function_name = function.name
        cqsearch_args = [
            "-s",
            self.CODEQUERY_DB,
            "-p",
            "6",
            "-t",
            function_name,
            "-e",
            "-u",
        ]

        results = self._run_cqsearch(*cqsearch_args)

        callers: set[Function] = set()
        for result in results:
            functions = self.get_functions(result.value, Path(result.file), result.line)
            callers.update(functions)

        logger.info("Found %d callers for %s", len(callers), function_name)

        return self._rebase_functions_file_paths(list(callers))

    def get_callees(self, function: Function | str) -> list[Function]:
        """Get the callees of a function. File paths are based on the challenge
        task container structure (e.g. /src)."""
        if isinstance(function, str):
            function_name = function
        elif isinstance(function, Function):
            function_name = function.name
        cqsearch_args = [
            "-s",
            self.CODEQUERY_DB,
            "-p",
            "7",
            "-t",
            function_name,
            "-e",
            "-u",
        ]

        results = self._run_cqsearch(*cqsearch_args)

        callees: set[Function] = set()
        for result in results:
            functions = self.get_functions(result.value)
            unique_functions: list[Function] = []
            for f in functions:
                if not any(x for x in unique_functions if x.has_same_source(f)):
                    unique_functions.append(f)
            callees.update(functions)

        # TODO(boyan): if function is a str we should try to find the actual function
        # at the beginning of this function so we can use to do the filtering. If not
        # then we need the function file
        if isinstance(function, Function):
            callees = self._filter_callees(function, list(callees))  # type: ignore [assignment]
        logger.info("Found %d callees for %s", len(callees), function_name)
        res = self._rebase_functions_file_paths(list(callees))
        return res

    def get_types(
        self,
        type_name: str,
        file_path: Path | None = None,
        function_name: str | None = None,
        fuzzy: bool | None = False,
        fuzzy_threshold: int = 80,
    ) -> list[TypeDefinition]:
        """Finds and return the definition of type named `typename`. File paths
        are based on the challenge task container structure (e.g. /src).

        The order of the results is (1) exact matches and (2) fuzzy matches sorted in descending order of similarity.

        NOTE: Fuzzy search will be disabled if a file path is provided.
        """
        if fuzzy and file_path:
            logger.warning(
                "Fuzzy search will be disabled because file path %s was provided.",
                file_path,
            )

        cqsearch_args = [
            "-s",
            self.CODEQUERY_DB,
            "-p",
            "1",
            "-t",
            str(type_name),
            "-e",
            "-u",
        ]
        if file_path:
            cqsearch_args += ["-b", file_path.as_posix()]

        results: list[CQSearchResult] = list(self._run_cqsearch(*cqsearch_args))

        # Extended fuzzy matching
        if fuzzy and file_path is None:
            # Fuzzy match the function name against all functions in the codebase
            fuzzy_matches: list[tuple[CQSearchResult, float]] = sorted(
                [
                    (t, rapidfuzz.fuzz.ratio(type_name, t.value))
                    for t in self._get_all_types()
                    if t.value
                    and rapidfuzz.fuzz.ratio(type_name, t.value) > fuzzy_threshold
                ],
                key=lambda x: x[1],
                reverse=True,
            )
            fuzzy_matches = [t for t, _ in fuzzy_matches]
            logger.info(
                "Found %d fuzzy matches for type %s", len(fuzzy_matches), type_name
            )
            results.extend(fuzzy_matches)
        logger.info(
            "Found %d instances (fuzzy or exact) of type %s", len(results), type_name
        )

        logger.info("Found %d instances of type %s", len(results), type_name)

        res: list[TypeDefinition] = []
        results_by_file = groupby(results, key=lambda x: x.file)
        for file, file_results in results_by_file:
            file_results_list = list(file_results)
            types_found = list(set(result.value for result in file_results_list))

            if not fuzzy and not all(str(type_name) == str(t) for t in types_found):
                logger.warning(
                    "Type name mismatch, this should not happen: %s",
                    type_name,
                )
                continue

            typedefs: dict[str, TypeDefinition] = {}

            for typename in types_found:
                t = self.ts.parse_types_in_code(file, typename, fuzzy)
                if not t:
                    logger.warning(
                        "Type definition not found in tree-sitter: %s", typename
                    )
                    continue
                typedefs.update(t)

            if function_name:
                # Get the function definition to find its scope
                function = self.ts.get_function(function_name, file)
                if function:
                    # Filter type definitions to only include those within the function's scope
                    filtered_typedefs = {}
                    for name, typedef in typedefs.items():
                        # Check if the type definition is within the function's scope
                        for body in function.bodies:
                            if (
                                body.start_line
                                <= typedef.definition_line
                                <= body.end_line
                            ):
                                filtered_typedefs[name] = typedef
                                break
                    typedefs = filtered_typedefs
                else:
                    typedefs = {}

            res.extend(typedefs.values())

        logger.info("Found %d types for %s after filtering", len(res), type_name)

        # Sort in same order as results
        results_value = [r.value for r in results]
        res_sorted: list[TypeDefinition] = sorted(
            res, key=lambda x: results_value.index(x.name)
        )

        # Rebase the file paths
        return self._rebase_types_file_paths(res_sorted)

    def get_type_calls(self, type_definition: TypeDefinition) -> list[TypeUsageInfo]:
        """Get the calls to a type definition. File paths are based on the challenge
        task container structure (e.g. /src)."""
        cqsearch_args = [
            "-s",
            self.CODEQUERY_DB,
            "-p",
            "8",
            "-t",
            type_definition.name,
            "-e",
            "-u",
        ]

        results = self._run_cqsearch(*cqsearch_args)

        logger.info("Found %d calls to type %s", len(results), type_definition.name)

        calls: list[TypeUsageInfo] = []
        for result in results:
            calls.append(
                TypeUsageInfo(
                    name=type_definition.name,
                    file_path=result.file,
                    line_number=result.line,
                )
            )

        return self._rebase_type_usages_file_paths(calls)


@dataclass
class CodeQueryPersistent(CodeQuery):
    """CodeQuery that we persist the status of the db

    It saves the db in the same workdir used by the challenge task, and it uses
    a copy of the challenge task named with the task-id + suffix. In this way it
    can always retrieve the db given a challenge task (even if it's a rw copy
    used by another instance).
    """

    work_dir: Path

    def __post_init__(self) -> None:
        """Post init the persistent codequery db"""
        task_id = self.challenge.task_meta.task_id
        cqdb_path = self.work_dir.joinpath(task_id + ".cqdb")
        try:
            self.challenge = ChallengeTask(cqdb_path, local_task_dir=cqdb_path)
        except ChallengeTaskError:
            # This is the case where the cqdb is not yet created
            logger.debug("Creating new CodeQueryPersistent DB in %s", cqdb_path)
            with self.challenge.get_rw_copy(self.work_dir) as persistent_challenge:
                self.challenge = persistent_challenge
                super().__post_init__()

                try:
                    persistent_challenge.commit(".cqdb")
                    logger.debug(
                        f"Uploading cqdb {persistent_challenge.local_task_dir} to remote storage"
                    )
                except Exception as e:
                    logger.exception("Failed to commit the cqdb: %s", e)
                    raise e

            return

        super().__post_init__()