#!/usr/bin/env python3
"""
Script: analyze_python_deps.py
Purpose: Static analyzer to inventory and categorize Python dependencies for Go migration assessment.
Target Audience: Technical Security Operators & Software Architects in Canadian classified environments.
Output: JSON report and CLI summary for triage, suitable for inclusion in accreditation documentation.
"""

import ast
import os
import json
import argparse
import hashlib
from pathlib import Path
from collections import defaultdict

# Categorization Heuristics - Tailored for Canadian government security requirements
STDLIB_MODULES = {'json', 'math', 'http', 'os', 'sys', 'csv', 'datetime', 're', 'ssl', 'hashlib', 'socket', 'logging'}
HIGH_RISK_CEXTENSIONS = {'numpy', 'pandas', 'cryptography', 'PIL', 'lxml', 'psycopg2', 'MySQLdb', 'pyOpenSSL', 'cffi'}
HIGH_RISK_COMPLEX = {'django', 'flask', 'fastapi', 'tornado', 'scipy', 'tensorflow', 'torch', 'keras', 'scikit-learn'}
MEDIUM_RISK = {'requests', 'aiohttp', 'sqlalchemy', 'pydantic', 'jinja2', 'celery', 'redis', 'pika'}
GOVERNMENT_RESTRICTED = {'torch', 'tensorflow', 'keras', 'transformers'}  # AI/ML libraries with export restrictions

def get_imports_from_file(file_path):
    """Parse a Python file and extract all imported module names with enhanced security analysis."""
    imports = set()
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            tree = ast.parse(content, filename=file_path)
        
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for n in node.names:
                    base_module = n.name.split('.')[0]
                    imports.add(base_module)
            elif isinstance(node, ast.ImportFrom):
                if node.module is not None:
                    base_module = node.module.split('.')[0]
                    imports.add(base_module)
                    
    except (SyntaxError, UnicodeDecodeError):
        # Skip files that can't be parsed (e.g., binary, bad syntax)
        return set()
    return imports

def scan_project(root_path):
    """Walk through a project directory and find all imports with additional security metadata."""
    root = Path(root_path)
    all_imports = set()
    files_scanned = 0
    file_hashes = {}

    for py_file in root.rglob('*.py'):
        files_scanned += 1
        imports_in_file = get_imports_from_file(py_file)
        all_imports |= imports_in_file
        
        # Calculate file hash for integrity verification
        with open(py_file, 'rb') as f:
            file_content = f.read()
            file_hash = hashlib.sha256(file_content).hexdigest()
            file_hashes[str(py_file)] = file_hash

    # Remove standard library modules (false positives possible)
    third_party_imports = all_imports - STDLIB_MODULES
    return sorted(third_party_imports), files_scanned, file_hashes

def categorize_modules(modules):
    """Categorize modules based on pre-defined lists with government security considerations."""
    report = {
        'high_risk_c': [],       # C-Extensions
        'high_risk_py': [],      # Complex Pure-Python
        'medium_risk': [],
        'low_risk': [],          # Simple pure-python
        'restricted': [],        # Government restricted/export controlled
        'unknown_risk': []       # Modules not in any category
    }
    
    for module in modules:
        if module in GOVERNMENT_RESTRICTED:
            report['restricted'].append(module)
        elif module in HIGH_RISK_CEXTENSIONS:
            report['high_risk_c'].append(module)
        elif module in HIGH_RISK_COMPLEX:
            report['high_risk_py'].append(module)
        elif module in MEDIUM_RISK:
            report['medium_risk'].append(module)
        elif module.replace('_', '-') in MEDIUM_RISK:  # Handle naming variations
            report['medium_risk'].append(module)
        else:
            # Check if this might be a standard library module we missed
            try:
                __import__(module)
                if hasattr(__import__(module), '__file__') and 'site-packages' not in __import__(module).__file__:
                    # It's actually a standard library module
                    pass
                else:
                    report['unknown_risk'].append(module)
            except ImportError:
                report['unknown_risk'].append(module)
    
    return report

def generate_sbom(modules, file_hashes):
    """Generate a minimal Software Bill of Materials for security review."""
    sbom = {
        "bomFormat": "CycloneDX",
        "specVersion": "1.4",
        "version": 1,
        "components": [],
        "fileHashes": file_hashes
    }
    
    for module in modules:
        sbom["components"].append({
            "type": "library",
            "name": module,
            "version": "unknown",  # Would be extracted in a real implementation
            "purl": f"pkg:pypi/{module}@unknown"
        })
    
    return sbom

def main():
    parser = argparse.ArgumentParser(description='Analyze Python project dependencies for Go migration in classified environments.')
    parser.add_argument('project_path', help='Path to the root of the Python project to analyze.')
    parser.add_argument('-o', '--output', help='Output JSON file for the report.', default='migration_report.json')
    parser.add_argument('--sbom', help='Generate Software Bill of Materials', action='store_true')
    args = parser.parse_args()

    if not os.path.isdir(args.project_path):
        print(f"Error: Path '{args.project_path}' is not a valid directory.")
        return

    print(f"[+] Scanning project at: {args.project_path}")
    print("[+] This analysis is suitable for Canadian government security reviews")
    
    imports, file_count, file_hashes = scan_project(args.project_path)
    report = categorize_modules(imports)

    print(f"[+] Scanned {file_count} .py files.")
    print(f"[+] Found {len(imports)} unique third-party modules.")

    print("\n--- Categorization Report ---")
    print(f"Government Restricted: {report['restricted']}")
    print(f"High Risk (C Extensions): {report['high_risk_c']}")
    print(f"High Risk (Complex Pure-Python): {report['high_risk_py']}")
    print(f"Medium Risk: {report['medium_risk']}")
    print(f"Low Risk (Simple): {report['low_risk']}")
    print(f"Unknown Risk: {report['unknown_risk']}")

    # Calculate a comprehensive migration complexity score
    complexity_score = (len(report['restricted']) * 5 + len(report['high_risk_c']) * 4 +
                        len(report['high_risk_py']) * 3 + len(report['medium_risk']) * 2 +
                        len(report['low_risk']) * 1 + len(report['unknown_risk']) * 2)

    # Generate recommendations based on Canadian government context
    if complexity_score < 15:
        recommendation = 'High Priority Candidate - Suitable for pilot program'
    elif complexity_score < 30:
        recommendation = 'Medium Priority - Requires detailed security assessment'
    else:
        recommendation = 'High Risk / Low Priority - Consider alternative approaches'

    summary = {
        'project_path': args.project_path,
        'files_scanned': file_count,
        'third_party_modules': imports,
        'categorized_modules': report,
        'migration_complexity_score': complexity_score,
        'recommendation': recommendation,
        'security_considerations': [
            'Review restricted modules with Export Controls office',
            'Assess C extensions for potential memory safety issues',
            'Validate all cryptographic implementations against CSE guidance'
        ]
    }

    # Include SBOM if requested
    if args.sbom:
        summary['software_bill_of_materials'] = generate_sbom(imports, file_hashes)

    with open(args.output, 'w') as f:
        json.dump(summary, f, indent=4)

    print(f"\n[+] Migration Complexity Score: {complexity_score}")
    print(f"[+] Recommendation: {recommendation}")
    print(f"[+] Security Considerations: {summary['security_considerations']}")
    print(f"[+] Full report written to: {args.output}")
    
    if args.sbom:
        print(f"[+] Software Bill of Materials included in report")

if __name__ == '__main__':
    main()