#!/usr/bin/env python3
"""
SQLMap CLI - A beautiful CLI wrapper for sqlmap
Automates comprehensive SQL injection testing with a single command
"""

import subprocess
import sys
import argparse
from pathlib import Path
from typing import List, Dict, Tuple, Optional, TypedDict, Any
from datetime import datetime

try:
    from rich.console import Console
    from rich.panel import Panel
    from rich.table import Table
    from rich.progress import (
        Progress,
        SpinnerColumn,
        TextColumn,
        BarColumn,
        TimeElapsedColumn,
    )
    from rich.prompt import Prompt, Confirm
    from rich import box
except ImportError:
    print("Error: 'rich' library is required. Install it with: pip install rich")
    sys.exit(1)

console = Console()

SQLMAP_PATH = Path(__file__).parent / "sqlmap.py"

SQL_TECHNIQUES = {
    "B": "Boolean-based blind",
    "E": "Error-based",
    "U": "Union query-based",
    "S": "Stacked queries",
    "T": "Time-based blind",
    "Q": "Inline queries",
}


class ScanResult(TypedDict):
    total_tests: int
    vulnerabilities: List[Dict[str, str]]
    start_time: Optional[datetime]
    end_time: Optional[datetime]
    target: Optional[str]


class SQLMapCLI:
    def __init__(self):
        self.console = Console()
        self.results: ScanResult = {
            "total_tests": 0,
            "vulnerabilities": [],
            "start_time": None,
            "end_time": None,
            "target": None,
        }

    def print_banner(self):
        """Display a beautiful banner"""
        banner = """
╔═══════════════════════════════════════════════════════════════╗
║                                                               ║
║   ███████╗ ██████╗ ██╗     ███╗   ███╗ █████╗ ██████╗         ║
║   ██╔════╝██╔═══██╗██║     ████╗ ████║██╔══██╗██╔══██╗        ║
║   ███████╗██║   ██║██║     ██╔████╔██║███████║██████╔╝        ║
║   ╚════██║██║▄▄ ██║██║     ██║╚██╔╝██║██╔══██║██╔═══╝         ║
║   ███████║╚██████╔╝███████╗██║ ╚═╝ ██║██║  ██║██║             ║
║   ╚══════╝ ╚══▀▀═╝ ╚══════╝╚═╝     ╚═╝╚═╝  ╚═╝╚═╝             ║
║                                                               ║
║              CLI - Automated SQL Injection Testing            ║
║                                                               ║
╚═══════════════════════════════════════════════════════════════╝
        """
        self.console.print(banner, style="bold cyan")
        self.console.print(
            Panel(
                "[yellow]⚠️  Legal Disclaimer: Only use on targets you have permission to test[/yellow]",
                border_style="yellow",
                box=box.ROUNDED,
            )
        )
        self.console.print()

    def run_sqlmap_test(
        self,
        url: str,
        level: int,
        risk: int,
        technique: str = "BEUSTQ",
        batch: bool = True,
        data: Optional[str] = None,
        verbose: int = 1,
        extra_args: Optional[List[str]] = None,
    ) -> Tuple[bool, str]:
        """Run sqlmap with specified parameters"""
        cmd = [
            sys.executable,
            str(SQLMAP_PATH),
            "-u",
            url,
            f"--level={level}",
            f"--risk={risk}",
            f"--technique={technique}",
            "-v",
            str(verbose),
        ]

        if batch:
            cmd.append("--batch")

        if data:
            cmd.extend(["--data", data, "--method", "POST"])

        if extra_args:
            cmd.extend(extra_args)

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=600,  # 10 minute timeout per test
            )
            return result.returncode == 0, result.stdout + result.stderr
        except subprocess.TimeoutExpired:
            return False, "Test timed out after 10 minutes"
        except Exception as e:
            return False, str(e)

    def parse_results(self, output: str) -> Dict[str, Any]:
        """Parse sqlmap output for vulnerabilities"""
        vulns = []

        # Look for vulnerability indicators
        if "sqlmap identified the following injection point" in output:
            # Extract injection details
            lines = output.split("\n")
            current_param = "Unknown"  # Default parameter name

            for i, line in enumerate(lines):
                if "Parameter:" in line:
                    current_param = line.split("Parameter:")[1].strip()
                elif "Type:" in line:
                    vuln_type = line.split("Type:")[1].strip()
                    # Check if next line contains the title
                    if i + 1 < len(lines) and "Title:" in lines[i + 1]:
                        title = lines[i + 1].split("Title:")[1].strip()
                        vulns.append(
                            {
                                "parameter": current_param,
                                "type": vuln_type,
                                "title": title,
                            }
                        )

        # Check for backend DBMS detection
        backend_dbms = None
        if "back-end DBMS:" in output.lower():
            for line in output.split("\n"):
                if "back-end DBMS:" in line.lower():
                    backend_dbms = line.split(":", 1)[1].strip()
                    break

        return {
            "vulnerabilities": vulns,
            "backend_dbms": backend_dbms,
            "is_vulnerable": len(vulns) > 0 or "vulnerable" in output.lower(),
        }

    def comprehensive_scan(
        self,
        url: str,
        max_level: int = 5,
        max_risk: int = 3,
        techniques: str = "BEUSTQ",
        data: Optional[str] = None,
        verbose: int = 1,
    ):
        """Run comprehensive scan with all levels and risks"""
        self.results["target"] = url
        self.results["start_time"] = datetime.now()

        # Create results table
        results_table = Table(title="Scan Results", box=box.ROUNDED)
        results_table.add_column("Level", style="cyan", justify="center")
        results_table.add_column("Risk", style="yellow", justify="center")
        results_table.add_column("Status", justify="center")
        results_table.add_column("Findings", style="magenta")

        total_tests = max_level * max_risk

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeElapsedColumn(),
            console=self.console,
        ) as progress:
            overall_task = progress.add_task(
                f"[cyan]Scanning {url}...", total=total_tests
            )

            for level in range(1, max_level + 1):
                for risk in range(1, max_risk + 1):
                    progress.update(
                        overall_task,
                        description=f"[cyan]Testing Level {level}, Risk {risk}...",
                    )

                    success, output = self.run_sqlmap_test(
                        url, level, risk, techniques, data=data, verbose=verbose
                    )
                    parsed = self.parse_results(output)

                    status = "✓" if success else "✗"
                    status_style = "green" if success else "red"

                    findings = (
                        "No vulnerabilities"
                        if not parsed["is_vulnerable"]
                        else f"{len(parsed['vulnerabilities'])} found!"
                    )
                    findings_style = (
                        "green" if not parsed["is_vulnerable"] else "bold red"
                    )

                    if parsed["is_vulnerable"]:
                        self.results["vulnerabilities"].extend(
                            parsed["vulnerabilities"]
                        )

                    results_table.add_row(
                        str(level),
                        str(risk),
                        f"[{status_style}]{status}[/{status_style}]",
                        f"[{findings_style}]{findings}[/{findings_style}]",
                    )

                    progress.update(overall_task, advance=1)
                    self.results["total_tests"] += 1

        self.results["end_time"] = datetime.now()

        # Display results
        self.console.print()
        self.console.print(results_table)
        self.display_summary()

    def quick_scan(
        self,
        url: str,
        level: int = 1,
        risk: int = 1,
        data: Optional[str] = None,
        raw: bool = False,
        verbose: int = 1,
    ):
        """Run a quick scan with default settings"""
        self.results["target"] = url
        self.results["start_time"] = datetime.now()

        if not raw:
            scan_info = f"[cyan]Running quick scan on:[/cyan]\n[yellow]{url}[/yellow]\n[dim]Level: {level}, Risk: {risk}[/dim]"
            if data:
                scan_info += f"\n[dim]POST Data: {data}[/dim]"

            self.console.print(Panel(scan_info, border_style="cyan", box=box.ROUNDED))

        if raw:
            # Raw mode - just show sqlmap output directly
            self.console.print("[cyan]Running sqlmap...[/cyan]\n")
            success, output = self.run_sqlmap_test(
                url, level, risk, data=data, verbose=verbose
            )
            self.console.print(output)
            return

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            TimeElapsedColumn(),
            console=self.console,
        ) as progress:
            task = progress.add_task(
                "[cyan]Scanning for vulnerabilities...", total=None
            )
            success, output = self.run_sqlmap_test(
                url, level, risk, data=data, verbose=verbose
            )
            progress.update(task, completed=True)

        parsed = self.parse_results(output)
        self.results["vulnerabilities"] = parsed["vulnerabilities"]
        self.results["total_tests"] = 1
        self.results["end_time"] = datetime.now()

        self.display_summary()

    def display_summary(self):
        """Display a comprehensive summary of results"""
        self.console.print()

        # Calculate duration
        duration = 0.0
        if self.results["end_time"] and self.results["start_time"]:
            duration = (
                self.results["end_time"] - self.results["start_time"]
            ).total_seconds()

        # Create summary panel
        summary_text = f"""
[cyan]Target:[/cyan] {self.results["target"] or "N/A"}
[cyan]Total Tests:[/cyan] {self.results["total_tests"]}
[cyan]Duration:[/cyan] {duration:.2f} seconds
[cyan]Vulnerabilities Found:[/cyan] {len(self.results["vulnerabilities"])}
        """

        self.console.print(
            Panel(
                summary_text.strip(),
                title="[bold]Scan Summary[/bold]",
                border_style="green"
                if len(self.results["vulnerabilities"]) == 0
                else "red",
                box=box.DOUBLE,
            )
        )

        # Display vulnerabilities if found
        if self.results["vulnerabilities"]:
            self.console.print()
            vuln_table = Table(title="⚠️  Vulnerabilities Detected", box=box.HEAVY)
            vuln_table.add_column("Parameter", style="cyan")
            vuln_table.add_column("Type", style="yellow")
            vuln_table.add_column("Title", style="red")

            for vuln in self.results["vulnerabilities"]:
                vuln_table.add_row(
                    vuln.get("parameter", "N/A"),
                    vuln.get("type", "N/A"),
                    vuln.get("title", "N/A"),
                )

            self.console.print(vuln_table)
            self.console.print()
            self.console.print(
                "[bold red]⚠️  SQL injection vulnerabilities detected! Take immediate action.[/bold red]"
            )
        else:
            self.console.print()
            self.console.print(
                "[bold green]✓ No SQL injection vulnerabilities detected.[/bold green]"
            )

        self.console.print()

    def interactive_mode(self):
        """Interactive mode for user input"""
        self.console.print()
        self.console.print(
            Panel(
                "[cyan]Interactive Mode[/cyan]\n[dim]Enter target details for SQL injection testing[/dim]",
                border_style="cyan",
            )
        )

        url = Prompt.ask("\n[cyan]Enter target URL[/cyan]")

        # Ask if this is a POST request
        has_data = Confirm.ask(
            "[cyan]Does this request require POST data/body?[/cyan]", default=False
        )

        data = None
        if has_data:
            self.console.print("\n[dim]Examples:[/dim]")
            self.console.print(
                '[dim]  JSON: {"email":"test@example.com","password":"pass123"}[/dim]'
            )
            self.console.print("[dim]  Form: username=admin&password=secret[/dim]")
            data = Prompt.ask("\n[cyan]Enter POST data/body[/cyan]")

        scan_type = Prompt.ask(
            "\n[cyan]Select scan type[/cyan]",
            choices=["quick", "comprehensive"],
            default="quick",
        )

        if scan_type == "quick":
            level = int(Prompt.ask("[cyan]Test level (1-5)[/cyan]", default="1"))
            risk = int(Prompt.ask("[cyan]Test risk (1-3)[/cyan]", default="1"))
            self.quick_scan(url, level, risk, data=data)
        else:
            max_level = int(
                Prompt.ask("[cyan]Maximum test level (1-5)[/cyan]", default="5")
            )
            max_risk = int(
                Prompt.ask("[cyan]Maximum test risk (1-3)[/cyan]", default="3")
            )
            self.comprehensive_scan(url, max_level, max_risk, data=data)


def main():
    parser = argparse.ArgumentParser(
        description="SQLMap CLI - Beautiful automated SQL injection testing",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Quick scan with default settings (GET parameter)
  python sqlmapcli.py -u "https://demo.owasp-juice.shop/rest/products/search?q=test"
  
  # Test with POST data (JSON)
  python sqlmapcli.py -u "https://demo.owasp-juice.shop/rest/user/login" --data='{"email":"test@example.com","password":"pass123"}'
  
  # Comprehensive scan (all risk and level combinations)
  python sqlmapcli.py -u "https://demo.owasp-juice.shop/rest/products/search?q=test" --comprehensive
  
  # Custom level and risk with POST data
  python sqlmapcli.py -u "https://demo.owasp-juice.shop/rest/user/login" --data='{"email":"test@example.com","password":"pass123"}' --level 3 --risk 2
  
  # Interactive mode
  python sqlmapcli.py --interactive
        """,
    )

    parser.add_argument(
        "-u", "--url", help='Target URL (e.g., "http://example.com/page?id=1")'
    )

    parser.add_argument(
        "--comprehensive",
        action="store_true",
        help="Run comprehensive scan with all risk/level combinations",
    )

    parser.add_argument(
        "--level",
        type=int,
        default=1,
        choices=[1, 2, 3, 4, 5],
        help="Level of tests to perform (1-5, default: 1)",
    )

    parser.add_argument(
        "--risk",
        type=int,
        default=1,
        choices=[1, 2, 3],
        help="Risk of tests to perform (1-3, default: 1)",
    )

    parser.add_argument(
        "--max-level",
        type=int,
        default=5,
        choices=[1, 2, 3, 4, 5],
        help="Maximum level for comprehensive scan (default: 5)",
    )

    parser.add_argument(
        "--max-risk",
        type=int,
        default=3,
        choices=[1, 2, 3],
        help="Maximum risk for comprehensive scan (default: 3)",
    )

    parser.add_argument(
        "--technique",
        type=str,
        default="BEUSTQ",
        help="SQL injection techniques to use (default: BEUSTQ)",
    )

    parser.add_argument(
        "--data",
        type=str,
        help='Data string to be sent through POST (e.g., "username=test&password=test")',
    )

    parser.add_argument(
        "--raw", action="store_true", help="Show raw sqlmap output without formatting"
    )

    parser.add_argument(
        "--verbose",
        type=int,
        choices=[0, 1, 2, 3, 4, 5, 6],
        help="Sqlmap verbosity level (0-6, default: 1)",
    )

    parser.add_argument(
        "-i", "--interactive", action="store_true", help="Run in interactive mode"
    )

    args = parser.parse_args()

    cli = SQLMapCLI()
    cli.print_banner()

    # Check if sqlmap exists
    if not SQLMAP_PATH.exists():
        console.print(
            f"[bold red]Error: sqlmap.py not found at {SQLMAP_PATH}[/bold red]",
            style="bold red",
        )
        console.print(
            "[yellow]Make sure you're running this script from the sqlmap directory[/yellow]"
        )
        sys.exit(1)

    # Interactive mode
    if args.interactive:
        cli.interactive_mode()
        return

    # Check if URL is provided
    if not args.url:
        console.print(
            "[bold red]Error: URL is required (use -u or --interactive)[/bold red]"
        )
        parser.print_help()
        sys.exit(1)

    # Run appropriate scan
    verbose_level = args.verbose if args.verbose is not None else 1

    if args.comprehensive:
        cli.comprehensive_scan(
            args.url,
            max_level=args.max_level,
            max_risk=args.max_risk,
            techniques=args.technique,
            data=args.data,
            verbose=verbose_level,
        )
    else:
        cli.quick_scan(
            args.url,
            level=args.level,
            risk=args.risk,
            data=args.data,
            raw=args.raw,
            verbose=verbose_level,
        )


if __name__ == "__main__":
    main()