moved server files to its own directory and update tests

email-server/auth.py

@@ -0,0 +1,106 @@
+"""
+Authentication modules for the SMTP server.
+"""
+
+import logging
+from datetime import datetime
+from aiosmtpd.smtp import AuthResult, LoginPassword
+from models import Session, User, Domain, WhitelistedIP, AuthLog, check_password
+
+logger = logging.getLogger(__name__)
+
+class Authenticator:
+    """Username/password authenticator."""
+    
+    def __call__(self, server, session, envelope, mechanism, auth_data):
+        if not isinstance(auth_data, LoginPassword):
+            logger.warning(f'Invalid auth data format: {type(auth_data)}')
+            return AuthResult(success=False, handled=True, message='535 Authentication failed')
+        
+        # Decode bytes to string if necessary
+        username = auth_data.login
+        password = auth_data.password
+        
+        if isinstance(username, bytes):
+            username = username.decode('utf-8')
+        if isinstance(password, bytes):
+            password = password.decode('utf-8')
+            
+        session_db = Session()
+        
+        try:
+            peer_ip = session.peer[0]
+            logger.debug(f'Authentication attempt: {username} from {peer_ip}')
+            
+            # Look up user in database
+            user = session_db.query(User).filter_by(email=username).first()
+            
+            if user and check_password(password, user.password_hash):
+                domain = session_db.query(Domain).filter_by(id=user.domain_id).first()
+                auth_log = AuthLog(
+                    timestamp=datetime.now(),
+                    peer=str(session.peer),
+                    username=username,
+                    success=True,
+                    message=f'Successful login for {username}'
+                )
+                session_db.add(auth_log)
+                session_db.commit()
+                
+                logger.info(f'Authenticated user: {username} from domain {domain.domain_name if domain else "unknown"}')
+                # Don't include the SMTP response code in the message - let aiosmtpd handle it
+                return AuthResult(success=True, handled=True)
+            else:
+                auth_log = AuthLog(
+                    timestamp=datetime.now(),
+                    peer=str(session.peer),
+                    username=username,
+                    success=False,
+                    message=f'Failed login for {username}: invalid credentials'
+                )
+                session_db.add(auth_log)
+                session_db.commit()
+                
+                logger.warning(f'Authentication failed for {username}: invalid credentials')
+                return AuthResult(success=False, handled=True, message='535 Authentication failed')
+            
+        except Exception as e:
+            session_db.rollback()
+            logger.error(f'Authentication error: {e}')
+            return AuthResult(success=False, handled=True, message='451 Internal server error')
+        finally:
+            session_db.close()
+
+class IPAuthenticator:
+    """IP-based authenticator for clients that don't provide credentials."""
+    
+    def __call__(self, server, session, envelope, mechanism, auth_data):
+        session_db = Session()
+        try:
+            peer_ip = session.peer[0]
+            logger.debug(f'IP-based authentication attempt from: {peer_ip}')
+
+            # Check if IP is whitelisted
+            whitelist = session_db.query(WhitelistedIP).filter_by(ip_address=peer_ip).first()
+            if whitelist:
+                domain = session_db.query(Domain).filter_by(id=whitelist.domain_id).first()
+                if domain:
+                    auth_log = AuthLog(
+                        timestamp=datetime.now(),
+                        peer=str(session.peer),
+                        username=None,
+                        success=True,
+                        message=f'Authenticated via whitelisted IP for domain {domain.domain_name}'
+                    )
+                    session_db.add(auth_log)
+                    session_db.commit()
+                    logger.info(f'Authenticated via whitelist: IP {peer_ip} for {domain.domain_name}')
+                    return AuthResult(success=True, handled=True, message='Authenticated via whitelist')
+            
+            return AuthResult(success=False, handled=True, message='IP not whitelisted')
+        except Exception as e:
+            session_db.rollback()
+            logger.error(f'IP Authentication error: {e}')
+            return AuthResult(success=False, handled=True, message='Server error')
+        finally:
+            session_db.close()

email-server/cli_tools.py

@@ -0,0 +1,215 @@
+"""
+Command-line tools for managing the SMTP server.
+"""
+
+import argparse
+import sys
+from models import Session, Domain, User, WhitelistedIP, hash_password, create_tables
+from dkim_manager import DKIMManager
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def add_domain(domain_name, requires_auth=True):
+    """Add a new domain to the database."""
+    session = Session()
+    try:
+        existing = session.query(Domain).filter_by(domain_name=domain_name).first()
+        if existing:
+            print(f"Domain {domain_name} already exists")
+            return False
+        
+        domain = Domain(domain_name=domain_name, requires_auth=requires_auth)
+        session.add(domain)
+        session.commit()
+        print(f"Added domain: {domain_name}")
+        return True
+    except Exception as e:
+        session.rollback()
+        print(f"Error adding domain: {e}")
+        return False
+    finally:
+        session.close()
+
+def add_user(email, password, domain_name):
+    """Add a new user to the database."""
+    session = Session()
+    try:
+        domain = session.query(Domain).filter_by(domain_name=domain_name).first()
+        if not domain:
+            print(f"Domain {domain_name} not found")
+            return False
+        
+        existing = session.query(User).filter_by(email=email).first()
+        if existing:
+            print(f"User {email} already exists")
+            return False
+        
+        user = User(
+            email=email,
+            password_hash=hash_password(password),
+            domain_id=domain.id
+        )
+        session.add(user)
+        session.commit()
+        print(f"Added user: {email}")
+        return True
+    except Exception as e:
+        session.rollback()
+        print(f"Error adding user: {e}")
+        return False
+    finally:
+        session.close()
+
+def add_whitelisted_ip(ip_address, domain_name):
+    """Add an IP to the whitelist for a domain."""
+    session = Session()
+    try:
+        domain = session.query(Domain).filter_by(domain_name=domain_name).first()
+        if not domain:
+            print(f"Domain {domain_name} not found")
+            return False
+        
+        existing = session.query(WhitelistedIP).filter_by(ip_address=ip_address).first()
+        if existing:
+            print(f"IP {ip_address} already whitelisted")
+            return False
+        
+        whitelist = WhitelistedIP(
+            ip_address=ip_address,
+            domain_id=domain.id
+        )
+        session.add(whitelist)
+        session.commit()
+        print(f"Added whitelisted IP: {ip_address} for domain {domain_name}")
+        return True
+    except Exception as e:
+        session.rollback()
+        print(f"Error adding whitelisted IP: {e}")
+        return False
+    finally:
+        session.close()
+
+def generate_dkim_key(domain_name):
+    """Generate DKIM key for a domain."""
+    dkim_manager = DKIMManager()
+    if dkim_manager.generate_dkim_keypair(domain_name):
+        print(f"Generated DKIM key for domain: {domain_name}")
+        
+        # Show DNS record
+        dns_record = dkim_manager.get_dkim_public_key_record(domain_name)
+        if dns_record:
+            print("\nAdd this DNS TXT record:")
+            print(f"Name: {dns_record['name']}")
+            print(f"Value: {dns_record['value']}")
+        return True
+    else:
+        print(f"Failed to generate DKIM key for domain: {domain_name}")
+        return False
+
+def list_dkim_keys():
+    """List all DKIM keys."""
+    dkim_manager = DKIMManager()
+    keys = dkim_manager.list_dkim_keys()
+    
+    if not keys:
+        print("No DKIM keys found")
+        return
+    
+    print("DKIM Keys:")
+    print("-" * 60)
+    for key in keys:
+        status = "ACTIVE" if key['active'] else "INACTIVE"
+        print(f"Domain: {key['domain']}")
+        print(f"Selector: {key['selector']}")
+        print(f"Status: {status}")
+        print(f"Created: {key['created_at']}")
+        print("-" * 60)
+
+def show_dns_records():
+    """Show DNS records for all domains."""
+    dkim_manager = DKIMManager()
+    session = Session()
+    try:
+        domains = session.query(Domain).all()
+        if not domains:
+            print("No domains found")
+            return
+        
+        print("DNS Records for DKIM:")
+        print("=" * 80)
+        
+        for domain in domains:
+            dns_record = dkim_manager.get_dkim_public_key_record(domain.domain_name)
+            if dns_record:
+                print(f"\nDomain: {domain.domain_name}")
+                print(f"Record Name: {dns_record['name']}")
+                print(f"Record Type: {dns_record['type']}")
+                print(f"Record Value: {dns_record['value']}")
+                print("-" * 80)
+    finally:
+        session.close()
+
+def main():
+    """Main CLI function."""
+    parser = argparse.ArgumentParser(description="SMTP Server Management Tool")
+    subparsers = parser.add_subparsers(dest='command', help='Available commands')
+    
+    # Initialize command
+    init_parser = subparsers.add_parser('init', help='Initialize database')
+    
+    # Domain commands
+    domain_parser = subparsers.add_parser('add-domain', help='Add a domain')
+    domain_parser.add_argument('domain', help='Domain name')
+    domain_parser.add_argument('--no-auth', action='store_true', help='Domain does not require authentication')
+    
+    # User commands
+    user_parser = subparsers.add_parser('add-user', help='Add a user')
+    user_parser.add_argument('email', help='User email')
+    user_parser.add_argument('password', help='User password')
+    user_parser.add_argument('domain', help='Domain name')
+    
+    # IP whitelist commands
+    ip_parser = subparsers.add_parser('add-ip', help='Add whitelisted IP')
+    ip_parser.add_argument('ip', help='IP address')
+    ip_parser.add_argument('domain', help='Domain name')
+    
+    # DKIM commands
+    dkim_parser = subparsers.add_parser('generate-dkim', help='Generate DKIM key for domain')
+    dkim_parser.add_argument('domain', help='Domain name')
+    
+    list_dkim_parser = subparsers.add_parser('list-dkim', help='List DKIM keys')
+    
+    dns_parser = subparsers.add_parser('show-dns', help='Show DNS records for DKIM')
+    
+    args = parser.parse_args()
+    
+    if not args.command:
+        parser.print_help()
+        return
+    
+    if args.command == 'init':
+        create_tables()
+        print("Database tables created successfully")
+    
+    elif args.command == 'add-domain':
+        add_domain(args.domain, not args.no_auth)
+    
+    elif args.command == 'add-user':
+        add_user(args.email, args.password, args.domain)
+    
+    elif args.command == 'add-ip':
+        add_whitelisted_ip(args.ip, args.domain)
+    
+    elif args.command == 'generate-dkim':
+        generate_dkim_key(args.domain)
+    
+    elif args.command == 'list-dkim':
+        list_dkim_keys()
+    
+    elif args.command == 'show-dns':
+        show_dns_records()
+
+if __name__ == '__main__':
+    main()

email-server/config.py

@@ -0,0 +1,25 @@
+"""
+Configuration settings for the SMTP server.
+"""
+
+# Server settings
+SMTP_PORT = 4025
+SMTP_TLS_PORT = 40587
+HOSTNAME = 'localhost'
+
+# Database settings
+DATABASE_URL = 'sqlite:///email-server/server_data/smtp_server.db'
+
+# Logging settings
+LOG_LEVEL = 'INFO'
+
+# Email relay settings
+RELAY_TIMEOUT = 10
+
+# TLS settings
+TLS_CERT_FILE = 'email-server/ssl_certs/server.crt'
+TLS_KEY_FILE = 'email-server/ssl_certs/server.key'
+
+# DKIM settings
+DKIM_SELECTOR = 'default'
+DKIM_KEY_SIZE = 2048

email-server/dkim_manager.py

@@ -0,0 +1,212 @@
+"""
+DKIM key management and email signing functionality.
+"""
+
+import logging
+import dkim
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import rsa
+from datetime import datetime
+from models import Session, Domain, DKIMKey
+from config import DKIM_SELECTOR, DKIM_KEY_SIZE
+
+logger = logging.getLogger(__name__)
+
+class DKIMManager:
+    """Manages DKIM keys and email signing."""
+    
+    def __init__(self):
+        self.selector = DKIM_SELECTOR
+    
+    def generate_dkim_keypair(self, domain_name):
+        """Generate DKIM key pair for a domain."""
+        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
+            
+            # Check if DKIM key already exists
+            existing_key = session.query(DKIMKey).filter_by(domain_id=domain.id, is_active=True).first()
+            if existing_key:
+                logger.info(f"DKIM key already exists for domain {domain_name}")
+                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=self.selector,
+                private_key=private_pem,
+                public_key=public_pem,
+                created_at=datetime.now(),
+                is_active=True
+            )
+            session.add(dkim_key)
+            session.commit()
+            
+            logger.info(f"Generated DKIM key for domain: {domain_name}")
+            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_dkim_private_key(self, domain_name):
+        """Get DKIM private key for a domain."""
+        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()
+            
+            if dkim_key:
+                return dkim_key.private_key
+            return None
+            
+        except Exception as e:
+            logger.error(f"Error getting DKIM private key for {domain_name}: {e}")
+            return None
+        finally:
+            session.close()
+    
+    def get_dkim_public_key_record(self, domain_name):
+        """Get DKIM public key DNS record for a domain."""
+        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()
+            
+            if dkim_key:
+                # Extract public key from PEM format for DNS record
+                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'{self.selector}._domainkey.{domain_name}',
+                    'type': 'TXT',
+                    'value': f'v=DKIM1; k=rsa; p={public_key_data}'
+                }
+            return None
+            
+        except Exception as e:
+            logger.error(f"Error getting DKIM public key record for {domain_name}: {e}")
+            return None
+        finally:
+            session.close()
+    
+    def sign_email(self, email_content, domain_name):
+        """Sign email content with DKIM."""
+        try:
+            private_key = self.get_dkim_private_key(domain_name)
+            if not private_key:
+                logger.warning(f"No DKIM key found for domain: {domain_name}")
+                return email_content
+            
+            # 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
+            
+            # Sign the email
+            signature = dkim.sign(
+                email_bytes,
+                self.selector.encode('utf-8'),
+                domain_name.encode('utf-8'),
+                private_key.encode('utf-8'),
+                include_headers=[b'from', b'to', b'subject', b'date', b'message-id']
+            )
+            
+            # Combine signature with original content
+            signed_content = signature + email_bytes
+            
+            logger.info(f"Successfully signed email for domain: {domain_name}")
+            
+            # Return as string if input was string
+            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.info(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()

email-server/email_relay.py

@@ -0,0 +1,71 @@
+"""
+Email relay functionality for the SMTP server.
+"""
+
+import dns.resolver
+import smtplib
+import logging
+from datetime import datetime
+from models import Session, EmailLog
+
+logger = logging.getLogger(__name__)
+
+class EmailRelay:
+    """Handles relaying emails to recipient mail servers."""
+    
+    def __init__(self):
+        self.timeout = 10
+    
+    def relay_email(self, mail_from, rcpt_tos, content):
+        """Relay email to recipient's mail server."""
+        try:
+            for rcpt in rcpt_tos:
+                domain = rcpt.split('@')[1]
+                try:
+                    mx_records = dns.resolver.resolve(domain, 'MX')
+                    mx_host = mx_records[0].exchange.to_text().rstrip('.')
+                except Exception as e:
+                    logger.error(f'Failed to resolve MX for {domain}: {e}')
+                    return False
+
+                try:
+                    with smtplib.SMTP(mx_host, 25, timeout=self.timeout) as relay_server:
+                        relay_server.set_debuglevel(1)
+                        relay_server.sendmail(mail_from, rcpt, content)
+                    logger.info(f'Relayed email to {rcpt} via {mx_host}')
+                except Exception as e:
+                    logger.error(f'Failed to relay email to {rcpt}: {e}')
+                    return False
+            return True
+        except Exception as e:
+            logger.error(f'General relay error: {e}')
+            return False
+    
+    def log_email(self, message_id, peer, mail_from, rcpt_tos, content, status, dkim_signed=False):
+        """Log email activity to database."""
+        session_db = Session()
+        try:
+            # Convert content to string if it's bytes
+            if isinstance(content, bytes):
+                content_str = content.decode('utf-8', errors='replace')
+            else:
+                content_str = content
+            
+            email_log = EmailLog(
+                message_id=message_id,
+                timestamp=datetime.now(),
+                peer=str(peer),
+                mail_from=mail_from,
+                rcpt_tos=', '.join(rcpt_tos),
+                content=content_str,
+                status=status,
+                dkim_signed=dkim_signed
+            )
+            session_db.add(email_log)
+            session_db.commit()
+            logger.debug(f'Logged email: {message_id}')
+        except Exception as e:
+            session_db.rollback()
+            logger.error(f'Error logging email: {e}')
+        finally:
+            session_db.close()

email-server/main.py

@@ -0,0 +1,157 @@
+"""
+Modular SMTP Server with DKIM support.
+Main server file that ties all modules together.
+"""
+
+import asyncio
+import logging
+import sys
+import os
+
+# Import our modules
+from config import SMTP_PORT, SMTP_TLS_PORT, HOSTNAME, LOG_LEVEL
+from models import create_tables
+from smtp_handler import CustomSMTPHandler, PlainController
+from tls_utils import generate_self_signed_cert, create_ssl_context
+from dkim_manager import DKIMManager
+from aiosmtpd.controller import Controller
+from aiosmtpd.smtp import SMTP as AIOSMTP
+
+# Configure logging
+logging.basicConfig(
+    level=getattr(logging, LOG_LEVEL),
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Enable asyncio debugging
+try:
+    loop = asyncio.get_running_loop()
+    loop.set_debug(True)
+except RuntimeError:
+    # No running loop, set debug when we create one
+    pass
+
+async def main():
+    """Main server function."""
+    logger.info("Starting SMTP Server with DKIM support...")
+    
+    # Initialize database
+    logger.info("Initializing database...")
+    create_tables()
+    
+    # Initialize DKIM manager and generate keys for domains without them
+    logger.info("Initializing DKIM manager...")
+    dkim_manager = DKIMManager()
+    dkim_manager.initialize_default_keys()
+    
+    # Add test data if needed
+    from models import Session, Domain, User, WhitelistedIP, hash_password
+    session = Session()
+    try:
+        # Add example.com domain if not exists
+        domain = session.query(Domain).filter_by(domain_name='example.com').first()
+        if not domain:
+            domain = Domain(domain_name='example.com', requires_auth=True)
+            session.add(domain)
+            session.commit()
+            logger.info("Added example.com domain")
+        
+        # Add test user if not exists
+        user = session.query(User).filter_by(email='test@example.com').first()
+        if not user:
+            user = User(
+                email='test@example.com',
+                password_hash=hash_password('testpass123'),
+                domain_id=domain.id
+            )
+            session.add(user)
+            session.commit()
+            logger.info("Added test user: test@example.com")
+        
+        # Add whitelisted IP if not exists
+        whitelist = session.query(WhitelistedIP).filter_by(ip_address='127.0.0.1').first()
+        if not whitelist:
+            whitelist = WhitelistedIP(ip_address='127.0.0.1', domain_id=domain.id)
+            session.add(whitelist)
+            session.commit()
+            logger.info("Added whitelisted IP: 127.0.0.1")
+    except Exception as e:
+        session.rollback()
+        logger.error(f"Error adding test data: {e}")
+    finally:
+        session.close()
+    
+    # Generate SSL certificate if it doesn't exist
+    logger.info("Checking SSL certificates...")
+    if not generate_self_signed_cert():
+        logger.error("Failed to generate SSL certificate")
+        return
+    
+    # Create SSL context
+    ssl_context = create_ssl_context()
+    if not ssl_context:
+        logger.error("Failed to create SSL context")
+        return
+    
+    # Start plain SMTP server (with IP whitelist fallback)
+    handler_plain = CustomSMTPHandler()
+    controller_plain = PlainController(
+        handler_plain,
+        hostname=HOSTNAME,
+        port=SMTP_PORT
+    )
+    controller_plain.start()
+    logger.info(f'Starting plain SMTP server on {HOSTNAME}:{SMTP_PORT}...')
+    
+    # Start TLS SMTP server using closure pattern like the original
+    handler_tls = CustomSMTPHandler()
+    
+    # Define TLS controller class with ssl_context in closure (like original)
+    class TLSController(Controller):
+        def factory(self):
+            return AIOSMTP(
+                self.handler,
+                tls_context=ssl_context,  # Use ssl_context from closure
+                require_starttls=False,  # Don't force STARTTLS, but make it available
+                auth_require_tls=True,   # If auth is used, require TLS
+                authenticator=self.handler.combined_authenticator,
+                decode_data=True,
+                hostname=self.hostname
+            )
+    
+    controller_tls = TLSController(
+        handler_tls,
+        hostname=HOSTNAME,
+        port=SMTP_TLS_PORT
+    )
+    controller_tls.start()
+    logger.info(f'Starting STARTTLS SMTP server on {HOSTNAME}:{SMTP_TLS_PORT}...')
+    
+    logger.info('Both SMTP servers are running:')
+    logger.info(f'  - Plain SMTP (IP whitelist): {HOSTNAME}:{SMTP_PORT}')
+    logger.info(f'  - STARTTLS SMTP (auth required): {HOSTNAME}:{SMTP_TLS_PORT}')
+    logger.info('  - DKIM signing enabled for configured domains')
+    logger.info('')
+    logger.info('Management commands:')
+    logger.info('  python cli_tools.py --help')
+    logger.info('')
+    logger.info('Press Ctrl+C to stop the servers...')
+    
+    try:
+        await asyncio.Event().wait()
+    except KeyboardInterrupt:
+        logger.info('Shutting down SMTP servers...')
+        controller_plain.stop()
+        controller_tls.stop()
+        logger.info('SMTP servers stopped.')
+
+if __name__ == '__main__':
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        logger.info('Server interrupted by user')
+        sys.exit(0)
+    except Exception as e:
+        logger.error(f'Server error: {e}')
+        sys.exit(1)

email-server/models.py

@@ -0,0 +1,76 @@
+"""
+Database models for the SMTP server.
+"""
+
+from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, Boolean
+from sqlalchemy.orm import declarative_base, sessionmaker
+from datetime import datetime
+import bcrypt
+from config import DATABASE_URL
+
+# SQLAlchemy setup
+Base = declarative_base()
+engine = create_engine(DATABASE_URL, echo=False)
+Session = sessionmaker(bind=engine)
+
+class Domain(Base):
+    __tablename__ = 'domains'
+    id = Column(Integer, primary_key=True)
+    domain_name = Column(String, unique=True, nullable=False)
+    requires_auth = Column(Boolean, default=True)
+
+class User(Base):
+    __tablename__ = 'users'
+    id = Column(Integer, primary_key=True)
+    email = Column(String, unique=True, nullable=False)
+    password_hash = Column(String, nullable=False)
+    domain_id = Column(Integer, nullable=False)
+
+class WhitelistedIP(Base):
+    __tablename__ = 'whitelisted_ips'
+    id = Column(Integer, primary_key=True)
+    ip_address = Column(String, unique=True, nullable=False)
+    domain_id = Column(Integer, nullable=False)
+
+class EmailLog(Base):
+    __tablename__ = 'email_logs'
+    id = Column(Integer, primary_key=True)
+    message_id = Column(String, unique=True, nullable=False)
+    timestamp = Column(DateTime, nullable=False)
+    peer = Column(String, nullable=False)
+    mail_from = Column(String, nullable=False)
+    rcpt_tos = Column(String, nullable=False)
+    content = Column(Text, nullable=False)
+    status = Column(String, nullable=False)
+    dkim_signed = Column(Boolean, default=False)
+
+class AuthLog(Base):
+    __tablename__ = 'auth_logs'
+    id = Column(Integer, primary_key=True)
+    timestamp = Column(DateTime, nullable=False)
+    peer = Column(String, nullable=False)
+    username = Column(String)
+    success = Column(Boolean, nullable=False)
+    message = Column(String, nullable=False)
+
+class DKIMKey(Base):
+    __tablename__ = 'dkim_keys'
+    id = Column(Integer, primary_key=True)
+    domain_id = Column(Integer, nullable=False)
+    selector = Column(String, nullable=False)
+    private_key = Column(Text, nullable=False)
+    public_key = Column(Text, nullable=False)
+    created_at = Column(DateTime, default=datetime.now)
+    is_active = Column(Boolean, default=True)
+
+def create_tables():
+    """Create all database tables."""
+    Base.metadata.create_all(engine)
+
+def hash_password(password):
+    """Hash a password using bcrypt."""
+    return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
+
+def check_password(password, hashed):
+    """Check a password against its hash."""
+    return bcrypt.checkpw(password.encode('utf-8'), hashed.encode('utf-8'))

email-server/smtp_handler.py

@@ -0,0 +1,144 @@
+"""
+SMTP handler for processing incoming emails.
+"""
+
+import logging
+import uuid
+from datetime import datetime
+from aiosmtpd.smtp import SMTP as AIOSMTP, AuthResult
+from aiosmtpd.controller import Controller
+from auth import Authenticator, IPAuthenticator
+from email_relay import EmailRelay
+from dkim_manager import DKIMManager
+
+logger = logging.getLogger(__name__)
+
+class CombinedAuthenticator:
+    """Combined authenticator that tries username/password first, then falls back to IP whitelist."""
+    
+    def __init__(self):
+        self.user_auth = Authenticator()
+        self.ip_auth = IPAuthenticator()
+    
+    def __call__(self, server, session, envelope, mechanism, auth_data):
+        from aiosmtpd.smtp import LoginPassword
+        
+        # If auth_data is provided (username/password), try user authentication first
+        if auth_data and isinstance(auth_data, LoginPassword):
+            result = self.user_auth(server, session, envelope, mechanism, auth_data)
+            if result.success:
+                return result
+            # If user auth fails, don't try IP auth - return the failure
+            return result
+        
+        # If no auth_data provided, try IP-based authentication
+        return self.ip_auth(server, session, envelope, mechanism, auth_data)
+
+class CustomSMTPHandler:
+    """Custom SMTP handler for processing emails."""
+    
+    def __init__(self):
+        self.authenticator = Authenticator()
+        self.ip_authenticator = IPAuthenticator()
+        self.combined_authenticator = CombinedAuthenticator()
+        self.email_relay = EmailRelay()
+        self.dkim_manager = DKIMManager()
+        self.auth_require_tls = False
+        self.auth_methods = ['LOGIN', 'PLAIN']
+
+    async def handle_DATA(self, server, session, envelope):
+        """Handle incoming email data."""
+        try:
+            message_id = str(uuid.uuid4())
+            logger.info(f'Received email {message_id} from {envelope.mail_from} to {envelope.rcpt_tos}')
+            
+            # Convert content to string if it's bytes
+            if isinstance(envelope.content, bytes):
+                content = envelope.content.decode('utf-8', errors='replace')
+            else:
+                content = envelope.content
+            
+            # Extract domain from sender for DKIM signing
+            sender_domain = envelope.mail_from.split('@')[1] if '@' in envelope.mail_from else None
+            
+            # Sign with DKIM if domain is configured
+            signed_content = content
+            dkim_signed = False
+            if sender_domain:
+                signed_content = self.dkim_manager.sign_email(content, sender_domain)
+                # Check if signing was successful (content changed)
+                dkim_signed = signed_content != content
+                if dkim_signed:
+                    logger.info(f'Email {message_id} signed with DKIM for domain {sender_domain}')
+            
+            # Relay the email
+            success = self.email_relay.relay_email(
+                envelope.mail_from, 
+                envelope.rcpt_tos, 
+                signed_content
+            )
+            
+            # Log the email
+            status = 'relayed' if success else 'failed'
+            self.email_relay.log_email(
+                message_id=message_id,
+                peer=session.peer,
+                mail_from=envelope.mail_from,
+                rcpt_tos=envelope.rcpt_tos,
+                content=content,  # Log original content, not signed
+                status=status,
+                dkim_signed=dkim_signed
+            )
+            
+            if success:
+                logger.info(f'Email {message_id} successfully relayed')
+                return '250 Message accepted for delivery'
+            else:
+                logger.error(f'Email {message_id} failed to relay')
+                return '550 Message relay failed'
+                
+        except Exception as e:
+            logger.error(f'Error handling email: {e}')
+            return '550 Internal server error'
+
+    async def handle_RCPT(self, server, session, envelope, address, rcpt_options):
+        """Handle RCPT TO command - validate recipients."""
+        logger.debug(f'RCPT TO: {address}')
+        envelope.rcpt_tos.append(address)
+        return '250 OK'
+
+    async def handle_MAIL(self, server, session, envelope, address, mail_options):
+        """Handle MAIL FROM command - validate sender."""
+        logger.debug(f'MAIL FROM: {address}')
+        envelope.mail_from = address
+        return '250 OK'
+
+class TLSController(Controller):
+    """Custom controller with TLS support - modeled after the working original."""
+    
+    def __init__(self, handler, ssl_context, hostname='localhost', port=40587):
+        self.ssl_context = ssl_context
+        super().__init__(handler, hostname=hostname, port=port)
+    
+    def factory(self):
+        return AIOSMTP(
+            self.handler,
+            tls_context=self.ssl_context,
+            require_starttls=False,  # Don't force STARTTLS, but make it available
+            auth_require_tls=True,   # If auth is used, require TLS
+            authenticator=self.handler.combined_authenticator,
+            decode_data=True,
+            hostname=self.hostname
+        )
+
+class PlainController(Controller):
+    """Controller for plain SMTP with username/password and IP-based authentication."""
+    
+    def factory(self):
+        return AIOSMTP(
+            self.handler,
+            authenticator=self.handler.combined_authenticator,
+            auth_require_tls=False,  # Allow AUTH over plain text (not recommended for production)
+            decode_data=True,
+            hostname=self.hostname
+        )

email-server/tls_utils.py

@@ -0,0 +1,61 @@
+"""
+TLS utilities for the SMTP server.
+"""
+
+import ssl
+import os
+import logging
+from OpenSSL import crypto
+from config import TLS_CERT_FILE, TLS_KEY_FILE
+
+logger = logging.getLogger(__name__)
+
+def generate_self_signed_cert():
+    """Generate self-signed SSL certificate if it doesn't exist."""
+    if os.path.exists(TLS_CERT_FILE) and os.path.exists(TLS_KEY_FILE):
+        logger.info("SSL certificate already exists")
+        return True
+    
+    try:
+        logger.info("Generating self-signed SSL certificate...")
+        
+        # Generate private key
+        k = crypto.PKey()
+        k.generate_key(crypto.TYPE_RSA, 2048)
+        
+        # Generate certificate
+        cert = crypto.X509()
+        cert.get_subject().CN = 'localhost'
+        cert.set_serial_number(1000)
+        cert.gmtime_adj_notBefore(0)
+        cert.gmtime_adj_notAfter(365 * 24 * 60 * 60)  # Valid for 1 year
+        cert.set_issuer(cert.get_subject())
+        cert.set_pubkey(k)
+        cert.sign(k, 'sha256')
+        
+        # Write certificate
+        with open(TLS_CERT_FILE, 'wb') as f:
+            f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert))
+        
+        # Write private key
+        with open(TLS_KEY_FILE, 'wb') as f:
+            f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k))
+        
+        logger.info(f"SSL certificate generated: {TLS_CERT_FILE}, {TLS_KEY_FILE}")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to generate SSL certificate: {e}")
+        return False
+
+def create_ssl_context():
+    """Create SSL context for TLS support."""
+    try:
+        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+        ssl_context.load_cert_chain(certfile=TLS_CERT_FILE, keyfile=TLS_KEY_FILE)
+        ssl_context.set_ciphers('DEFAULT')  # Relax ciphers for compatibility
+        logger.info('SSL context created successfully')
+        return ssl_context
+    except Exception as e:
+        logger.error(f'Failed to create SSL context: {e}')
+        return None