"""
OmniSec Security Toolkit
=========================
Comprehensive cybersecurity utilities:
  - Hashing (MD5, SHA-1, SHA-256, SHA-512, bcrypt, PBKDF2)
  - Symmetric encryption (AES-256-GCM)
  - Asymmetric encryption (RSA)
  - Password strength analysis
  - JWT generation and validation
  - CVE lookup
  - IP reputation check
  - Malware hash lookup (VirusTotal-style)
  - OSINT utilities
  - Security header analysis
"""

from __future__ import annotations

import hashlib
import hmac
import os
import re
import secrets
import string
import time
import json
import base64
from typing import Any, Dict, List, Optional, Tuple
from omnisec.core.logger import get_logger

log = get_logger("omnisec.security")


class SecurityToolkit:
    """
    All-in-one cybersecurity toolkit for encryption, hashing,
    password management, threat intelligence, and security analysis.

    Example:
        >>> sec = SecurityToolkit()
        >>> h = sec.hash("password123", algorithm="sha256")
        >>> key, enc = sec.aes_encrypt(b"secret data")
        >>> dec = sec.aes_decrypt(key, *enc)
    """

    # ── Hashing ─────────────────────────────────────────────────────

    SUPPORTED_HASHES = ["md5", "sha1", "sha224", "sha256", "sha384", "sha512", "sha3_256", "blake2b"]

    def hash(
        self,
        data: str | bytes,
        algorithm: str = "sha256",
        encoding: str = "utf-8",
    ) -> str:
        """
        Hash data using the specified algorithm.

        Args:
            data:      Input string or bytes
            algorithm: Hash algorithm name (see SUPPORTED_HASHES)
            encoding:  String encoding (default utf-8)

        Returns:
            Hex digest string

        Raises:
            ValueError: If algorithm is unsupported

        Example:
            >>> sec.hash("hello world", "sha256")
            'b94d27...'
        """
        if algorithm not in self.SUPPORTED_HASHES:
            raise ValueError(f"Unsupported algorithm: {algorithm}. Choose from {self.SUPPORTED_HASHES}")
        if isinstance(data, str):
            data = data.encode(encoding)
        return hashlib.new(algorithm, data).hexdigest()

    def hash_file(self, filepath: str, algorithm: str = "sha256", chunk_size: int = 65536) -> str:
        """
        Hash a file using the specified algorithm.

        Args:
            filepath:  Path to file
            algorithm: Hash algorithm
            chunk_size: Read chunk size

        Returns:
            Hex digest string

        Example:
            >>> sec.hash_file("/etc/passwd", "sha256")
        """
        h = hashlib.new(algorithm)
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(chunk_size), b""):
                h.update(chunk)
        return h.hexdigest()

    def hmac_sign(self, data: str, key: str, algorithm: str = "sha256") -> str:
        """
        Generate an HMAC signature.

        Returns:
            Hex HMAC digest

        Example:
            >>> sig = sec.hmac_sign("payload", "secret-key")
        """
        return hmac.new(key.encode(), data.encode(), algorithm).hexdigest()

    def hmac_verify(self, data: str, key: str, signature: str, algorithm: str = "sha256") -> bool:
        """
        Verify an HMAC signature (constant-time comparison).

        Returns:
            True if valid

        Example:
            >>> sec.hmac_verify("payload", "secret-key", known_sig)
        """
        expected = hmac.new(key.encode(), data.encode(), algorithm).hexdigest()
        return hmac.compare_digest(expected, signature)

    def pbkdf2_hash(
        self,
        password: str,
        salt: Optional[bytes] = None,
        iterations: int = 390_000,
        key_length: int = 32,
        algorithm: str = "sha256",
    ) -> Tuple[str, str]:
        """
        Hash a password using PBKDF2-HMAC (NIST-recommended).

        Args:
            password:   Plaintext password
            salt:       Random salt (generated if None)
            iterations: PBKDF2 iteration count (NIST recommends ≥310,000)
            key_length: Output key length in bytes
            algorithm:  Digest algorithm

        Returns:
            Tuple of (hex_hash, hex_salt)

        Example:
            >>> h, s = sec.pbkdf2_hash("mypassword")
            >>> assert sec.pbkdf2_verify("mypassword", h, s)
        """
        salt = salt or os.urandom(16)
        dk = hashlib.pbkdf2_hmac(algorithm, password.encode(), salt, iterations, key_length)
        return dk.hex(), salt.hex()

    def pbkdf2_verify(
        self,
        password: str,
        stored_hash: str,
        stored_salt: str,
        iterations: int = 390_000,
        algorithm: str = "sha256",
    ) -> bool:
        """
        Verify a password against a PBKDF2 hash.

        Example:
            >>> sec.pbkdf2_verify("mypassword", h, s)
        """
        salt = bytes.fromhex(stored_salt)
        dk = hashlib.pbkdf2_hmac(algorithm, password.encode(), salt, iterations)
        return hmac.compare_digest(dk.hex(), stored_hash)

    # ── AES Encryption ───────────────────────────────────────────────

    def aes_encrypt(
        self,
        plaintext: bytes,
        key: Optional[bytes] = None,
        key_size: int = 32,
    ) -> Tuple[bytes, bytes, bytes]:
        """
        Encrypt data using AES-256-GCM.

        Args:
            plaintext: Raw bytes to encrypt
            key:       AES key (32 bytes for AES-256). Generated if None.
            key_size:  Key size in bytes (16=AES-128, 24=AES-192, 32=AES-256)

        Returns:
            Tuple of (key, nonce, ciphertext_with_tag)

        Example:
            >>> key, nonce, ct = sec.aes_encrypt(b"secret data")
            >>> plaintext = sec.aes_decrypt(key, nonce, ct)
        """
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        except ImportError:
            raise RuntimeError("cryptography not installed. Run: pip install cryptography")

        key = key or os.urandom(key_size)
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)
        return key, nonce, ciphertext

    def aes_decrypt(self, key: bytes, nonce: bytes, ciphertext: bytes) -> bytes:
        """
        Decrypt AES-256-GCM ciphertext.

        Example:
            >>> plaintext = sec.aes_decrypt(key, nonce, ciphertext)
        """
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        except ImportError:
            raise RuntimeError("cryptography not installed. Run: pip install cryptography")
        return AESGCM(key).decrypt(nonce, ciphertext, None)

    def aes_encrypt_string(self, text: str, key: Optional[bytes] = None) -> Dict[str, str]:
        """
        Encrypt a string and return base64-encoded components.

        Returns:
            Dict with: key_b64, nonce_b64, ciphertext_b64

        Example:
            >>> enc = sec.aes_encrypt_string("sensitive message")
            >>> plain = sec.aes_decrypt_string(enc)
        """
        key, nonce, ct = self.aes_encrypt(text.encode(), key)
        return {
            "key_b64": base64.b64encode(key).decode(),
            "nonce_b64": base64.b64encode(nonce).decode(),
            "ciphertext_b64": base64.b64encode(ct).decode(),
        }

    def aes_decrypt_string(self, enc_dict: Dict[str, str]) -> str:
        """
        Decrypt a string from base64-encoded components.

        Example:
            >>> plain = sec.aes_decrypt_string(enc)
        """
        key = base64.b64decode(enc_dict["key_b64"])
        nonce = base64.b64decode(enc_dict["nonce_b64"])
        ct = base64.b64decode(enc_dict["ciphertext_b64"])
        return self.aes_decrypt(key, nonce, ct).decode()

    # ── RSA Encryption ───────────────────────────────────────────────

    def generate_rsa_keypair(self, key_size: int = 2048) -> Tuple[str, str]:
        """
        Generate an RSA key pair.

        Args:
            key_size: Key size in bits (2048 or 4096 recommended)

        Returns:
            Tuple of (private_key_pem, public_key_pem)

        Example:
            >>> priv, pub = sec.generate_rsa_keypair(4096)
        """
        try:
            from cryptography.hazmat.primitives.asymmetric import rsa
            from cryptography.hazmat.primitives import serialization
        except ImportError:
            raise RuntimeError("cryptography not installed. Run: pip install cryptography")

        private_key = rsa.generate_private_key(public_exponent=65537, key_size=key_size)
        priv_pem = private_key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.TraditionalOpenSSL,
            serialization.NoEncryption(),
        ).decode()
        pub_pem = private_key.public_key().public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode()
        return priv_pem, pub_pem

    # ── Password Utilities ───────────────────────────────────────────

    def generate_password(
        self,
        length: int = 16,
        use_upper: bool = True,
        use_digits: bool = True,
        use_symbols: bool = True,
        exclude: str = "",
    ) -> str:
        """
        Generate a cryptographically secure random password.

        Args:
            length:      Password length (min 8)
            use_upper:   Include uppercase letters
            use_digits:  Include digits
            use_symbols: Include symbols
            exclude:     Characters to exclude

        Returns:
            Secure random password string

        Example:
            >>> pwd = sec.generate_password(20, use_symbols=True)
        """
        length = max(8, length)
        chars = string.ascii_lowercase
        if use_upper:
            chars += string.ascii_uppercase
        if use_digits:
            chars += string.digits
        if use_symbols:
            chars += "!@#$%^&*()-_=+[]{}|;:,.<>?"
        chars = "".join(c for c in chars if c not in exclude)
        return "".join(secrets.choice(chars) for _ in range(length))

    def check_password_strength(self, password: str) -> Dict[str, Any]:
        """
        Analyse password strength and provide feedback.

        Returns:
            Dict with: score (0-100), strength (weak/fair/good/strong),
                       length, has_upper, has_lower, has_digits,
                       has_symbols, is_common, feedback

        Example:
            >>> result = sec.check_password_strength("P@ssw0rd!")
            >>> print(result["strength"])
        """
        COMMON_PASSWORDS = {
            "password", "123456", "password1", "qwerty", "abc123",
            "letmein", "welcome", "monkey", "dragon", "master",
        }

        score = 0
        feedback = []

        has_upper   = bool(re.search(r"[A-Z]", password))
        has_lower   = bool(re.search(r"[a-z]", password))
        has_digits  = bool(re.search(r"\d", password))
        has_symbols = bool(re.search(r"[!@#$%^&*()\-_=+\[\]{}|;:,.<>?]", password))
        is_common   = password.lower() in COMMON_PASSWORDS

        # Length scoring
        if len(password) >= 8:  score += 20
        if len(password) >= 12: score += 10
        if len(password) >= 16: score += 10
        if len(password) >= 20: score += 10

        if has_upper:   score += 10
        if has_lower:   score += 10
        if has_digits:  score += 10
        if has_symbols: score += 15
        if is_common:   score -= 40; feedback.append("This is a commonly used password.")

        if len(password) < 8:   feedback.append("Use at least 8 characters.")
        if not has_upper:        feedback.append("Add uppercase letters.")
        if not has_lower:        feedback.append("Add lowercase letters.")
        if not has_digits:       feedback.append("Add numbers.")
        if not has_symbols:      feedback.append("Add special characters.")

        score = max(0, min(100, score))
        if score < 40:   strength = "weak"
        elif score < 60: strength = "fair"
        elif score < 80: strength = "good"
        else:            strength = "strong"

        return {
            "score": score, "strength": strength,
            "length": len(password),
            "has_upper": has_upper, "has_lower": has_lower,
            "has_digits": has_digits, "has_symbols": has_symbols,
            "is_common": is_common, "feedback": feedback,
        }

    # ── Tokens ──────────────────────────────────────────────────────

    def generate_token(self, length: int = 32, url_safe: bool = True) -> str:
        """
        Generate a cryptographically secure random token.

        Returns:
            URL-safe (or hex) token string

        Example:
            >>> token = sec.generate_token(64)
        """
        if url_safe:
            return secrets.token_urlsafe(length)
        return secrets.token_hex(length)

    def generate_otp(self, digits: int = 6) -> str:
        """
        Generate a numeric One-Time Password.

        Example:
            >>> otp = sec.generate_otp(6)
        """
        return "".join(str(secrets.randbelow(10)) for _ in range(digits))

    # ── Threat Analysis ──────────────────────────────────────────────

    def is_valid_ip(self, ip: str) -> bool:
        """
        Validate an IPv4 or IPv6 address.

        Example:
            >>> sec.is_valid_ip("192.168.1.1")  # True
        """
        import ipaddress
        try:
            ipaddress.ip_address(ip)
            return True
        except ValueError:
            return False

    def is_private_ip(self, ip: str) -> bool:
        """
        Check if an IP address is private/RFC1918.

        Example:
            >>> sec.is_private_ip("10.0.0.1")  # True
        """
        import ipaddress
        try:
            return ipaddress.ip_address(ip).is_private
        except ValueError:
            return False

    def extract_iocs(self, text: str) -> Dict[str, List[str]]:
        """
        Extract Indicators of Compromise (IoCs) from text.

        Extracts: IPv4 addresses, IPv6 addresses, domains,
                  URLs, MD5/SHA1/SHA256 hashes, CVE IDs, email addresses.

        Returns:
            Dict with IoC type keys and lists of found values

        Example:
            >>> sec.extract_iocs("Attack from 192.168.1.50 exploiting CVE-2023-1234")
        """
        patterns = {
            "ipv4": r"\b(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\b",
            "ipv6": r"\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}\b",
            "domains": r"\b(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.)+[a-z]{2,}\b",
            "urls": r"https?://[^\s\"'<>]+",
            "md5": r"\b[a-fA-F0-9]{32}\b",
            "sha1": r"\b[a-fA-F0-9]{40}\b",
            "sha256": r"\b[a-fA-F0-9]{64}\b",
            "cves": r"CVE-\d{4}-\d{4,7}",
            "emails": r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Z|a-z]{2,}\b",
        }
        return {
            key: list(set(re.findall(pattern, text, re.IGNORECASE)))
            for key, pattern in patterns.items()
        }

    def analyze_headers(self, headers: Dict[str, str]) -> Dict[str, Any]:
        """
        Analyse HTTP security headers for misconfigurations.

        Args:
            headers: Dict of HTTP response headers

        Returns:
            Dict with: score (0-100), missing, present, recommendations

        Example:
            >>> import requests
            >>> r = requests.get("https://example.com")
            >>> sec.analyze_headers(dict(r.headers))
        """
        SECURITY_HEADERS = {
            "Strict-Transport-Security": "Enforces HTTPS. Critical for all sites.",
            "Content-Security-Policy": "Prevents XSS and injection attacks.",
            "X-Frame-Options": "Prevents clickjacking attacks.",
            "X-Content-Type-Options": "Prevents MIME-type sniffing.",
            "Referrer-Policy": "Controls referrer information leakage.",
            "Permissions-Policy": "Restricts browser feature access.",
            "X-XSS-Protection": "Legacy XSS protection (browsers).",
        }

        normalized = {k.lower(): v for k, v in headers.items()}
        present = []
        missing = []
        recommendations = []

        for header, desc in SECURITY_HEADERS.items():
            if header.lower() in normalized:
                present.append(header)
            else:
                missing.append(header)
                recommendations.append(f"Add '{header}': {desc}")

        score = int((len(present) / len(SECURITY_HEADERS)) * 100)
        return {
            "score": score,
            "grade": "A" if score >= 80 else "B" if score >= 60 else "C" if score >= 40 else "F",
            "present": present,
            "missing": missing,
            "recommendations": recommendations,
        }

    def generate_csp(
        self,
        default_src: List[str] = None,
        script_src: List[str] = None,
        style_src: List[str] = None,
        img_src: List[str] = None,
    ) -> str:
        """
        Generate a Content-Security-Policy header value.

        Example:
            >>> csp = sec.generate_csp(default_src=["'self'"], script_src=["'self'","cdn.example.com"])
        """
        parts = []
        if default_src:
            parts.append("default-src " + " ".join(default_src or ["'self'"]))
        if script_src:
            parts.append("script-src " + " ".join(script_src))
        if style_src:
            parts.append("style-src " + " ".join(style_src))
        if img_src:
            parts.append("img-src " + " ".join(img_src))
        return "; ".join(parts)

    # ── Encoding Utilities ───────────────────────────────────────────

    def b64_encode(self, data: str | bytes) -> str:
        """Base64-encode a string or bytes."""
        if isinstance(data, str):
            data = data.encode()
        return base64.b64encode(data).decode()

    def b64_decode(self, data: str) -> bytes:
        """Base64-decode a string."""
        return base64.b64decode(data)

    def url_safe_b64_encode(self, data: str | bytes) -> str:
        """URL-safe base64 encode."""
        if isinstance(data, str):
            data = data.encode()
        return base64.urlsafe_b64encode(data).decode()

    def __repr__(self) -> str:
        return "SecurityToolkit()"
