"""Vulnerability scanner module for CTF-H"""

import re
import os
from typing import List, Tuple
from ctfh.utils import (
    print_section, print_colored, Fore, Style, get_file_path, get_input,
    print_table, ask_save_to_file, export_to_json, export_to_csv, export_to_html,
    show_progress, add_to_history, get_input_with_example, validate_file_path, get_config
)
from ctfh.menu import Menu


# Vulnerability patterns
VULN_PATTERNS = [
    (r'eval\s*\(', 'eval() - Code injection risk'),
    (r'innerHTML\s*=', 'innerHTML - XSS risk'),
    (r'document\.write\s*\(', 'document.write() - XSS risk'),
    (r'shell\s*=\s*True', 'shell=True - Command injection risk'),
    (r'pickle\.loads\s*\(', 'pickle.loads() - Deserialization risk'),
    (r'os\.system\s*\(', 'os.system() - Command injection risk'),
    (r'exec\s*\(', 'exec() - Code injection risk'),
    (r'subprocess\.call\s*\(', 'subprocess.call() - Command injection risk'),
    (r'eval\s*\(', 'eval() - Code injection risk'),
    (r'system\s*\(', 'system() - Command injection risk'),
    (r'SQL.*\+.*\+', 'SQL string concatenation - SQL injection risk'),
    (r'\.innerHTML', 'innerHTML usage - XSS risk'),
    (r'\.outerHTML', 'outerHTML usage - XSS risk'),
]


def scan_file(file_path: str) -> List[Tuple[int, str, str]]:
    """Scan file for vulnerability patterns"""
    findings = []
    
    # Check if path exists
    if not os.path.exists(file_path):
        print_colored(f"Error: Path does not exist: {file_path}", Fore.RED)
        return findings
    
    # Check if it's a directory
    if os.path.isdir(file_path):
        print_colored(f"Error: Path is a directory, not a file: {file_path}", Fore.RED)
        print_colored("Tip: Use 'Scan Directory (Recursive)' option for directories", Fore.YELLOW)
        return findings
    
    # Check if it's a file
    if not os.path.isfile(file_path):
        print_colored(f"Error: Path is not a valid file: {file_path}", Fore.RED)
        return findings
    
    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            for line_num, line in enumerate(f, 1):
                for pattern, description in VULN_PATTERNS:
                    if re.search(pattern, line, re.IGNORECASE):
                        findings.append((line_num, line.strip(), description))
    except PermissionError:
        print_colored(f"Error: Permission denied. Cannot read file: {file_path}", Fore.RED)
        print_colored("Tip: Check file permissions or run with appropriate privileges", Fore.YELLOW)
    except Exception as e:
        print_colored(f"Error reading file: {e}", Fore.RED)
    
    return findings


def scan_directory(directory: str, extensions: List[str] = None) -> dict:
    """Scan directory for vulnerable patterns"""
    if extensions is None:
        extensions = ['.py', '.js', '.html', '.php', '.java', '.c', '.cpp']
    
    results = {}
    
    try:
        for root, dirs, files in os.walk(directory):
            for file in files:
                if any(file.endswith(ext) for ext in extensions):
                    file_path = os.path.join(root, file)
                    findings = scan_file(file_path)
                    if findings:
                        results[file_path] = findings
    except Exception as e:
        print_colored(f"Error scanning directory: {e}", Fore.RED)
    
    return results


def scan_single_file() -> None:
    """Scan a single file for vulnerabilities"""
    print_section("Vulnerability Scanner - Single File")
    
    file_path = get_input_with_example("Enter file path to scan", "C:\\Users\\Desktop\\script.py")
    
    if not file_path:
        print_colored("No file path provided.", Fore.RED)
        input("\nPress Enter to continue...")
        return
    
    # Validate path
    is_valid, error_msg = validate_file_path(file_path, must_exist=True, must_be_file=True)
    if not is_valid:
        print_colored(f"Error: {error_msg}", Fore.RED)
        if "directory" in error_msg.lower():
            print_colored("Tip: Use option 2 'Scan Directory (Recursive)' for directories", Fore.YELLOW)
        input("\nPress Enter to continue...")
        return
    
    print_colored(f"\nScanning: {file_path}\n", Fore.YELLOW)
    findings = scan_file(file_path)
    
    if findings:
        print_colored(f"Found {len(findings)} potential vulnerabilities:\n", Fore.RED)
        
        # Use table format
        table_rows = [[str(f[0]), f[1], f[2][:60]] for f in findings]
        print_table(["Line", "Code", "Vulnerability"], table_rows)
        
        # Add to history
        add_to_history("scan", {
            "file": file_path,
            "findings_count": len(findings),
            "findings": findings
        })
        
        # Export options
        export_data = {
            "file": file_path,
            "scan_date": str(os.path.getmtime(file_path)),
            "findings": [{"line": f[0], "code": f[1], "vulnerability": f[2]} for f in findings]
        }
        
        # Save as text
        text_output = f"Vulnerability Scan Results\n{'='*60}\nFile: {file_path}\nFindings: {len(findings)}\n\n"
        for line_num, line, desc in findings:
            text_output += f"Line {line_num}: {desc}\n  {line}\n\n"
        
        ask_save_to_file(text_output, "vulnerability_scan", "txt")
        
        # Export as JSON
        if get_input("Export as JSON? (y/n)", "n").lower() == 'y':
            json_file = get_input("Enter JSON filename", "vulnerability_scan.json")
            if json_file and export_to_json(export_data, json_file):
                print_colored(f"✓ Exported to {json_file}", Fore.GREEN)
        
        # Export as CSV
        if get_input("Export as CSV? (y/n)", "n").lower() == 'y':
            csv_file = get_input("Enter CSV filename", "vulnerability_scan.csv")
            if csv_file:
                csv_rows = [[str(f[0]), f[1], f[2]] for f in findings]
                if export_to_csv(["Line", "Code", "Vulnerability"], csv_rows, csv_file):
                    print_colored(f"✓ Exported to {csv_file}", Fore.GREEN)
    else:
        print_colored("No vulnerabilities detected.", Fore.GREEN)
        add_to_history("scan", {"file": file_path, "findings_count": 0})
    
    input("\nPress Enter to continue...")


def scan_directory_recursive() -> None:
    """Scan directory recursively for vulnerabilities"""
    print_section("Vulnerability Scanner - Directory")
    directory = get_input_with_example("Enter directory path to scan", "C:\\Users\\Desktop\\project")
    
    if not directory or not os.path.isdir(directory):
        print_colored("Invalid directory.", Fore.RED)
        input("\nPress Enter to continue...")
        return
    
    config = get_config()
    default_exts = ",".join(config.get("default_extensions", [".py", ".js", ".html"]))
    extensions_input = get_input("Enter file extensions to scan (comma-separated)", default_exts)
    extensions = [ext.strip() for ext in extensions_input.split(',')]
    
    print_colored(f"\nScanning directory: {directory}\n", Fore.YELLOW)
    print_colored("This may take a while...\n", Fore.CYAN)
    
    # Count files first for progress
    file_count = 0
    for root, dirs, files in os.walk(directory):
        for file in files:
            if any(file.endswith(ext) for ext in extensions):
                file_count += 1
    
    print_colored(f"Found {file_count} files to scan\n", Fore.CYAN)
    
    results = {}
    scanned = 0
    try:
        for root, dirs, files in os.walk(directory):
            for file in files:
                if any(file.endswith(ext) for ext in extensions):
                    file_path = os.path.join(root, file)
                    findings = scan_file(file_path)
                    scanned += 1
                    if findings:
                        results[file_path] = findings
                    
                    # Show progress
                    if config.get("show_progress", True):
                        show_progress(scanned, file_count, "Scanning")
    except KeyboardInterrupt:
        print_colored("\n\nScan interrupted by user.", Fore.YELLOW)
    
    print()  # New line after progress
    
    if results:
        total_findings = sum(len(v) for v in results.values())
        print_colored(f"\nFound {total_findings} potential vulnerabilities in {len(results)} files:\n", Fore.RED)
        
        # Summary table
        summary_rows = [[os.path.basename(path), str(len(findings))] for path, findings in results.items()]
        print_table(["File", "Vulnerabilities"], summary_rows[:20])  # Show first 20
        
        if len(results) > 20:
            print_colored(f"\n... and {len(results) - 20} more files with vulnerabilities", Fore.YELLOW)
        
        # Export options
        export_data = {
            "directory": directory,
            "extensions": extensions,
            "files_scanned": scanned,
            "files_with_vulns": len(results),
            "total_findings": total_findings,
            "results": {path: [{"line": f[0], "code": f[1], "vulnerability": f[2]} for f in findings] 
                       for path, findings in results.items()}
        }
        
        add_to_history("scan", {"directory": directory, "findings_count": total_findings, "files": len(results)})
        
        # Save results
        text_output = f"Vulnerability Scan Results - Directory\n{'='*60}\nDirectory: {directory}\nFiles with vulnerabilities: {len(results)}\nTotal findings: {total_findings}\n\n"
        for file_path, findings in results.items():
            text_output += f"\n{file_path}:\n"
            for line_num, line, desc in findings:
                text_output += f"  Line {line_num}: {desc}\n    {line}\n"
        
        ask_save_to_file(text_output, "directory_scan_results", "txt")
        
        # Export formats
        if get_input("\nExport as JSON? (y/n)", "n").lower() == 'y':
            json_file = get_input("Enter JSON filename", "scan_results.json")
            if json_file and export_to_json(export_data, json_file):
                print_colored(f"✓ Exported to {json_file}", Fore.GREEN)
        
        if get_input("Export as HTML report? (y/n)", "n").lower() == 'y':
            html_file = get_input("Enter HTML filename", "scan_report.html")
            if html_file and export_to_html("Vulnerability Scan Report", text_output, html_file):
                print_colored(f"✓ Exported to {html_file}", Fore.GREEN)
    else:
        print_colored(f"No vulnerabilities detected in {scanned} files.", Fore.GREEN)
    
    input("\nPress Enter to continue...")


def scan_text_input() -> None:
    """Scan user-provided text for vulnerabilities"""
    print_section("Vulnerability Scanner - Text Input")
    text = get_input("Enter code/text to scan (or 'file:path' to load from file)")
    
    if text.startswith('file:'):
        file_path = text[5:].strip()
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                lines = f.readlines()
        except Exception as e:
            print_colored(f"Error reading file: {e}", Fore.RED)
            input("\nPress Enter to continue...")
            return
    else:
        lines = text.split('\n')
    
    findings = []
    for line_num, line in enumerate(lines, 1):
        for pattern, description in VULN_PATTERNS:
            if re.search(pattern, line, re.IGNORECASE):
                findings.append((line_num, line.strip(), description))
    
    if findings:
        print_colored(f"\nFound {len(findings)} potential vulnerabilities:\n", Fore.RED)
        for line_num, line, description in findings:
            print_colored(f"Line {line_num}: {description}", Fore.RED)
            print_colored(f"  {line[:80]}", Fore.CYAN)
    else:
        print_colored("No vulnerabilities detected.", Fore.GREEN)
    
    input("\nPress Enter to continue...")


def vulnerability_menu() -> None:
    """Vulnerability scanner module menu"""
    def handle_help():
        help_text = """
Vulnerability Scanner Module Help
=================================

This module scans code for common security vulnerabilities and dangerous patterns.

Scan Single File:
  - Scans one file for vulnerability patterns
  - Shows results in a formatted table
  - Can export results as TXT, JSON, or CSV
  - Example: C:\\Users\\Desktop\\script.py

Scan Directory (Recursive):
  - Scans all files in a directory and subdirectories
  - Shows progress bar during scanning
  - Can filter by file extensions (.py, .js, .html, etc.)
  - Exports comprehensive reports

Scan Text Input:
  - Paste code directly or load from file
  - Quick check without saving files
  - Useful for code snippets

Detected Patterns:
  - eval(), exec() - Code injection
  - innerHTML, document.write - XSS
  - shell=True, os.system - Command injection
  - pickle.loads - Deserialization
  - SQL string concatenation - SQL injection
  - And more...

Export Formats:
  - TXT: Human-readable text file
  - JSON: Structured data for automation
  - CSV: Spreadsheet-compatible format
  - HTML: Formatted report for sharing

Tips:
  - Use directory scan for full projects
  - Export JSON for automated analysis
  - Check file permissions if you get errors
        """
        print_colored(help_text, Fore.CYAN)
        input("\nPress Enter to continue...")
    
    options = [
        (1, "Scan Single File", scan_single_file),
        (2, "Scan Directory (Recursive)", scan_directory_recursive),
        (3, "Scan Text Input", scan_text_input),
        (99, "Help (?)", handle_help),
        (0, "Back to Main Menu", lambda: None),
    ]
    
    menu = Menu("Vulnerability Scanner Module", options)
    result = menu.run()
    # Return None to signal "go back to main menu" when submenu exits
    return None if result else False

