test_setup.py

#!/usr/bin/env python3
"""
Setup verification script for Volatility3 Triage Tool.
Run this to check if your environment is properly configured.
"""

import sys
import os
import subprocess
from pathlib import Path


def check_python_version():
    """Check Python version."""
    print("Checking Python version...")
    version = sys.version_info
    if version.major >= 3 and version.minor >= 8:
        print(f"  ✓ Python {version.major}.{version.minor}.{version.micro}")
        return True
    else:
        print(f"  ✗ Python {version.major}.{version.minor}.{version.micro} (need 3.8+)")
        return False


def check_volatility():
    """Check if Volatility3 is installed."""
    print("\nChecking Volatility3 installation...")
    try:
        result = subprocess.run(
            ['vol', '--help'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode == 0:
            print("  ✓ Volatility3 'vol' command found")
            return True
        else:
            print("  ✗ Volatility3 'vol' command not working")
            return False
    except FileNotFoundError:
        print("  ✗ Volatility3 'vol' command not found in PATH")
        return False
    except Exception as e:
        print(f"  ✗ Error checking Volatility3: {e}")
        return False


def check_dependencies():
    """Check if Python dependencies are installed."""
    print("\nChecking Python dependencies...")

    required = [
        'fastapi',
        'uvicorn',
        'pydantic',
    ]

    optional_llm = [
        ('google.generativeai', 'Google Gemini'),
        ('anthropic', 'Anthropic Claude'),
    ]

    all_ok = True
    for package in required:
        try:
            __import__(package)
            print(f"  ✓ {package}")
        except ImportError:
            print(f"  ✗ {package} (not installed)")
            all_ok = False

    # Check for at least one LLM provider
    print("\nChecking LLM providers (need at least one)...")
    llm_available = False
    for package, name in optional_llm:
        try:
            __import__(package)
            print(f"  ✓ {package} ({name})")
            llm_available = True
        except ImportError:
            print(f"  ✗ {package} ({name}) (not installed)")

    if not llm_available:
        print("  ! ERROR: No LLM provider installed!")
        print("    Install at least one: pip install google-generativeai")
        all_ok = False

    return all_ok


def check_api_key():
    """Check if LLM API key is set."""
    print("\nChecking LLM API configuration...")

    # First check for .env file
    env_file = Path('.env')
    if env_file.exists():
        print(f"  ✓ .env file found")
        try:
            from dotenv import load_dotenv
            load_dotenv()
        except ImportError:
            print("  ! Warning: python-dotenv not installed, .env file won't be loaded")
    else:
        print(f"  ! .env file not found (optional)")

    # Determine provider
    provider = os.environ.get('LLM_PROVIDER', 'google')
    print(f"  LLM Provider: {provider}")

    # Check if appropriate API key is set
    if provider == 'google':
        api_key = os.environ.get('GOOGLE_API_KEY')
        if api_key:
            print(f"  ✓ GOOGLE_API_KEY is set ({api_key[:10]}...)")
            return True
        else:
            print("  ✗ GOOGLE_API_KEY not set")
            print("     Get your key from: https://aistudio.google.com/apikey")
            print("     Set in .env: GOOGLE_API_KEY=your-key-here")
            return False
    elif provider == 'anthropic':
        api_key = os.environ.get('ANTHROPIC_API_KEY')
        if api_key:
            print(f"  ✓ ANTHROPIC_API_KEY is set ({api_key[:10]}...)")
            return True
        else:
            print("  ✗ ANTHROPIC_API_KEY not set")
            print("     Get your key from: https://console.anthropic.com/")
            print("     Set in .env: ANTHROPIC_API_KEY=your-key-here")
            return False
    else:
        print(f"  ! Unknown provider: {provider}")
        print("    Valid providers: google, anthropic")
        return False


def check_directories():
    """Check if required directories exist."""
    print("\nChecking project structure...")

    required_dirs = [
        'backend',
        'static',
        'templates',
    ]

    # Optional directory (will be auto-created)
    optional_dirs = [
        'runs',
    ]

    all_ok = True
    for dirname in required_dirs:
        path = Path(dirname)
        if path.exists() and path.is_dir():
            print(f"  ✓ {dirname}/")
        else:
            print(f"  ✗ {dirname}/ (missing)")
            all_ok = False

    # Check optional directories (informational only)
    for dirname in optional_dirs:
        path = Path(dirname)
        if path.exists() and path.is_dir():
            print(f"  ✓ {dirname}/ (optional)")
        else:
            print(f"  ℹ {dirname}/ (will be auto-created)")

    return all_ok


def check_backend_modules():
    """Check if backend modules are present."""
    print("\nChecking backend modules...")

    required_files = [
        'backend/__init__.py',
        'backend/config.py',
        'backend/runner.py',
        'backend/triage.py',
        'backend/reports.py',
        'backend/llm_interface.py',
        'backend/api.py',
    ]

    all_ok = True
    for filepath in required_files:
        path = Path(filepath)
        if path.exists() and path.is_file():
            print(f"  ✓ {filepath}")
        else:
            print(f"  ✗ {filepath} (missing)")
            all_ok = False

    return all_ok


def main():
    """Run all checks."""
    print("=" * 60)
    print("Volatility3 Triage Tool - Setup Verification")
    print("=" * 60)

    results = []

    results.append(("Python Version", check_python_version()))
    results.append(("Volatility3", check_volatility()))
    results.append(("Python Dependencies", check_dependencies()))
    results.append(("LLM API Key", check_api_key()))
    results.append(("Project Structure", check_directories()))
    results.append(("Backend Modules", check_backend_modules()))

    print("\n" + "=" * 60)
    print("Summary")
    print("=" * 60)

    for name, passed in results:
        status = "✓ PASS" if passed else "✗ FAIL"
        print(f"{status:>8}  {name}")

    all_passed = all(result[1] for result in results)

    print("\n" + "=" * 60)
    if all_passed:
        print("✓ All checks passed! You're ready to run the tool.")
        print("\nStart the server with:")
        print("  ./run.sh            (Linux/macOS)")
        print("  run.bat             (Windows)")
        print("  python -m backend.api  (manual)")
    else:
        print("✗ Some checks failed. Please fix the issues above.")
        print("\nCommon fixes:")
        print("  - Install dependencies: pip install -r requirements.txt")
        print("  - Install Volatility3: Follow README.md instructions")
        print("  - For Google Gemini: export GOOGLE_API_KEY='your-key'")
        print("  - For Anthropic Claude: export ANTHROPIC_API_KEY='your-key'")
    print("=" * 60)

    return 0 if all_passed else 1


if __name__ == '__main__':
    sys.exit(main())