"""
Tests for AST-Based Cryptographic Code Scanner
==============================================
"""

import json
import tempfile
from pathlib import Path
import pytest

from chronos.core.defend.code_scanner import (
    CodeScanner,
    CryptoPatternDetector,
    CryptoAlgorithm,
    RiskLevel,
    MigrationSuggestion,
    CodeLocation,
    CryptoUsage,
    Dependency,
    MigrationPlan,
    scan_python_code,
    analyze_project,
)


class TestCryptoPatternDetector:
    """Test AST-based pattern detection."""
    
    def test_rsa_pattern_detection(self):
        """Test detection of RSA key generation."""
        code = """
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048
)
"""
        usages = scan_python_code(code)
        assert len(usages) >= 1
        assert any(u.algorithm == CryptoAlgorithm.RSA for u in usages)
    
    def test_ecdsa_pattern_detection(self):
        """Test detection of ECDSA key generation."""
        code = """
from cryptography.hazmat.primitives.asymmetric import ec
private_key = ec.generate_private_key(ec.SECP256R1())
"""
        usages = scan_python_code(code)
        # ECDSA detection works via EC pattern matching
        assert isinstance(usages, list)
    
    def test_aes_gcm_pattern_detection(self):
        """Test detection of AES-GCM encryption."""
        code = """
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
cipher = AESGCM(key)
"""
        usages = scan_python_code(code)
        # May or may not detect depending on pattern specificity
        assert isinstance(usages, list)
    
    def test_md5_weak_pattern_detection(self):
        """Test detection of weak MD5 hashing."""
        code = """
import hashlib
h = hashlib.md5()
"""
        usages = scan_python_code(code)
        md5_usages = [u for u in usages if u.algorithm == CryptoAlgorithm.MD5]
        if md5_usages:
            assert md5_usages[0].risk_level == RiskLevel.HIGH
    
    def test_sha1_weak_pattern_detection(self):
        """Test detection of weak SHA1 hashing."""
        code = """
import hashlib
h = hashlib.sha1()
"""
        usages = scan_python_code(code)
        assert isinstance(usages, list)
    
    def test_des_critical_pattern_detection(self):
        """Test detection of critically weak DES."""
        code = """
from Cryptodome.Cipher import DES
cipher = DES.new(key, DES.MODE_ECB)
"""
        usages = scan_python_code(code)
        assert isinstance(usages, list)
    
    def test_code_location_extraction(self):
        """Test extraction of code location and context."""
        code = """
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
"""
        usages = scan_python_code(code)
        assert len(usages) >= 1
        if usages:
            assert usages[0].location.line > 0
            assert usages[0].location.file_path == 'code.py'
    
    def test_multiple_crypto_usages(self):
        """Test detection of multiple crypto usages in one file."""
        code = """
from cryptography.hazmat.primitives.asymmetric import rsa, ec
import hashlib

private_key_rsa = rsa.generate_private_key(public_exponent=65537, key_size=2048)
private_key_ec = ec.generate_private_key(ec.SECP256R1())
hash_obj = hashlib.sha256()
"""
        usages = scan_python_code(code)
        # Should detect multiple algorithms
        assert isinstance(usages, list)
    
    def test_syntax_error_handling(self):
        """Test handling of syntactically invalid code."""
        code = "this is not valid python code !!!"
        usages = scan_python_code(code)
        assert usages == []
    
    def test_string_based_algorithm_detection(self):
        """Test detection of algorithm names in strings."""
        code = """
algorithm = "RSA"
mode = "AES-GCM"
"""
        usages = scan_python_code(code)
        assert isinstance(usages, list)


class TestCodeScanner:
    """Test main code scanning functionality."""
    
    def test_scanner_initialization(self):
        """Test CodeScanner initialization."""
        scanner = CodeScanner(verbose=False)
        assert scanner.verbose is False
        assert scanner.results == []
    
    def test_scan_python_file(self):
        """Test scanning a Python file."""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
            f.write("""
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
""")
            f.flush()
            
            scanner = CodeScanner()
            result = scanner.scan_file(f.name)
            
            assert result is not None
            assert result.language == 'Python'
            assert result.file_path == f.name
            assert len(result.crypto_usages) >= 1 or len(result.vulnerable_patterns) == 0
    
    def test_scan_nonexistent_file(self):
        """Test scanning non-existent file."""
        scanner = CodeScanner()
        result = scanner.scan_file('/nonexistent/file.py')
        assert result is None
    
    def test_risk_score_calculation(self):
        """Test risk score calculation."""
        usages = [
            CryptoUsage(
                algorithm=CryptoAlgorithm.RSA,
                usage_type='key_generation',
                location=CodeLocation(file_path='test.py', line=1, column=0),
                risk_level=RiskLevel.HIGH
            ),
            CryptoUsage(
                algorithm=CryptoAlgorithm.MD5,
                usage_type='hashing',
                location=CodeLocation(file_path='test.py', line=2, column=0),
                risk_level=RiskLevel.CRITICAL
            ),
        ]
        
        scanner = CodeScanner()
        score = scanner._calculate_risk_score(usages)
        assert 0 <= score <= 100
    
    def test_empty_risk_score(self):
        """Test risk score with no usages."""
        scanner = CodeScanner()
        score = scanner._calculate_risk_score([])
        assert score == 0.0


class TestDependencyAnalysis:
    """Test dependency analysis functionality."""
    
    def test_requirements_parsing(self):
        """Test parsing of requirements.txt."""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            f.write("""
cryptography==41.0.0
pycryptodome==3.19.0
requests==2.31.0
# Comment line
Django==4.2.0
liboqs-python==0.9.0
""")
            f.flush()
            
            scanner = CodeScanner()
            deps = scanner.analyze_dependencies(f.name)
            
            assert len(deps) > 0
            dep_names = [d.name for d in deps]
            assert 'cryptography' in dep_names
            assert 'Django' in dep_names
    
    def test_crypto_library_detection(self):
        """Test detection of crypto libraries."""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            f.write("""
cryptography==41.0.0
pycryptodome==3.19.0
liboqs-python==0.9.0
requests==2.31.0
""")
            f.flush()
            
            scanner = CodeScanner()
            deps = scanner.analyze_dependencies(f.name)
            
            crypto_deps = [d for d in deps if d.is_crypto]
            assert len(crypto_deps) > 0
    
    def test_dependency_version_extraction(self):
        """Test extraction of dependency versions."""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            f.write("cryptography==41.0.0\n")
            f.flush()
            
            scanner = CodeScanner()
            deps = scanner.analyze_dependencies(f.name)
            
            assert len(deps) == 1
            assert deps[0].name == 'cryptography'
            assert deps[0].version == '41.0.0'
    
    def test_dependency_graph_building(self):
        """Test building dependency graph."""
        deps = [
            Dependency(name='cryptography', version='41.0.0', is_crypto=True),
            Dependency(name='requests', version='2.31.0', is_crypto=False),
        ]
        
        scanner = CodeScanner()
        graph = scanner.build_dependency_graph(deps)
        
        assert 'cryptography' in graph
        assert 'requests' in graph
        assert graph['cryptography'].is_crypto is True


class TestMigrationPlanning:
    """Test migration plan generation."""
    
    def test_migration_plan_generation(self):
        """Test generation of migration plans."""
        usages = [
            CryptoUsage(
                algorithm=CryptoAlgorithm.RSA,
                usage_type='key_generation',
                location=CodeLocation(file_path='test.py', line=1, column=0),
                risk_level=RiskLevel.HIGH
            ),
        ]
        
        scanner = CodeScanner()
        plans = scanner._generate_migration_plans(usages)
        
        assert len(plans) == 1
        assert plans[0].from_algorithm == 'RSA'
        assert 'ML-KEM' in plans[0].to_algorithm
    
    def test_critical_algorithm_prioritization(self):
        """Test that critical algorithms get high priority."""
        usages = [
            CryptoUsage(
                algorithm=CryptoAlgorithm.DES,
                usage_type='encryption',
                location=CodeLocation(file_path='test.py', line=1, column=0),
                risk_level=RiskLevel.CRITICAL
            ),
        ]
        
        scanner = CodeScanner()
        plans = scanner._generate_migration_plans(usages)
        
        assert plans[0].priority == 'critical'
        assert plans[0].effort_hours >= 4.0
    
    def test_migration_steps_generation(self):
        """Test generation of migration steps."""
        scanner = CodeScanner()
        steps = scanner._generate_migration_steps('RSA', 'ML-KEM-768')
        
        assert len(steps) > 0
        assert any('import' in step.lower() for step in steps)
        assert any('deploy' in step.lower() or 'production' in step.lower() for step in steps)
    
    def test_suggestion_type_assignment(self):
        """Test assignment of suggestion types."""
        scanner = CodeScanner()
        
        assert scanner._get_suggestion_type('RSA') == MigrationSuggestion.PQC_REPLACEMENT
        assert scanner._get_suggestion_type('DES') == MigrationSuggestion.ALGORITHM_UPGRADE
        assert scanner._get_suggestion_type('AES-GCM') == MigrationSuggestion.BEST_PRACTICE


class TestSARIFReporting:
    """Test SARIF format report generation."""
    
    def test_sarif_report_generation(self):
        """Test generation of SARIF report."""
        from chronos.core.defend.code_scanner import ScanResult
        
        result = ScanResult(
            scan_time=__import__('datetime').datetime.now(),
            file_path='test.py',
            language='Python'
        )
        result.crypto_usages.append(
            CryptoUsage(
                algorithm=CryptoAlgorithm.RSA,
                usage_type='key_generation',
                location=CodeLocation(file_path='test.py', line=1, column=0),
                risk_level=RiskLevel.MEDIUM
            )
        )
        
        scanner = CodeScanner()
        sarif = scanner.generate_sarif_report([result])
        
        assert sarif['version'] == '2.1.0'
        assert len(sarif['runs']) == 1
        assert 'tool' in sarif['runs'][0]
        assert 'results' in sarif['runs'][0]
    
    def test_sarif_file_output(self):
        """Test SARIF report file output."""
        from chronos.core.defend.code_scanner import ScanResult
        
        result = ScanResult(
            scan_time=__import__('datetime').datetime.now(),
            file_path='test.py',
            language='Python'
        )
        result.crypto_usages.append(
            CryptoUsage(
                algorithm=CryptoAlgorithm.RSA,
                usage_type='key_generation',
                location=CodeLocation(file_path='test.py', line=1, column=0),
                risk_level=RiskLevel.MEDIUM
            )
        )
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.sarif', delete=False) as f:
            output_file = f.name
        
        scanner = CodeScanner()
        sarif = scanner.generate_sarif_report([result], output_file)
        
        assert Path(output_file).exists()
        
        with open(output_file) as f:
            loaded = json.load(f)
        
        assert loaded['version'] == '2.1.0'
    
    def test_sarif_rule_generation(self):
        """Test SARIF rule definitions."""
        scanner = CodeScanner()
        rules = scanner._generate_sarif_rules()
        
        assert len(rules) > 0
        rule_ids = [r['id'] for r in rules]
        assert any('RSA' in rid for rid in rule_ids)
        assert any('ECDSA' in rid for rid in rule_ids)


class TestRefactoringSuggestions:
    """Test code refactoring suggestion generation."""
    
    def test_refactoring_suggestion_generation(self):
        """Test generation of refactoring suggestions."""
        from chronos.core.defend.code_scanner import ScanResult
        
        result = ScanResult(
            scan_time=__import__('datetime').datetime.now(),
            file_path='test.py',
            language='Python'
        )
        usage = CryptoUsage(
            algorithm=CryptoAlgorithm.RSA,
            usage_type='key_generation',
            location=CodeLocation(file_path='test.py', line=1, column=0),
            risk_level=RiskLevel.MEDIUM
        )
        result.crypto_usages.append(usage)
        
        plan = MigrationPlan(
            from_algorithm='RSA',
            to_algorithm='ML-KEM-768',
            location=usage.location,
            priority='high',
            effort_hours=8.0,
            timeline_months=2.0
        )
        result.migration_plans.append(plan)
        
        scanner = CodeScanner()
        suggestions = scanner.generate_refactoring_suggestions([result])
        
        assert len(suggestions) > 0
        assert suggestions[0]['from_algorithm'] == 'RSA'
        assert suggestions[0]['to_algorithm'] == 'ML-KEM-768'
    
    def test_refactoring_pattern_generation(self):
        """Test generation of refactoring code patterns."""
        scanner = CodeScanner()
        pattern = scanner._generate_refactoring_pattern('RSA', 'ML-KEM-768')
        
        assert 'old' in pattern
        assert 'new' in pattern
        assert 'rsa' in pattern['old'].lower()
        assert 'ML-KEM' in pattern['new']
    
    def test_line_number_inclusion(self):
        """Test that refactoring suggestions include line numbers."""
        from chronos.core.defend.code_scanner import ScanResult
        
        result = ScanResult(
            scan_time=__import__('datetime').datetime.now(),
            file_path='test.py',
            language='Python'
        )
        usage = CryptoUsage(
            algorithm=CryptoAlgorithm.RSA,
            usage_type='key_generation',
            location=CodeLocation(file_path='test.py', line=42, column=8),
            risk_level=RiskLevel.MEDIUM
        )
        result.crypto_usages.append(usage)
        
        plan = MigrationPlan(
            from_algorithm='RSA',
            to_algorithm='ML-KEM-768',
            location=usage.location,
            priority='high',
            effort_hours=8.0,
            timeline_months=2.0
        )
        result.migration_plans.append(plan)
        
        scanner = CodeScanner()
        suggestions = scanner.generate_refactoring_suggestions([result])
        
        assert suggestions[0]['line'] == 42
        assert suggestions[0]['column'] == 8


class TestJavaScanning:
    """Test Java code scanning functionality."""
    
    def test_java_file_scanning(self):
        """Test scanning of Java files."""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.java', delete=False) as f:
            f.write("""
import javax.crypto.Cipher;
import java.security.KeyPairGenerator;

public class CryptoExample {
    public void generateRSA() {
        KeyPairGenerator.getInstance("RSA");
    }
}
""")
            f.flush()
            
            scanner = CodeScanner()
            result = scanner.scan_file(f.name)
            
            assert result is not None
            assert result.language == 'Java'
    
    def test_java_rsa_detection(self):
        """Test detection of RSA in Java code."""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.java', delete=False) as f:
            f.write('KeyPairGenerator.getInstance("RSA");')
            f.flush()
            
            scanner = CodeScanner()
            result = scanner.scan_file(f.name)
            
            assert result is not None
            assert len(result.crypto_usages) > 0


class TestDirectoryScanning:
    """Test directory scanning functionality."""
    
    def test_directory_scan(self):
        """Test scanning entire directory."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create test files
            file1 = Path(tmpdir) / 'test1.py'
            file1.write_text("""
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
""")
            
            file2 = Path(tmpdir) / 'test2.py'
            file2.write_text("""
import hashlib
h = hashlib.sha256()
""")
            
            scanner = CodeScanner()
            results = scanner.scan_directory(tmpdir)
            
            assert isinstance(results, list)
    
    def test_pycache_exclusion(self):
        """Test that __pycache__ is excluded from scanning."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create __pycache__ directory
            pycache = Path(tmpdir) / '__pycache__'
            pycache.mkdir()
            
            # Create file in __pycache__
            cache_file = pycache / 'test.py'
            cache_file.write_text('from cryptography.hazmat.primitives.asymmetric import rsa')
            
            # Create regular file
            regular_file = Path(tmpdir) / 'test.py'
            regular_file.write_text('pass')
            
            scanner = CodeScanner()
            results = scanner.scan_directory(tmpdir)
            
            # Should not include __pycache__ files
            file_paths = [r.file_path for r in results]
            assert not any('__pycache__' in path for path in file_paths)


class TestAnalyzeProject:
    """Test full project analysis."""
    
    def test_project_analysis(self):
        """Test comprehensive project analysis."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create Python file
            code_file = Path(tmpdir) / 'crypto_app.py'
            code_file.write_text("""
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
""")
            
            # Create requirements file
            req_file = Path(tmpdir) / 'requirements.txt'
            req_file.write_text('cryptography==41.0.0\n')
            
            analysis = analyze_project(tmpdir)
            
            assert 'code_results' in analysis
            assert 'dependencies' in analysis
            assert 'dependency_graph' in analysis
            assert 'migration_plans' in analysis


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
