# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
DKIM (DomainKeys Identified Mail) Manager.

Handles:
- RSA key pair generation
- Message signing
- DNS record generation
"""

import base64
import hashlib
import logging
from datetime import datetime
from email.mime.multipart import MIMEMultipart
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# Optional cryptography for RSA
try:
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import padding, rsa

    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False
    logger.debug("cryptography not installed - DKIM signing disabled")


class DKIMManager:
    """
    DKIM key management and message signing.

    DKIM adds a cryptographic signature to outgoing emails,
    allowing receiving servers to verify the message wasn't
    modified in transit and came from an authorized sender.
    """

    def __init__(
        self,
        domain: str,
        key_path: Path,
        selector: str = "default",
        key_size: int = 2048,
    ):
        self.domain = domain
        self.key_path = Path(key_path)
        self.selector = selector
        self.key_size = key_size

        self._private_key = None
        self._public_key = None

    async def initialize(self):
        """Initialize DKIM keys (generate if needed)."""
        if not HAS_CRYPTO:
            logger.debug("DKIM disabled - cryptography library not installed")
            return

        self.key_path.mkdir(parents=True, exist_ok=True)

        private_key_file = self.key_path / "private.pem"
        public_key_file = self.key_path / "public.pem"

        if private_key_file.exists() and public_key_file.exists():
            # Load existing keys
            await self._load_keys(private_key_file, public_key_file)
            logger.info(f"Loaded DKIM keys for {self.domain}")
        else:
            # Generate new keys
            await self._generate_keys()
            await self._save_keys(private_key_file, public_key_file)
            logger.info(f"Generated new DKIM keys for {self.domain}")

    async def _generate_keys(self):
        """Generate new RSA key pair."""
        self._private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=self.key_size,
            backend=default_backend(),
        )
        self._public_key = self._private_key.public_key()

    async def _load_keys(self, private_path: Path, public_path: Path):
        """Load existing keys from files."""
        private_pem = private_path.read_bytes()
        self._private_key = serialization.load_pem_private_key(
            private_pem,
            password=None,
            backend=default_backend(),
        )
        self._public_key = self._private_key.public_key()

    async def _save_keys(self, private_path: Path, public_path: Path):
        """Save keys to files."""
        # Private key
        private_pem = self._private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        private_path.write_bytes(private_pem)
        private_path.chmod(0o600)

        # Public key
        public_pem = self._public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        public_path.write_bytes(public_pem)

    async def get_public_key_dns(self) -> Optional[str]:
        """
        Get public key formatted for DNS TXT record.

        Returns the value to put in:
        selector._domainkey.domain.com TXT "..."
        """
        if not self._public_key:
            return None

        # Get public key in DER format
        public_der = self._public_key.public_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )

        # Base64 encode
        public_b64 = base64.b64encode(public_der).decode("ascii")

        # Format as DKIM DNS record
        return f"v=DKIM1; k=rsa; p={public_b64}"

    async def sign_message(self, message: MIMEMultipart) -> MIMEMultipart:
        """
        Add DKIM signature to a message.

        Args:
            message: The email message to sign

        Returns:
            The message with DKIM-Signature header added
        """
        if not self._private_key:
            return message

        # Headers to sign
        headers_to_sign = ["from", "to", "subject", "date", "message-id"]

        # Build canonicalized header data
        header_data = []
        signed_headers = []

        for header in headers_to_sign:
            value = message.get(header)
            if value:
                # Relaxed canonicalization
                canonical = f"{header.lower()}:{self._canonicalize_header(value)}"
                header_data.append(canonical)
                signed_headers.append(header)

        # Body hash (relaxed canonicalization)
        body = message.get_payload()
        if isinstance(body, list):
            # Multipart - serialize
            body_text = ""
            for part in body:
                body_text += part.as_string()
        else:
            body_text = str(body) if body else ""

        body_canonical = self._canonicalize_body(body_text)
        body_hash = base64.b64encode(hashlib.sha256(body_canonical.encode()).digest()).decode(
            "ascii"
        )

        # Build DKIM header (without signature)
        timestamp = int(datetime.utcnow().timestamp())

        dkim_header = (
            f"v=1; "
            f"a=rsa-sha256; "
            f"c=relaxed/relaxed; "
            f"d={self.domain}; "
            f"s={self.selector}; "
            f"t={timestamp}; "
            f"h={':'.join(signed_headers)}; "
            f"bh={body_hash}; "
            f"b="
        )

        # Add DKIM header to data to sign
        header_data.append(f"dkim-signature:{self._canonicalize_header(dkim_header)}")

        # Sign
        data_to_sign = "\r\n".join(header_data)

        signature = self._private_key.sign(
            data_to_sign.encode(),
            padding.PKCS1v15(),
            hashes.SHA256(),
        )

        signature_b64 = base64.b64encode(signature).decode("ascii")

        # Add signature to header
        dkim_header += signature_b64

        # Insert DKIM-Signature header at the top
        message._headers.insert(0, ("DKIM-Signature", dkim_header))

        return message

    def _canonicalize_header(self, value: str) -> str:
        """
        Relaxed header canonicalization.

        - Convert to lowercase
        - Unfold header continuation
        - Reduce whitespace to single space
        - Trim trailing whitespace
        """
        # Unfold continuation lines
        value = value.replace("\r\n", "").replace("\n", "")

        # Reduce whitespace
        import re

        value = re.sub(r"\s+", " ", value)

        # Trim
        value = value.strip()

        return value

    def _canonicalize_body(self, body: str) -> str:
        """
        Relaxed body canonicalization.

        - Remove trailing whitespace from lines
        - Reduce whitespace to single space
        - Remove trailing empty lines
        """
        import re

        lines = body.split("\n")
        canonical_lines = []

        for line in lines:
            # Remove trailing whitespace
            line = line.rstrip()
            # Reduce whitespace
            line = re.sub(r"[ \t]+", " ", line)
            canonical_lines.append(line)

        # Remove trailing empty lines
        while canonical_lines and not canonical_lines[-1]:
            canonical_lines.pop()

        # Add final CRLF
        return "\r\n".join(canonical_lines) + "\r\n"

    async def verify_signature(self, message_data: bytes) -> bool:
        """
        Verify DKIM signature of a received message.

        Note: Full verification requires DNS lookup of sender's public key.
        This is a simplified check.
        """
        # Extract DKIM-Signature header
        import email

        msg = email.message_from_bytes(message_data)

        dkim_sig = msg.get("DKIM-Signature")
        if not dkim_sig:
            return False

        # Parse signature
        params = {}
        for part in dkim_sig.split(";"):
            part = part.strip()
            if "=" in part:
                key, value = part.split("=", 1)
                params[key.strip()] = value.strip()

        # Would need to:
        # 1. Fetch public key from DNS (params['s']._domainkey.params['d'])
        # 2. Reconstruct signed data
        # 3. Verify signature

        # For now, just check that required fields exist
        required = ["v", "a", "d", "s", "h", "bh", "b"]
        return all(p in params for p in required)

    def get_dns_record_info(self) -> dict:
        """Get DNS record information for setup."""
        return {
            "host": f"{self.selector}._domainkey",
            "type": "TXT",
            "selector": self.selector,
            "domain": self.domain,
            "full_domain": f"{self.selector}._domainkey.{self.domain}",
        }
