"""
Tests for the malware analyzer functionality
"""
import pytest
import tempfile
import os
from unittest.mock import Mock, patch
from analyzer import MalwareAnalyzer

class TestMalwareAnalyzer:
    
    def setup_method(self):
        """Set up test fixtures"""
        self.analyzer = MalwareAnalyzer()
    
    def test_init(self):
        """Test analyzer initialization"""
        assert self.analyzer.sandbox_timeout == 300
        assert self.analyzer.virustotal_api_key is None or isinstance(self.analyzer.virustotal_api_key, str)
    
    def test_get_file_type(self):
        """Test file type detection"""
        with tempfile.NamedTemporaryFile(suffix='.txt', delete=False) as tmp_file:
            tmp_file.write(b"test content")
            tmp_file.flush()
            
            file_type = self.analyzer._get_file_type(tmp_file.name)
            assert file_type is not None
            assert isinstance(file_type, str)
            
            os.unlink(tmp_file.name)
    
    def test_is_executable(self):
        """Test executable file detection"""
        # Test with a text file (should not be executable)
        with tempfile.NamedTemporaryFile(suffix='.txt', delete=False) as tmp_file:
            tmp_file.write(b"not executable content")
            tmp_file.flush()
            
            assert not self.analyzer._is_executable(tmp_file.name)
            os.unlink(tmp_file.name)
    
    def test_static_analysis(self):
        """Test static analysis functionality"""
        # Create a test file with suspicious content
        test_content = b"""
        This is a test file with some suspicious content.
        It contains IP addresses like 192.168.1.1 and 10.0.0.1
        It also has URLs like https://malicious.com/download
        And suspicious commands like cmd.exe and powershell
        """
        
        with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
            tmp_file.write(test_content)
            tmp_file.flush()
            
            result = self.analyzer._static_analysis(tmp_file.name)
            
            assert 'features' in result
            assert 'ip_addresses' in result
            assert 'attack_vectors' in result
            
            # Should find IP addresses
            assert len(result['ip_addresses']) >= 2
            assert '192.168.1.1' in result['ip_addresses']
            assert '10.0.0.1' in result['ip_addresses']
            
            # Should find suspicious features
            assert any('suspicious_command' in feature for feature in result['features'])
            
            os.unlink(tmp_file.name)
    
    def test_categorize_context_words(self):
        """Test context word categorization"""
        words = ['company', 'software', 'bank', 'school', 'government', 'unknown_word']
        
        categorized = self.analyzer._categorize_context_words(words)
        
        assert categorized['company'] == 'business'
        assert categorized['software'] == 'technology'
        assert categorized['bank'] == 'finance'
        assert categorized['school'] == 'education'
        assert categorized['government'] == 'government'
        assert categorized['unknown_word'] == 'general'
    
    def test_classify_malware(self):
        """Test malware classification logic"""
        # Test backdoor classification
        analysis_result = {
            'features': ['network_communication', 'persistence_mechanism'],
            'attack_vectors': ['persistence']
        }
        
        classification = self.analyzer._classify_malware(analysis_result)
        assert classification == "Backdoor/Trojan"
        
        # Test ransomware classification
        analysis_result = {
            'features': ['encryption_obfuscation'],
            'attack_vectors': []
        }
        
        classification = self.analyzer._classify_malware(analysis_result)
        assert classification == "Ransomware"
        
        # Test RAT classification
        analysis_result = {
            'features': ['network_communication'],
            'attack_vectors': ['command_execution']
        }
        
        classification = self.analyzer._classify_malware(analysis_result)
        assert classification == "Remote Access Tool"
    
    @patch('analyzer.requests.get')
    def test_virustotal_lookup(self, mock_get):
        """Test VirusTotal API lookup"""
        # Mock successful response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            'response_code': 1,
            'positives': 5,
            'total': 60,
            'scan_date': '2023-01-01 12:00:00'
        }
        mock_get.return_value = mock_response
        
        # Set API key
        self.analyzer.virustotal_api_key = "test-api-key"
        
        result = self.analyzer._virustotal_lookup("test-hash")
        
        assert result is not None
        assert result['virustotal_detected'] is True
        assert result['virustotal_engines'] == 5
        assert result['virustotal_total'] == 60
    
    @patch('analyzer.requests.get')
    def test_virustotal_lookup_no_api_key(self, mock_get):
        """Test VirusTotal lookup without API key"""
        self.analyzer.virustotal_api_key = None
        
        result = self.analyzer._virustotal_lookup("test-hash")
        
        assert result is None
        mock_get.assert_not_called()
    
    def test_extract_context_words(self):
        """Test context word extraction"""
        file_info = {
            'page_url': 'https://business-company.com/software-download',
            'link_text': 'Download Business Software'
        }
        
        # Mock database session
        with patch('analyzer.db_session') as mock_db:
            mock_db.query.return_value.filter_by.return_value.first.return_value = None
            mock_db.query.return_value.scalar.return_value = "2023-01-01 12:00:00"
            
            context_words = self.analyzer._extract_context_words(file_info)
            
            assert len(context_words) > 0
            assert 'business' in context_words
            assert 'company' in context_words
            assert 'software' in context_words
    
    def test_analyze_file_nonexistent(self):
        """Test analysis of non-existent file"""
        file_info = {
            'local_path': '/nonexistent/file.exe',
            'filename': 'test.exe'
        }
        
        result = self.analyzer.analyze_file(file_info)
        
        assert result is None
    
    def test_analyze_file_success(self):
        """Test successful file analysis"""
        # Create a test file
        test_content = b"test executable content with IP 192.168.1.1"
        
        with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
            tmp_file.write(test_content)
            tmp_file.flush()
            
            file_info = {
                'local_path': tmp_file.name,
                'filename': 'test.exe',
                'file_hash': 'test-hash-123'
            }
            
            # Mock database and external API calls
            with patch('analyzer.db_session') as mock_db, \
                 patch.object(self.analyzer, '_virustotal_lookup', return_value=None), \
                 patch.object(self.analyzer, '_check_known_malware', return_value=None):
                
                mock_db.query.return_value.filter_by.return_value.first.return_value = None
                mock_db.query.return_value.scalar.return_value = "2023-01-01 12:00:00"
                
                result = self.analyzer.analyze_file(file_info)
                
                assert result is not None
                assert result['hash'] == 'test-hash-123'
                assert result['filename'] == 'test.exe'
                assert 'features' in result
                assert 'ip_addresses' in result
                assert 'attack_vectors' in result
                assert 'context_words' in result
                
            os.unlink(tmp_file.name)
