"""
Malware analysis and sandboxing system
"""
import os
import re
import json
import hashlib
import subprocess
import tempfile
import shutil
import time
import requests
import yara
import pefile
import magic
from pathlib import Path
from config import Config
from database import db_session
from models import KnownMalware, ContextWord
import logging

logger = logging.getLogger(__name__)

class MalwareAnalyzer:
    def __init__(self):
        self.yara_rules = self._load_yara_rules()
        self.known_malware_db = self._load_known_malware()
        self.virustotal_api_key = Config.VIRUSTOTAL_API_KEY
        self.sandbox_timeout = Config.SANDBOX_TIMEOUT
        
    def analyze_file(self, file_info):
        """
        Comprehensive malware analysis of a downloaded file
        """
        if not file_info.get('local_path') or not os.path.exists(file_info['local_path']):
            logger.error(f"File not found: {file_info.get('local_path')}")
            return None
        
        analysis_result = {
            'hash': file_info.get('file_hash'),
            'filename': file_info.get('filename'),
            'file_size': file_info.get('file_size'),
            'file_type': self._get_file_type(file_info['local_path']),
            'name': None,
            'confidence': 0.0,
            'features': [],
            'ip_addresses': [],
            'attack_vectors': [],
            'context_words': []
        }
        
        try:
            # 1. Static analysis
            static_analysis = self._static_analysis(file_info['local_path'])
            analysis_result.update(static_analysis)
            
            # 2. YARA rule matching
            yara_matches = self._yara_analysis(file_info['local_path'])
            if yara_matches:
                analysis_result['features'].extend(yara_matches)
                analysis_result['confidence'] += 0.3
            
            # 3. Known malware database lookup
            known_match = self._check_known_malware(analysis_result['hash'])
            if known_match:
                analysis_result.update(known_match)
                analysis_result['confidence'] += 0.4
            
            # 4. VirusTotal lookup
            if self.virustotal_api_key:
                vt_result = self._virustotal_lookup(analysis_result['hash'])
                if vt_result:
                    analysis_result.update(vt_result)
            
            # 5. Sandbox analysis (if executable)
            if self._is_executable(file_info['local_path']):
                sandbox_result = self._sandbox_analysis(file_info['local_path'])
                if sandbox_result:
                    analysis_result.update(sandbox_result)
            
            # 6. Extract context words from the page
            context_words = self._extract_context_words(file_info)
            analysis_result['context_words'] = context_words
            
            # 7. Determine if it's malware
            if analysis_result['confidence'] > 0.5:
                analysis_result['name'] = self._classify_malware(analysis_result)
            
            return analysis_result
            
        except Exception as e:
            logger.error(f"Error analyzing file {file_info['local_path']}: {str(e)}")
            return analysis_result

    def _static_analysis(self, file_path):
        """Perform static analysis on the file"""
        result = {
            'features': [],
            'ip_addresses': [],
            'attack_vectors': []
        }
        
        try:
            # Read file content
            with open(file_path, 'rb') as f:
                content = f.read()
            
            # Look for IP addresses
            ip_pattern = r'\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b'
            ips = re.findall(ip_pattern, content.decode('utf-8', errors='ignore'))
            result['ip_addresses'] = list(set(ips))
            
            # Look for URLs
            url_pattern = r'https?://[^\s<>"{}|\\^`\[\]]+'
            urls = re.findall(url_pattern, content.decode('utf-8', errors='ignore'))
            if urls:
                result['features'].append('network_communication')
            
            # Look for suspicious strings
            suspicious_strings = [
                'cmd.exe', 'powershell', 'regsvr32', 'rundll32', 'wscript',
                'cscript', 'certutil', 'bitsadmin', 'wmic', 'net user',
                'net localgroup', 'schtasks', 'at ', 'sc create'
            ]
            
            content_lower = content.lower()
            for string in suspicious_strings:
                if string.encode() in content_lower:
                    result['features'].append(f'suspicious_command_{string.replace(".", "_")}')
                    result['attack_vectors'].append('command_execution')
            
            # Look for encryption/obfuscation indicators
            if b'xor' in content_lower or b'encrypt' in content_lower:
                result['features'].append('encryption_obfuscation')
            
            # Look for persistence mechanisms
            persistence_indicators = [
                b'runonce', b'run', b'startup', b'autorun', b'shell:startup',
                b'registry', b'reg add', b'regedit'
            ]
            
            for indicator in persistence_indicators:
                if indicator in content_lower:
                    result['features'].append('persistence_mechanism')
                    result['attack_vectors'].append('persistence')
                    break
            
        except Exception as e:
            logger.error(f"Error in static analysis: {str(e)}")
        
        return result

    def _yara_analysis(self, file_path):
        """Analyze file using YARA rules"""
        matches = []
        
        try:
            if self.yara_rules:
                yara_matches = self.yara_rules.match(file_path)
                for match in yara_matches:
                    matches.append(f"yara_{match.rule}")
                    matches.extend([f"yara_tag_{tag}" for tag in match.tags])
        except Exception as e:
            logger.error(f"Error in YARA analysis: {str(e)}")
        
        return matches

    def _check_known_malware(self, file_hash):
        """Check against known malware database"""
        try:
            known_malware = db_session.query(KnownMalware).filter(
                (KnownMalware.md5_hash == file_hash) |
                (KnownMalware.sha1_hash == file_hash) |
                (KnownMalware.sha256_hash == file_hash)
            ).first()
            
            if known_malware:
                return {
                    'name': known_malware.name,
                    'features': known_malware.features or [],
                    'attack_vectors': known_malware.attack_vectors or [],
                    'ip_addresses': known_malware.c2_servers or []
                }
        except Exception as e:
            logger.error(f"Error checking known malware: {str(e)}")
        
        return None

    def _virustotal_lookup(self, file_hash):
        """Look up file hash in VirusTotal"""
        if not self.virustotal_api_key:
            return None
        
        try:
            url = f"https://www.virustotal.com/vtapi/v2/file/report"
            params = {
                'apikey': self.virustotal_api_key,
                'resource': file_hash
            }
            
            response = requests.get(url, params=params, timeout=10)
            if response.status_code == 200:
                data = response.json()
                
                if data.get('response_code') == 1:  # File found
                    positives = data.get('positives', 0)
                    total = data.get('total', 0)
                    
                    return {
                        'virustotal_detected': positives > 0,
                        'virustotal_engines': positives,
                        'virustotal_total': total,
                        'virustotal_scan_date': data.get('scan_date')
                    }
        except Exception as e:
            logger.error(f"Error in VirusTotal lookup: {str(e)}")
        
        return None

    def _sandbox_analysis(self, file_path):
        """Perform sandbox analysis on executable files"""
        result = {
            'features': [],
            'attack_vectors': []
        }
        
        try:
            # Create temporary sandbox directory
            with tempfile.TemporaryDirectory() as sandbox_dir:
                # Copy file to sandbox
                sandbox_file = os.path.join(sandbox_dir, os.path.basename(file_path))
                shutil.copy2(file_path, sandbox_file)
                
                # Make file executable
                os.chmod(sandbox_file, 0o755)
                
                # Run in sandbox with timeout
                try:
                    process = subprocess.Popen(
                        [sandbox_file],
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        cwd=sandbox_dir,
                        timeout=self.sandbox_timeout
                    )
                    
                    stdout, stderr = process.communicate()
                    
                    # Analyze behavior
                    if process.returncode != 0:
                        result['features'].append('execution_error')
                    
                    # Check for network activity (simplified)
                    if b'http' in stdout.lower() or b'https' in stdout.lower():
                        result['features'].append('network_activity')
                        result['attack_vectors'].append('network_communication')
                    
                    # Check for file system activity
                    if b'file' in stdout.lower() or b'directory' in stdout.lower():
                        result['features'].append('file_system_activity')
                        result['attack_vectors'].append('file_manipulation')
                    
                except subprocess.TimeoutExpired:
                    result['features'].append('execution_timeout')
                    result['attack_vectors'].append('resource_exhaustion')
                    process.kill()
                
        except Exception as e:
            logger.error(f"Error in sandbox analysis: {str(e)}")
        
        return result

    def _extract_context_words(self, file_info):
        """Extract context words from the webpage where file was found"""
        context_words = []
        
        try:
            page_url = file_info.get('page_url', '')
            link_text = file_info.get('link_text', '')
            
            # Extract words from URL
            url_words = re.findall(r'\b\w+\b', page_url.lower())
            context_words.extend(url_words)
            
            # Extract words from link text
            if link_text:
                link_words = re.findall(r'\b\w+\b', link_text.lower())
                context_words.extend(link_words)
            
            # Categorize and store context words
            categorized_words = self._categorize_context_words(context_words)
            
            # Store in database
            for word, category in categorized_words.items():
                existing_word = db_session.query(ContextWord).filter_by(word=word).first()
                if existing_word:
                    existing_word.frequency += 1
                    existing_word.last_used = db_session.query(db_session.func.now()).scalar()
                else:
                    new_word = ContextWord(word=word, category=category)
                    db_session.add(new_word)
            
            db_session.commit()
            
            return list(categorized_words.keys())
            
        except Exception as e:
            logger.error(f"Error extracting context words: {str(e)}")
            return []

    def _categorize_context_words(self, words):
        """Categorize context words based on content"""
        categories = {
            'business': ['company', 'corp', 'inc', 'ltd', 'business', 'enterprise'],
            'technology': ['software', 'app', 'system', 'tech', 'computer', 'digital'],
            'finance': ['bank', 'payment', 'money', 'financial', 'credit', 'account'],
            'education': ['school', 'university', 'college', 'student', 'education'],
            'government': ['gov', 'government', 'official', 'public', 'state'],
            'healthcare': ['medical', 'health', 'hospital', 'doctor', 'patient'],
            'entertainment': ['game', 'music', 'video', 'movie', 'entertainment']
        }
        
        categorized = {}
        for word in words:
            if len(word) < 3:  # Skip short words
                continue
                
            for category, keywords in categories.items():
                if word in keywords:
                    categorized[word] = category
                    break
            else:
                categorized[word] = 'general'
        
        return categorized

    def _classify_malware(self, analysis_result):
        """Classify malware based on analysis results"""
        features = analysis_result.get('features', [])
        attack_vectors = analysis_result.get('attack_vectors', [])
        
        # Simple classification logic
        if 'network_communication' in features and 'persistence' in attack_vectors:
            return "Backdoor/Trojan"
        elif 'encryption_obfuscation' in features:
            return "Ransomware"
        elif 'command_execution' in attack_vectors:
            return "Remote Access Tool"
        elif 'file_system_activity' in features:
            return "File Infector"
        else:
            return "Suspicious File"

    def _get_file_type(self, file_path):
        """Get file type using python-magic"""
        try:
            return magic.from_file(file_path)
        except:
            return "unknown"

    def _is_executable(self, file_path):
        """Check if file is executable"""
        file_type = self._get_file_type(file_path).lower()
        executable_indicators = ['executable', 'pe32', 'elf', 'mach-o']
        return any(indicator in file_type for indicator in executable_indicators)

    def _load_yara_rules(self):
        """Load YARA rules for malware detection"""
        try:
            rules_dir = Path('yara_rules')
            if rules_dir.exists():
                rules = yara.compile(str(rules_dir))
                return rules
        except Exception as e:
            logger.error(f"Error loading YARA rules: {str(e)}")
        return None

    def _load_known_malware(self):
        """Load known malware signatures"""
        try:
            # This would load from the database or external file
            return db_session.query(KnownMalware).all()
        except Exception as e:
            logger.error(f"Error loading known malware: {str(e)}")
        return []
