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

import re
import os
from typing import List, Tuple
from ctfh.utils import print_section, print_colored, Fore, get_file_path, get_input
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 = []
    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 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_file_path("Enter file path to scan")
    
    if not file_path:
        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)
        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 scan_directory_recursive() -> None:
    """Scan directory recursively for vulnerabilities"""
    print_section("Vulnerability Scanner - Directory")
    directory = get_input("Enter directory path to scan")
    
    if not directory or not os.path.isdir(directory):
        print_colored("Invalid directory.", Fore.RED)
        input("\nPress Enter to continue...")
        return
    
    extensions_input = get_input("Enter file extensions to scan (comma-separated, e.g., .py,.js,.html)", ".py,.js,.html")
    extensions = [ext.strip() for ext in extensions_input.split(',')]
    
    print_colored(f"\nScanning directory: {directory}\n", Fore.YELLOW)
    results = scan_directory(directory, extensions)
    
    if results:
        total_findings = sum(len(v) for v in results.values())
        print_colored(f"Found {total_findings} potential vulnerabilities in {len(results)} files:\n", Fore.RED)
        
        for file_path, findings in results.items():
            print_colored(f"\n{file_path}:", Fore.YELLOW)
            for line_num, line, description in findings:
                print_colored(f"  Line {line_num}: {description}", Fore.RED)
                print_colored(f"    {line[:60]}", Fore.CYAN)
    else:
        print_colored("No vulnerabilities detected.", 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"""
    options = [
        (1, "Scan Single File", scan_single_file),
        (2, "Scan Directory (Recursive)", scan_directory_recursive),
        (3, "Scan Text Input", scan_text_input),
        (0, "Back to Main Menu", lambda: None),
    ]
    
    menu = Menu("Vulnerability Scanner Module", options)
    menu.run()

