cli.py

import os
import math
import hashlib
import time
import argparse
import logging
import ast
import re
from datetime import datetime
from typing import List, Dict, Optional, Set
from collections import Counter

class CodeAnalyzer:
    def __init__(self):
        # Suspicious imports that might indicate malicious behavior
        self.suspicious_imports = {
            'cryptography.fernet', 'pycrypto', 'cryptography', 
            'requests', 'socket', 'subprocess', 'winreg',
            'ctypes', 'pyHook', 'pythoncom'
        }
        
        # Suspicious function names/patterns
        self.suspicious_functions = {
            'encrypt', 'decrypt', 'ransom', 'payload',
            'shellcode', 'exploit', 'inject', 'hook'
        }
        
        # Suspicious string patterns
        self.suspicious_strings = [
            r'\.encrypt\(', r'\.decrypt\(', 
            r'bitcoin', r'wallet', r'ransom',
            r'delete.*shadow', r'taskkill',
            r'process.*kill', r'registry.*delete'
        ]

    def analyze_python_code(self, filepath: str) -> Dict:
        """Analyze Python source code for suspicious patterns."""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                content = f.read()
                
            result = {
                'suspicious_imports': set(),
                'suspicious_functions': set(),
                'suspicious_patterns': set(),
                'uses_encryption': False,
                'file_operations': False,
                'system_operations': False
            }

            # Parse and analyze the AST
            tree = ast.parse(content)
            
            # Analyze imports
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for name in node.names:
                        if any(susp in name.name for susp in self.suspicious_imports):
                            result['suspicious_imports'].add(name.name)
                
                elif isinstance(node, ast.ImportFrom):
                    module = node.module or ''
                    if any(susp in module for susp in self.suspicious_imports):
                        result['suspicious_imports'].add(module)
                
                # Check for suspicious function definitions
                elif isinstance(node, ast.FunctionDef):
                    if any(susp in node.name.lower() for susp in self.suspicious_functions):
                        result['suspicious_functions'].add(node.name)
                
                # Check for file operations
                elif isinstance(node, ast.Call):
                    if hasattr(node.func, 'id'):
                        if node.func.id in {'open', 'write', 'remove', 'unlink'}:
                            result['file_operations'] = True
                    elif hasattr(node.func, 'attr'):
                        if node.func.attr in {'encrypt', 'decrypt'}:
                            result['uses_encryption'] = True

            # Check for suspicious string patterns
            for pattern in self.suspicious_strings:
                if re.search(pattern, content, re.IGNORECASE):
                    result['suspicious_patterns'].add(pattern)

            return result
            
        except Exception as e:
            logging.error(f"Error analyzing Python code in {filepath}: {str(e)}")
            return {}

class FileScanner:
    def __init__(self, path: str):
        self.path = path
        self.suspicious_extensions = {
            '.encrypted', '.crypto', '.locked', '.crypted',
            '.cry', '.crown', '.crypt', '.beast',
            '.wannacry', '.wcry', '.wncry', '.tesla',
            '.py', '.pyc', '.pyw'  # Added Python extensions
        }
        self.known_malware_hashes = {
            # Example hashes of known malware/ransomware signatures
            "a123b456c789": "WannaCry",
            "d456e789f123": "Petya",
        }
        self.code_analyzer = CodeAnalyzer()
    
    def calculate_file_hash(self, filepath: str) -> str:
        """Calculate SHA-256 hash of a file."""
        sha256_hash = hashlib.sha256()
        try:
            with open(filepath, "rb") as f:
                for byte_block in iter(lambda: f.read(4096), b""):
                    sha256_hash.update(byte_block)
            return sha256_hash.hexdigest()
        except Exception as e:
            logging.error(f"Error calculating hash for {filepath}: {str(e)}")
            return ""

    def check_file_entropy(self, filepath: str) -> float:
        """Calculate Shannon entropy of file content."""
        try:
            with open(filepath, "rb") as f:
                data = f.read()
                if not data:
                    return 0
                
                entropy = 0
                for x in range(256):
                    p_x = data.count(bytes([x])) / len(data)
                    if p_x > 0:
                        entropy += - p_x * math.log2(p_x)
                return entropy
        except Exception as e:
            logging.error(f"Error calculating entropy for {filepath}: {str(e)}")
            return 0

    def is_suspicious_extension(self, filepath: str) -> bool:
        """Check if file has suspicious extension."""
        return any(filepath.lower().endswith(ext) for ext in self.suspicious_extensions)

    def analyze_file_content(self, filepath: str) -> Dict:
        """Analyze file content for suspicious patterns."""
        try:
            with open(filepath, 'rb') as f:
                content = f.read()
                
            result = {
                'has_executable_code': False,
                'has_base64': False,
                'has_encrypted_content': False
            }
            
            # Check for executable content
            executable_patterns = [b'MZ', b'PE\x00\x00', b'#!/', b'<?php']
            result['has_executable_code'] = any(pattern in content[:100] for pattern in executable_patterns)
            
            # Check for Base64 encoded content
            base64_pattern = rb'^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$'
            result['has_base64'] = bool(re.search(base64_pattern, content))
            
            # Check for potential encrypted content
            result['has_encrypted_content'] = self.check_file_entropy(filepath) > 7.0
            
            return result
            
        except Exception as e:
            logging.error(f"Error analyzing file content for {filepath}: {str(e)}")
            return {}

    def scan_file(self, filepath: str) -> Dict:
        """Scan individual file for malware indicators."""
        result = {
            "filepath": filepath,
            "suspicious": False,
            "reasons": [],
            "hash": "",
            "entropy": 0,
            "code_analysis": {},
            "content_analysis": {}
        }

        # Basic checks
        if self.is_suspicious_extension(filepath):
            result["suspicious"] = True
            result["reasons"].append("Suspicious file extension")

        file_hash = self.calculate_file_hash(filepath)
        result["hash"] = file_hash

        if file_hash in self.known_malware_hashes:
            result["suspicious"] = True
            result["reasons"].append(f"Matched known malware: {self.known_malware_hashes[file_hash]}")

        entropy = self.check_file_entropy(filepath)
        result["entropy"] = entropy
        
        if entropy > 7.5:
            result["suspicious"] = True
            result["reasons"].append("High entropy - possible encryption")

        # Advanced content analysis
        content_analysis = self.analyze_file_content(filepath)
        result["content_analysis"] = content_analysis
        
        if content_analysis.get('has_executable_code'):
            result["suspicious"] = True
            result["reasons"].append("Contains executable code")
        
        if content_analysis.get('has_base64'):
            result["suspicious"] = True
            result["reasons"].append("Contains Base64 encoded content")

        # Python-specific analysis
        if filepath.lower().endswith(('.py', '.pyw')):
            code_analysis = self.code_analyzer.analyze_python_code(filepath)
            result["code_analysis"] = code_analysis
            
            if code_analysis.get('suspicious_imports'):
                result["suspicious"] = True
                result["reasons"].append(f"Suspicious imports: {', '.join(code_analysis['suspicious_imports'])}")
            
            if code_analysis.get('suspicious_functions'):
                result["suspicious"] = True
                result["reasons"].append(f"Suspicious functions: {', '.join(code_analysis['suspicious_functions'])}")
            
            if code_analysis.get('suspicious_patterns'):
                result["suspicious"] = True
                result["reasons"].append(f"Suspicious code patterns detected")
            
            if code_analysis.get('uses_encryption') and code_analysis.get('file_operations'):
                result["suspicious"] = True
                result["reasons"].append("Combines encryption with file operations - possible ransomware")

        return result

class SystemScanner:
    def __init__(self):
        self.file_scanner = None
        self.scan_results = []
        self.start_time = None
        self.end_time = None
        self.total_files_scanned = 0
        self.suspicious_files_count = 0

    def scan_directory(self, path: str, recursive: bool = True) -> List[Dict]:
        """Scan directory for suspicious files."""
        self.file_scanner = FileScanner(path)
        self.scan_results = []
        self.start_time = datetime.now()
        self.total_files_scanned = 0
        self.suspicious_files_count = 0

        try:
            if recursive:
                for root, _, files in os.walk(path):
                    for file in files:
                        filepath = os.path.join(root, file)
                        self.total_files_scanned += 1
                        result = self.file_scanner.scan_file(filepath)
                        if result["suspicious"]:
                            self.suspicious_files_count += 1
                            self.scan_results.append(result)
            else:
                for file in os.listdir(path):
                    filepath = os.path.join(path, file)
                    if os.path.isfile(filepath):
                        self.total_files_scanned += 1
                        result = self.file_scanner.scan_file(filepath)
                        if result["suspicious"]:
                            self.suspicious_files_count += 1
                            self.scan_results.append(result)

        except Exception as e:
            logging.error(f"Error scanning directory {path}: {str(e)}")

        self.end_time = datetime.now()
        return self.scan_results

    def generate_report(self) -> str:
        """Generate detailed scan report."""
        report = []
        report.append("=== Enhanced Malware Scanner Report ===")
        report.append(f"Scan started: {self.start_time}")
        report.append(f"Scan completed: {self.end_time}")
        report.append(f"Duration: {self.end_time - self.start_time}")
        report.append(f"Total files scanned: {self.total_files_scanned}")
        report.append(f"Suspicious files found: {self.suspicious_files_count}")
        
        if self.suspicious_files_count > 0:
            report.append("\nDetailed Analysis of Suspicious Files:")
            
            for result in self.scan_results:
                report.append("\n" + "="*50)
                report.append(f"\nFile: {result['filepath']}")
                report.append(f"Hash: {result['hash']}")
                report.append(f"Entropy: {result['entropy']:.2f}")
                
                report.append("\nReasons for suspicion:")
                for reason in result['reasons']:
                    report.append(f"  - {reason}")
                
                if result.get('code_analysis'):
                    report.append("\nCode Analysis:")
                    code_analysis = result['code_analysis']
                    if code_analysis.get('suspicious_imports'):
                        report.append(f"  - Suspicious imports: {', '.join(code_analysis['suspicious_imports'])}")
                    if code_analysis.get('suspicious_functions'):
                        report.append(f"  - Suspicious functions: {', '.join(code_analysis['suspicious_functions'])}")
                    if code_analysis.get('suspicious_patterns'):
                        report.append(f"  - Suspicious patterns detected: {', '.join(code_analysis['suspicious_patterns'])}")
                
                if result.get('content_analysis'):
                    report.append("\nContent Analysis:")
                    content_analysis = result['content_analysis']
                    if content_analysis.get('has_executable_code'):
                        report.append("  - Contains executable code")
                    if content_analysis.get('has_base64'):
                        report.append("  - Contains Base64 encoded content")
                    if content_analysis.get('has_encrypted_content'):
                        report.append("  - Contains potentially encrypted content")

        return "\n".join(report)

def main():
    parser = argparse.ArgumentParser(description="Enhanced Malware and Ransomware Detection Tool")
    parser.add_argument("path", help="Path to scan")
    parser.add_argument("-r", "--recursive", action="store_true", help="Scan directories recursively")
    parser.add_argument("-o", "--output", help="Output file for report")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose logging")
    args = parser.parse_args()

    # Setup logging
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format='%(asctime)s - %(levelname)s - %(message)s')

    # Initialize and run scanner
    scanner = SystemScanner()
    logging.info(f"Starting scan of {args.path}")
    
    results = scanner.scan_directory(args.path, args.recursive)
    report = scanner.generate_report()

    # Output results
    if args.output:
        with open(args.output, 'w') as f:
            f.write(report)
        print(f"Report written to {args.output}")
    else:
        print(report)

if __name__ == "__main__":
    main()