trailofbits-buttercup/program-model/tests/common.py

import logging
from dataclasses import dataclass
from pathlib import Path

import pytest
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.project_yaml import Language

from buttercup.program_model.codequery import CodeQuery
from buttercup.program_model.utils.common import (
    Function,
    TypeDefinition,
    TypeDefinitionType,
    TypeUsageInfo,
)

logger = logging.getLogger(__name__)


def filter_project_context(
    container_src_dir: str,
    results: list[Function | TypeDefinition | TypeUsageInfo],
    language: Language,
):
    """Some challenge tasks result in multiple instances of the target project to
    be built in the /src/ directory. This in turn causes Codequery to return multiple
    matches for queried context because it gets matches from parallel instances of the
    project.

    This function filters out found functions and types using the target project name.
    It removes all matches that don't directly come from the target project. To identify these,
    it checks whether matches belong to a file that as the proper project_name in their path."""
    if language == Language.C:
        return [x for x in results if f"/{container_src_dir}/" in str(x.file_path)]
    return results


@dataclass(frozen=True)
class TestFunctionInfo:
    num_bodies: int
    body_excerpts: list[str]


# Prevent pytest from collecting this as a test
TestFunctionInfo.__test__ = False


def common_test_get_functions(codequery: CodeQuery, function_name, file_path, function_info):
    """Generic function for testing get_functions() in C codebases"""
    if file_path:
        file_path = Path(file_path)
    functions = codequery.get_functions(function_name, file_path=file_path)
    assert len(functions) == 1
    assert functions[0].name == function_name
    if file_path:
        assert Path(functions[0].file_path) == file_path
    assert len(functions[0].bodies) == function_info.num_bodies
    for body in function_info.body_excerpts:
        assert any([body in x.body for x in functions[0].bodies])


@dataclass(frozen=True)
class TestCallerInfo:
    name: str
    file_path: Path
    start_line: int


# Prevent pytest from collecting this as a test
TestCallerInfo.__test__ = False


def common_test_get_callers(
    fuzz_task: ChallengeTask,
    codequery: CodeQuery,
    function_name,
    file_path,
    line_number,
    fuzzy,
    expected_callers,
    num_callers: int | None = None,
):
    """Test that we can get function callers"""
    function = codequery.get_functions(
        function_name=function_name,
        file_path=Path(file_path),
        line_number=line_number,
        fuzzy=fuzzy,
    )[0]

    callers = codequery.get_callers(function)
    callers = filter_project_context(fuzz_task.container_src_dir(), callers, codequery._get_project_language())

    # Output for debugging
    if not callers:
        logger.info(f"No callers found for {function_name} in {file_path}")
    for c in callers:
        logger.info(f"Caller: {c.name} {c.file_path} {[b.start_line for b in c.bodies]}")

    # Validate each caller
    for expected_caller in expected_callers:
        caller_info = [
            c
            for c in callers
            if c.name == expected_caller.name
            and c.file_path == Path(expected_caller.file_path)
            and any(True for b in c.bodies if b.start_line <= expected_caller.start_line <= b.end_line)
        ]
        if len(caller_info) == 0:
            pytest.fail(f"Couldn't find expected caller: {expected_caller}")
        elif len(caller_info) > 1:
            pytest.fail(f"Found multiple identical callers for: {expected_caller}")

    # Make sure we get the right number of callers
    if num_callers and len(callers) != num_callers:
        pytest.fail(f"Expected {num_callers} callers, got {len(callers)}")


@dataclass(frozen=True)
class TestCalleeInfo:
    name: str
    file_path: Path
    start_line: int


# Prevent pytest from collecting this as a test
TestCalleeInfo.__test__ = False


# NOTE(boyan): this is similar to the common_test_get_callers
# but we don't factorize the code so far in case we need to
# make the tests diverge in the future
def common_test_get_callees(
    fuzz_task: ChallengeTask,
    codequery: CodeQuery,
    function_name,
    file_path,
    line_number,
    fuzzy,
    expected_callees,
    num_callees: int | None = None,
):
    """Test that we can get function callees."""
    function = codequery.get_functions(
        function_name=function_name,
        file_path=Path(file_path),
        line_number=line_number,
        fuzzy=fuzzy,
    )[0]

    callees = codequery.get_callees(function)
    callees = filter_project_context(fuzz_task.container_src_dir(), callees, codequery._get_project_language())

    # Output for debugging
    if not callees:
        logger.info(f"No callees found for {function_name} in {file_path}")
    for c in callees:
        logger.info(f"Callee: {c.name} {c.file_path} {[b.start_line for b in c.bodies]}")

    # Validate each callee
    for expected_callee in expected_callees:
        callee_info = [
            c
            for c in callees
            if c.name == expected_callee.name
            and c.file_path == Path(expected_callee.file_path)
            and any(True for b in c.bodies if b.start_line <= expected_callee.start_line <= b.end_line)
        ]
        if len(callee_info) == 0:
            pytest.fail(f"Couldn't find expected callee: {expected_callee}")
        elif len(callee_info) > 1:
            pytest.fail(f"Found multiple identical callees for: {expected_callee}")

    # Make sure we don't get more callees than expected
    if num_callees and len(callees) != num_callees:
        pytest.fail(f"Expected {num_callees} callees, got {len(callees)}")


@dataclass(frozen=True)
class TestTypeDefinitionInfo:
    name: str
    type: TypeDefinitionType
    definition: str
    definition_line: int
    file_path: str


# Prevent pytest from collecting this as a test
TestTypeDefinitionInfo.__test__ = False


def common_test_get_type_definitions(
    fuzz_task: ChallengeTask,
    codequery: CodeQuery,
    type_name,
    file_path,
    fuzzy,
    type_definition_info,
):
    """Test that we can get type defs"""
    type_definitions = codequery.get_types(
        type_name=type_name,
        file_path=Path(file_path) if file_path else None,
        fuzzy=fuzzy,
    )
    type_definitions = filter_project_context(
        fuzz_task.container_src_dir(),
        type_definitions,
        codequery._get_project_language(),
    )

    # Output for debugging
    for c in type_definitions:
        logger.info(f"Type definition: {c.name} {c.type} {c.file_path} {c.definition_line}")

    found = [
        c
        for c in type_definitions
        if c.name == type_name
        and c.type == type_definition_info.type
        and type_definition_info.definition in c.definition
        and c.definition_line == type_definition_info.definition_line
        and c.file_path == Path(type_definition_info.file_path)
    ]
    if len(found) == 0:
        pytest.fail(f"Couldn't find expected type definition: {type_definition_info}")
    elif len(found) > 1:
        pytest.fail(f"Found multiple identical type definitions for: {type_definition_info}")


@dataclass(frozen=True)
class TestTypeUsageInfo:
    file_path: Path
    line_number: int


# Prevent pytest from collecting this as a test
TestTypeUsageInfo.__test__ = False


def common_test_get_type_usages(
    fuzz_task: ChallengeTask,
    codequery: CodeQuery,
    type_name,
    file_path,
    fuzzy,
    type_usage_infos,
    num_type_usages: int | None = None,
):
    """Test that we can get type usages"""
    type_definition = codequery.get_types(
        type_name=type_name,
        file_path=Path(file_path) if file_path else None,
        fuzzy=fuzzy,
    )[0]

    call_sites = codequery.get_type_calls(type_definition)
    call_sites = filter_project_context(fuzz_task.container_src_dir(), call_sites, codequery._get_project_language())

    # Output for debugging
    for c in call_sites:
        logger.info(f"Type usage: {c.name} {c.file_path} {c.line_number}")

    for type_usage_info in type_usage_infos:
        type_usage = [
            c
            for c in call_sites
            if c.name == type_name
            and c.file_path == Path(type_usage_info.file_path)
            and c.line_number == type_usage_info.line_number
        ]
        if len(type_usage) == 0:
            pytest.fail(f"Couldn't find expected type usage: {type_usage_info}")
        elif len(type_usage) > 1:
            pytest.fail(f"Found multiple identical type usages for: {type_usage_info}")

    if num_type_usages and len(call_sites) != num_type_usages:
        pytest.fail(f"Expected {num_type_usages} type usages, got {len(call_sites)}")