main.py

#!/usr/bin/env python3
"""
GitHub Advisory Database Downloader and Analyzer.
This tool downloads security advisories and analyzes vulnerability fixes
using PyDriller to extract commit diffs and code modifications.
"""

import logging
import os
import shutil
import stat
import subprocess
import time

import psutil

from analysis import DatasetGenerator
from analysis.advisory_parser import AdvisoryParser
from analysis.repository_metadata_parser import RepositoryMetadataParser
from analysis.repository_test_checker import RepositoryTestChecker
from analysis.test_locator import TestLocator
from analysis.vulnerability_analyzer import VulnerabilityAnalyzer
from core.database import DB
from downloader.advisory_downloader import AdvisoryDownloader
from downloader.repository_metadata_downloader import RepositoryMetadataDownloader
from ui.cli import CLI
from utils import config

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.FileHandler("advisory_analyzer.log"), logging.StreamHandler()],
)
logger = logging.getLogger()


def setup_environment():
    """Create necessary directories for the application."""
    # Create data directory and subdirectories
    os.makedirs(config.DATA_DIR, exist_ok=True)
    os.makedirs(config.ADVISORIES_DIR, exist_ok=True)
    os.makedirs(config.ANALYSIS_DIR, exist_ok=True)
    os.makedirs(config.CWE_CACHE_DIR, exist_ok=True)
    os.makedirs(config.DATABASE_DIR, exist_ok=True)
    os.makedirs(config.REPOSITORIES_DIR, exist_ok=True)
    os.makedirs(config.TEMP_DIR, exist_ok=True)

    logger.info("Environment setup complete.")


def close_open_files(directory):
    """Find and close any processes using files in the given directory."""
    for proc in psutil.process_iter(["pid", "name"]):
        try:
            for file in proc.open_files():
                if directory in file.path:
                    logger.info(
                        f"Closing process {proc.pid} ({proc.name()}) using {file.path}"
                    )
                    proc.terminate()  # Terminate the process
                    proc.wait(timeout=5)  # Ensure process exits
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            continue  # Ignore processes that have already exited


def force_delete_git_files(directory):
    """Forcefully delete .git folders and make them writable."""
    for root, dirs, files in os.walk(directory, topdown=False):
        for file in files:
            file_path = os.path.join(root, file)
            try:
                os.chmod(file_path, stat.S_IWRITE)  # Ensure writable
                os.remove(file_path)
            except PermissionError:
                logger.warning(f"Could not delete file: {file_path}")

        for dir in dirs:
            dir_path = os.path.join(root, dir)
            try:
                os.chmod(dir_path, stat.S_IWRITE)
                shutil.rmtree(dir_path, ignore_errors=True)
            except PermissionError:
                logger.warning(f"Could not delete directory: {dir_path}")


def delete_cloned_repositories():
    """Safely delete the cloned_repositories directory, ensuring no file locks remain."""
    cloned_repos_path = config.TEMP_DIR

    if os.path.exists(cloned_repos_path):
        logger.info("Attempting to delete cloned repositories...")

        kill_git_processes()  # Kill Git processes before deletion
        # Ensure no processes are using files
        close_open_files(cloned_repos_path)

        for attempt in range(5):  # Retry mechanism
            try:
                # Remove .git files first
                force_delete_git_files(cloned_repos_path)
                shutil.rmtree(cloned_repos_path)  # Delete the directory
                logger.info("Successfully deleted cloned repositories.")
                break
            except PermissionError as e:
                logger.warning(
                    f"Permission error while deleting {cloned_repos_path}: {e}"
                )
                time.sleep(2)  # Wait and retry
            except Exception as e:
                logger.error(f"Unexpected error deleting {cloned_repos_path}: {e}")
                break


def kill_git_processes():
    """Kill all active Git processes to unlock files."""
    try:
        if os.name == "nt":  # Windows
            subprocess.run(
                ["taskkill", "/IM", "git.exe", "/F"],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        else:  # Mac/Linux
            subprocess.run(
                ["pkill", "-f", "git"],
                check=False,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        logger.info("Killed Git processes before deletion.")
    except Exception as e:
        logger.warning(f"Could not kill Git processes: {e}")


def main():
    """Main application flow"""

    print("GitHub Advisory Database Downloader and Analyzer")
    print("=" * 50)
    print("This tool downloads security advisories and analyzes vulnerability fixes")
    print("using PyDriller to extract commit diffs and code modifications.")
    print("=" * 50)

    # Setup environment (create directories)
    setup_environment()

    # Set the database path
    db_path = os.path.join(config.DATABASE_DIR, config.DATABASE_NAME)

    # Check if database already exists
    if os.path.exists(db_path):
        print(f"\nDatabase already exists at: {db_path}")
        choice = input(
            "Do you want to overwrite the existing database? (y/n): "
        ).lower()

        if choice == "y" or choice == "yes":
            logger.info("User chose to overwrite the existing database.")

            # TODO: Add a confirmation prompt before overwriting
            #  For now commented out

            # Remove the existing database
            # os.remove(db_path)
            logger.info("Removed existing database.")
        else:
            logger.info(
                "User chose to add data incrementally to the existing database."
            )
    db = DB(db_path, echo=False)
    db.initialize(logger)  # Initialize and return the session

    try:
        # Create instances of core components
        advisory_downloader = AdvisoryDownloader()
        advisory_parser = AdvisoryParser(db)  # For parsing advisories
        vulnerability_analyzer = VulnerabilityAnalyzer(db)  # For analyzing repositories and commits
        repository_metadata_downloader = RepositoryMetadataDownloader(db)
        repository_metadata_parser = RepositoryMetadataParser(db)
        repository_test_checker = RepositoryTestChecker(db)
        test_locator = TestLocator(db)
        dataset_generator = DatasetGenerator()

        # Run the CLI menu
        cli = CLI(
            db,
            advisory_downloader,
            advisory_parser,
            vulnerability_analyzer,
            repository_metadata_downloader,
            repository_metadata_parser,
            repository_test_checker,
            test_locator,
            dataset_generator,
        )  # Inject all components
        cli.menu()

    finally:
        # Ensure the database session is closed on exit
        db.close_session()

        # Delete cloned repositories
        # delete_cloned_repositories()


if __name__ == "__main__":
    main()