PyMTA-server/email_server/dkim_manager.py

"""
DKIM key management and email signing functionality.
"""

import dkim
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa
from datetime import datetime
from email_server.models import Session, Domain, DKIMKey
from email_server.settings_loader import load_settings
from email_server.tool_box import get_logger
import random
import string

settings = load_settings()
DKIM_SELECTOR = settings['DKIM']['DKIM_SELECTOR']
DKIM_KEY_SIZE = int(settings['DKIM']['DKIM_KEY_SIZE'])

logger = get_logger()

class DKIMManager:
    """Manages DKIM keys and email signing."""
    
    def __init__(self, selector: str = None):
        """Initialize DKIMManager with a selector. If not provided, use random."""
        if selector:
            self.selector = selector
        else:
            self.selector = self._generate_random_selector()
    
    @staticmethod
    def _generate_random_selector(length: int = 12) -> str:
        """Generate a random DKIM selector name (8-12 chars)."""
        return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

    def generate_dkim_keypair(self, domain_name, selector: str = None):
        """Generate DKIM key pair for a domain, optionally with a custom selector."""
        session = Session()
        try:
            # Check if domain exists
            domain = session.query(Domain).filter_by(domain_name=domain_name).first()
            if not domain:
                logger.error(f"Domain {domain_name} not found")
                return False
            
            # Use provided selector or instance selector
            use_selector = selector or self.selector
            
            # Check if DKIM key with this selector already exists
            existing_key = session.query(DKIMKey).filter_by(domain_id=domain.id, selector=use_selector, is_active=True).first()
            if existing_key:
                logger.debug(f"DKIM key already exists for domain {domain_name} and selector {use_selector}")
                return True
            
            # Generate RSA key pair
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=DKIM_KEY_SIZE
            )
            
            # Get private key in PEM format
            private_pem = private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            ).decode('utf-8')
            
            # Get public key in PEM format
            public_key = private_key.public_key()
            public_pem = public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            ).decode('utf-8')
            
            # Store in database
            dkim_key = DKIMKey(
                domain_id=domain.id,
                selector=use_selector,
                private_key=private_pem,
                public_key=public_pem,
                created_at=datetime.now(),
                is_active=True
            )
            session.add(dkim_key)
            session.commit()
            
            logger.debug(f"Generated DKIM key for domain: {domain_name} selector: {use_selector}")
            return True
            
        except Exception as e:
            session.rollback()
            logger.error(f"Error generating DKIM key for {domain_name}: {e}")
            return False
        finally:
            session.close()

    def get_active_dkim_key(self, domain_name):
        """Get the active DKIM key for a domain (only one active per selector)."""
        session = Session()
        try:
            domain = session.query(Domain).filter_by(domain_name=domain_name).first()
            if not domain:
                return None
            dkim_key = session.query(DKIMKey).filter_by(domain_id=domain.id, is_active=True).first()
            return dkim_key
        except Exception as e:
            logger.error(f"Error getting DKIM key for {domain_name}: {e}")
            return None
        finally:
            session.close()

    def get_dkim_private_key(self, domain_name):
        """Get DKIM private key for a domain."""
        dkim_key = self.get_active_dkim_key(domain_name)
        if dkim_key:
            return dkim_key.private_key
        return None
    
    def get_dkim_public_key_record(self, domain_name):
        """Get DKIM public key DNS record for a domain (active key only)."""
        dkim_key = self.get_active_dkim_key(domain_name)
        if dkim_key:
            public_key_lines = dkim_key.public_key.strip().split('\n')
            public_key_data = ''.join(public_key_lines[1:-1])  # Remove header/footer
            return {
                'name': f'{dkim_key.selector}._domainkey.{domain_name}',
                'type': 'TXT',
                'value': f'v=DKIM1; k=rsa; p={public_key_data}'
            }
        return None

    def sign_email(self, email_content, domain_name):
        """Sign email content with DKIM. Only add one DKIM header, after all modifications."""
        try:
            dkim_key = self.get_active_dkim_key(domain_name)
            if not dkim_key:
                logger.warning(f"No DKIM key found for domain: {domain_name}")
                return email_content
            private_key = dkim_key.private_key
            selector = dkim_key.selector
            # Convert content to bytes if it's a string
            if isinstance(email_content, str):
                email_bytes = email_content.encode('utf-8')
            else:
                email_bytes = email_content
            # Remove any existing DKIM-Signature header (robust multiline)
            import re
            email_bytes = re.sub(br'^DKIM-Signature:.*?(\r?\n[ \t].*?)*\r?\n', b'', email_bytes, flags=re.MULTILINE)
            # Canonicalization: relaxed/relaxed, add more headers to h=
            headers = [
                b'from', b'to', b'subject', b'date', b'message-id',
                b'mime-version', b'content-type', b'content-transfer-encoding'
            ]
            signature = dkim.sign(
                email_bytes,
                selector.encode('utf-8'),
                domain_name.encode('utf-8'),
                private_key.encode('utf-8'),
                include_headers=headers,
                canonicalize=(b'relaxed', b'relaxed')
            )
            signed_content = signature + email_bytes
            logger.debug(f"Successfully signed email for domain: {domain_name} selector: {selector}")
            if isinstance(email_content, str):
                return signed_content.decode('utf-8')
            else:
                return signed_content
        except Exception as e:
            logger.error(f"Error signing email for domain {domain_name}: {e}")
            return email_content
    
    def list_dkim_keys(self):
        """List all DKIM keys."""
        session = Session()
        try:
            keys = session.query(DKIMKey, Domain).join(Domain).all()
            result = []
            
            for dkim_key, domain in keys:
                result.append({
                    'domain': domain.domain_name,
                    'selector': dkim_key.selector,
                    'created_at': dkim_key.created_at,
                    'active': dkim_key.is_active
                })
            
            return result
            
        except Exception as e:
            logger.error(f"Error listing DKIM keys: {e}")
            return []
        finally:
            session.close()
    
    def initialize_default_keys(self):
        """Initialize DKIM keys for existing domains that don't have them."""
        session = Session()
        try:
            domains = session.query(Domain).all()
            for domain in domains:
                existing_key = session.query(DKIMKey).filter_by(
                    domain_id=domain.id, 
                    is_active=True
                ).first()
                
                if not existing_key:
                    logger.debug(f"Generating DKIM key for existing domain: {domain.domain_name}")
                    self.generate_dkim_keypair(domain.domain_name)
                    
        except Exception as e:
            logger.error(f"Error initializing default DKIM keys: {e}")
        finally:
            session.close()