vulnerability_scanner.py

#!/usr/bin/env python3
"""
Vulnerability Prioritization and Remediation System

This is the main script that orchestrates the entire workflow:
1. Network scanning with Nmap
2. Processing scan results with AI models
3. Generating prioritized vulnerability reports
"""

import os
import sys
import argparse
from pathlib import Path
import subprocess
import time
from datetime import datetime

def check_dependencies():
    """Check if required dependencies are installed"""
    try:
        # Check for nmap
        nmap_check = subprocess.run(['nmap', '-V'], capture_output=True, text=True)
        if nmap_check.returncode != 0:
            print("[!] Error: Nmap is not installed or not in PATH.")
            print("    Please install Nmap: https://nmap.org/download.html")
            return False
        
        # Check for python dependencies
        try:
            import nmap
            import pandas as pd
            import matplotlib.pyplot as plt
        except ImportError as e:
            print(f"[!] Error: Missing Python dependency: {e}")
            print("    Please install required packages: pip install python-nmap pandas matplotlib")
            return False
        
        return True
    except Exception as e:
        print(f"[!] Error checking dependencies: {e}")
        return False

def parse_args():
    """Parse command-line arguments"""
    parser = argparse.ArgumentParser(
        description="AI-Driven Vulnerability Prioritization and Remediation System",
        formatter_class=argparse.RawTextHelpFormatter
    )
    
    # Target specification
    parser.add_argument('target', help='Target IP address, hostname, or network (e.g., 192.168.1.1 or 192.168.1.0/24)')
    
    # Scan options
    parser.add_argument('-p', '--ports', help='Port specification (e.g., 22,80,443 or 1-1000)')
    parser.add_argument('--quick', action='store_true', help='Perform a quick scan (less comprehensive)')
    
    # Output options
    parser.add_argument('-o', '--output', help='Base name for output files')
    parser.add_argument('--show', action='store_true', help='Open the HTML report when completed')
    
    return parser.parse_args()

def run_scan(target, ports=None, quick=False):
    """Run network scan using the network_scanner module"""
    from network_scanner import NetworkScanner
    
    scanner = NetworkScanner()
    print(f"[*] Starting {'quick' if quick else 'comprehensive'} scan of {target}...")
    
    if quick:
        scan_result = scanner.basic_scan(target, ports)
    else:
        scan_result = scanner.full_scan(target, ports)
    
    if not scan_result:
        print("[!] Scan failed or returned no results")
        return None
    
    parsed_results = scanner.parse_results()
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    results_file = scanner.save_results(parsed_results, target)
    
    ai_data = scanner.export_for_ai(parsed_results)
    if not ai_data:
        print("[!] No vulnerabilities found or failed to parse results")
        return None
    
    ai_data_file = Path("scan_results") / f"ai_input_{target.replace('/', '_')}_{timestamp}.json"
    with open(ai_data_file, 'w') as f:
        import json
        json.dump(ai_data, f, indent=2)
    
    print(f"[+] Scan completed. Found {len(ai_data)} potential vulnerabilities.")
    return ai_data_file

def process_with_ai(scan_results_file):
    """Process scan results with AI models using the integration module"""
    from integration import VulnerabilityProcessor
    
    print("[*] Processing scan results with AI models...")
    processor = VulnerabilityProcessor()
    analyzed_results = processor.process_scan_results(scan_results_file)
    
    if not analyzed_results:
        print("[!] AI processing failed or found no vulnerabilities")
        return None
    
    output_file = scan_results_file.parent / f"analyzed_{scan_results_file.name}"
    with open(output_file, 'w') as f:
        import json
        json.dump(analyzed_results, f, indent=2)
    
    print(f"[+] AI analysis completed. Results saved to {output_file}")
    
    # Count by severity
    severity_counts = {}
    for result in analyzed_results:
        severity = result.get('severity', 'Unknown')
        severity_counts[severity] = severity_counts.get(severity, 0) + 1
    
    # Print summary
    print("[+] Vulnerability summary:")
    for severity in ["Critical", "High", "Medium", "Low", "Unknown"]:
        if severity in severity_counts:
            print(f"  - {severity}: {severity_counts[severity]}")
    
    return output_file

def generate_reports(analyzed_file, output_base=None, show_report=False):
    """Generate reports using the report_generator module"""
    from report_generator import VulnerabilityReporter
    
    print("[*] Generating vulnerability reports...")
    
    try:
        with open(analyzed_file, 'r') as f:
            import json
            analyzed_data = json.load(f)
        
        reporter = VulnerabilityReporter()
        
        if output_base:
            reports_dir = Path("reports")
            reports_dir.mkdir(exist_ok=True)
            csv_output = reports_dir / f"{output_base}.csv"
            html_output = reports_dir / f"{output_base}.html"
        else:
            csv_output = None
            html_output = None
        
        # Generate CSV report
        csv_report = reporter.generate_csv_report(analyzed_data, csv_output)
        
        # Generate HTML report
        html_report = reporter.generate_html_report(analyzed_data, html_output)
        
        if show_report:
            reporter.open_report(html_report)
        
        return html_report
    except Exception as e:
        print(f"[!] Error generating reports: {e}")
        return None

def main():
    """Main execution function"""
    print("=" * 80)
    print("  AI-Driven Vulnerability Prioritization and Remediation System")
    print("=" * 80)
    
    # Check dependencies
    if not check_dependencies():
        sys.exit(1)
    
    # Parse arguments
    args = parse_args()
    
    # Generate output base name if not provided
    output_base = args.output
    if not output_base:
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        target_clean = args.target.replace('/', '_').replace('.', '_')
        output_base = f"vuln_scan_{target_clean}_{timestamp}"
    
    try:
        # Step 1: Run network scan
        scan_results_file = run_scan(args.target, args.ports, args.quick)
        if not scan_results_file:
            print("[!] Scan produced no usable results. Exiting.")
            sys.exit(1)
        
        # Step 2: Process with AI
        analyzed_file = process_with_ai(scan_results_file)
        if not analyzed_file:
            print("[!] AI analysis produced no usable results. Exiting.")
            sys.exit(1)
        
        # Step 3: Generate reports
        report_file = generate_reports(analyzed_file, output_base, args.show)
        
        print("\n" + "=" * 80)
        if report_file:
            print(f"[+] Vulnerability assessment complete! Report available at: {report_file}")
        else:
            print("[!] Vulnerability assessment completed with errors.")
        print("=" * 80)
    
    except KeyboardInterrupt:
        print("\n[!] Process interrupted by user. Exiting.")
        sys.exit(1)
    except Exception as e:
        print(f"\n[!] An error occurred: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()