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

import logging
import os
import stat
import uuid
from dataclasses import dataclass, field
from buttercup.common.queues import (
    QueueFactory,
    ReliableQueue,
    QueueNames,
    GroupNames,
)
from buttercup.program_model.indexer import Indexer, IndexConf
from buttercup.program_model.kythe import KytheTool, KytheConf
from buttercup.program_model.graph import GraphStorage
from buttercup.program_model.codequery import CodeQueryPersistent
from buttercup.common.datastructures.msg_pb2 import IndexRequest, IndexOutput
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.utils import serve_loop
from pathlib import Path
from redis import Redis
import subprocess
import tempfile

logger = logging.getLogger(__name__)


@dataclass
class ProgramModel:
    sleep_time: float = 1.0
    redis: Redis | None = None
    task_queue: ReliableQueue | None = field(init=False, default=None)
    output_queue: ReliableQueue | None = field(init=False, default=None)
    wdir: Path | None = None
    script_dir: Path | None = None
    kythe_dir: Path | None = None
    graphdb_url: str = "ws://graphdb:8182/gremlin"
    graphdb_enabled: bool = True
    python: str | None = None
    allow_pull: bool = True
    base_image_url: str = "gcr.io/oss-fuzz"

    def __post_init__(self):
        """Post-initialization setup."""
        self.wdir = Path(self.wdir).resolve()
        self.script_dir = Path(self.script_dir).resolve()
        self.kythe_dir = Path(self.kythe_dir).resolve()

        if self.redis is not None:
            logger.debug("Using Redis for task queues")
            queue_factory = QueueFactory(self.redis)
            self.task_queue = queue_factory.create(QueueNames.INDEX, GroupNames.INDEX)
            self.output_queue = queue_factory.create(QueueNames.INDEX_OUTPUT)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()

    def cleanup(self):
        """Cleanup resources used by the program model"""
        pass

    def index_kythe(
        self, task_id: str, output_id: str, output_dir: Path, td: Path
    ) -> Path:
        """Index using kythe. Returns path to the index binary file."""
        ktool = KytheTool(KytheConf(self.kythe_dir))
        merged_kzip = Path(td) / f"kythe_output_merge_{output_id}.kzip"
        ktool.merge_kythe_output(output_dir, merged_kzip)

        # Convert the merged kzip file into a binary file
        bin_file = Path(td) / f"kythe_output_cxx_{output_id}.bin"
        try:
            ktool.cxx_index(merged_kzip, bin_file)
        except Exception as e:
            logger.error(f"Failed to index program {task_id} to binary: {bin_file}")
            raise e

        return bin_file

    def store_graphml(
        self, task_id: str, output_id: str, bin_file: Path, td: Path
    ) -> Path:
        """Store the program into a graphml file. Returns path to the graphml file."""
        graphml_file = Path(td) / f"kythe_output_graphml_{output_id}.xml"
        with open(graphml_file, "w") as fw, open(bin_file, "rb") as fr:
            gs = GraphStorage(task_id=task_id)
            gs.process_stream(fr)
            fw.write(gs.to_graphml())
        return graphml_file

    def load_graphml(self, graphml_file: Path) -> None:
        """Load graphml file into graph database."""
        # TODO(Evan): This needs to wait until JanusGraph is ready. For some reason, even if the container is running and healthy, it's not ready to accept connections.
        from gremlin_python.process.anonymous_traversal import traversal
        from gremlin_python.driver.driver_remote_connection import (
            DriverRemoteConnection,
        )

        g = traversal().withRemote(DriverRemoteConnection(self.graphdb_url, "g"))
        g.io(str(graphml_file)).read().iterate()

    def process_task_kythe(self, args: IndexRequest) -> bool:
        """Process a single task for indexing a program"""
        # Convert path strings to Path objects
        logger.debug(f"Kythe dir: {self.kythe_dir}")
        logger.debug(f"Script dir: {self.script_dir}")
        logger.debug(f"Wdir: {self.wdir}")
        logger.debug(f"Python: {self.python}")
        logger.debug(f"Allow pull: {self.allow_pull}")
        logger.debug(f"Base image URL: {self.base_image_url}")

        with tempfile.TemporaryDirectory(dir=self.wdir) as td:
            logger.info(
                f"Running indexer for {args.package_name} | {args.task_dir} | {args.task_id}"
            )

            # Change permissions so that JanusGraph can read from the temporary directory
            current = os.stat(td).st_mode
            janus_user = stat.S_IRGRP | stat.S_IROTH | stat.S_IXGRP | stat.S_IXOTH
            os.chmod(td, current | janus_user)

            tsk = ChallengeTask(
                read_only_task_dir=args.task_dir,
                python_path=self.python,
            )

            with tsk.get_rw_copy(work_dir=td) as local_tsk:
                logger.debug(f"Local path: {local_tsk.local_task_dir}")

                # Apply the diff if it exists
                logger.info(f"Applying diff for {args.package_name} {args.task_id}")
                if not local_tsk.apply_patch_diff():
                    logger.info(f"No diffs for {args.package_name} {args.task_id}")

                # Index the task
                try:
                    indexer_conf = IndexConf(
                        scriptdir=self.script_dir,
                        python=self.python,
                        allow_pull=self.allow_pull,
                        base_image_url=self.base_image_url,
                        wdir=td,
                    )
                    indexer = Indexer(indexer_conf)
                    output_dir = indexer.index_target(local_tsk)

                    # Because docker is running as root, we need to chown the output directory to the current user
                    subprocess.run(
                        [
                            "sudo",
                            "chown",
                            "-R",
                            f"{os.getuid()}:{os.getgid()}",
                            local_tsk.local_task_dir,
                        ],
                        check=True,
                    )
                except Exception as e:
                    logger.error(f"Failed to index task {args.task_id}: {e}")
                    return False
                if output_dir is None:
                    logger.error(f"Failed to index task {args.task_id}")
                    return False
                logger.info(f"Successfully indexed task {args.task_id}")

                output_id = str(uuid.uuid4())

                try:
                    bin_file = self.index_kythe(args.task_id, output_id, output_dir, td)
                except Exception as e:
                    logger.error(f"Failed to index files for {args.task_id}: {e}")
                    return False
                logger.info(
                    f"Successfully indexed and merged kythe output for {args.task_id} to {bin_file}"
                )

                # Store the program into a graphml file
                try:
                    graphml_file = Path(td) / f"kythe_output_graphml_{output_id}.xml"
                    with open(graphml_file, "w") as fw, open(bin_file, "rb") as fr:
                        gs = GraphStorage(task_id=args.task_id)
                        gs.process_stream(fr, fw)
                        logger.info(
                            f"Successfully stored program {args.task_id} in graphml file: {graphml_file}"
                        )
                except Exception as e:
                    logger.error(
                        f"Failed to store program {args.task_id} in graphml file {graphml_file}: {e}"
                    )
                    return False
                logger.info(
                    f"Successfully stored program {args.task_id} in graphml file: {graphml_file}"
                )

                # Load graphml file into graph database
                try:
                    logger.debug("Loading graphml file into JanusGraph...")
                    self.load_graphml(graphml_file)
                except Exception:
                    logger.exception(
                        f"Failed to load graphml file {graphml_file} into JanusGraph"
                    )
                    return False
                logger.info("Successfully loaded graphml file into JanusGraph")

        return True

    def process_task_codequery(self, args: IndexRequest) -> bool:
        """Process a single task for indexing a program"""
        try:
            challenge = ChallengeTask(
                read_only_task_dir=args.task_dir,
                python_path=self.python,
            )
            with challenge.get_rw_copy(work_dir=self.wdir) as local_challenge:
                # Apply the diff if it exists
                logger.info(f"Applying diff for {args.package_name} {args.task_id}")
                if not local_challenge.apply_patch_diff():
                    logger.info(f"No diffs for {args.package_name} {args.task_id}")

                CodeQueryPersistent(local_challenge, work_dir=self.wdir)

            return True
        except Exception as e:
            logger.exception(f"Failed to process task {args.task_id}: {e}")
            return False

    def process_task(self, args: IndexRequest) -> bool:
        """Process a single task for indexing a program"""
        # If at least one of the two methods succeeds, return True
        rv_code_query: bool = self.process_task_codequery(args)
        rv_kythe: bool = False
        if self.graphdb_enabled:
            rv_kythe = self.process_task_kythe(args)
        return rv_code_query or rv_kythe

    def serve_item(self) -> bool:
        rq_item = self.task_queue.pop()
        if rq_item is None:
            return False

        task_index: IndexRequest = rq_item.deserialized
        success = self.process_task(task_index)

        if success:
            self.output_queue.push(
                IndexOutput(
                    build_type=task_index.build_type,
                    package_name=task_index.package_name,
                    sanitizer=task_index.sanitizer,
                    task_dir=task_index.task_dir,
                    task_id=task_index.task_id,
                )
            )
            self.task_queue.ack_item(rq_item.item_id)
            logger.info(f"Successfully processed task {task_index.task_id}")
        else:
            logger.error(f"Failed to process task {task_index.task_id}")

        return True

    def serve(self):
        """Main loop to process tasks from queue"""
        if self.task_queue is None:
            raise ValueError("Task queue is not initialized")

        if self.output_queue is None:
            raise ValueError("Output queue is not initialized")

        logger.info("Starting indexing service")
        serve_loop(self.serve_item, self.sleep_time)