Files
PyMTA-server/email_server/smtp_handler.py
2025-06-01 12:09:12 +01:00

228 lines
8.9 KiB
Python

"""
Enhanced SMTP handler for processing incoming emails with security controls.
Security Features:
- Users can only send as their own email or domain emails (if permitted)
- IP authentication is domain-specific
- Sender authorization validation
- Enhanced header management
"""
import uuid
import email.utils
from datetime import datetime
from aiosmtpd.smtp import SMTP as AIOSMTP, AuthResult
from aiosmtpd.controller import Controller
from email_server.auth import EnhancedAuthenticator, EnhancedIPAuthenticator, validate_sender_authorization, get_authenticated_domain_id
from email_server.email_relay import EmailRelay
from email_server.dkim_manager import DKIMManager
from email_server.tool_box import get_logger
logger = get_logger()
class EnhancedCombinedAuthenticator:
"""
Enhanced combined authenticator with sender validation support.
Features:
- User authentication with session storage
- IP-based authentication with domain validation
- Fallback authentication logic
"""
def __init__(self):
self.user_auth = EnhancedAuthenticator()
self.ip_auth = EnhancedIPAuthenticator()
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, IP auth will be validated during MAIL FROM
# For now, allow the connection to proceed
return AuthResult(success=True, handled=True)
class EnhancedCustomSMTPHandler:
"""Enhanced custom SMTP handler with security controls."""
def __init__(self):
self.authenticator = EnhancedAuthenticator()
self.ip_authenticator = EnhancedIPAuthenticator()
self.combined_authenticator = EnhancedCombinedAuthenticator()
self.email_relay = EmailRelay()
self.dkim_manager = DKIMManager()
self.auth_require_tls = False
self.auth_methods = ['LOGIN', 'PLAIN']
def _ensure_required_headers(self, content: str, envelope, message_id: str) -> str:
"""Ensure all required email headers are present and properly formatted.
Args:
content (str): Email content.
envelope: SMTP envelope.
message_id (str): Generated message ID.
Returns:
str: Email content with all required headers.
"""
import email
from email.parser import Parser
from email.policy import default
# Parse the message using the email library
msg = Parser(policy=default).parsestr(content)
# Set or add required headers if missing
if not msg.get('Message-ID'):
msg['Message-ID'] = f"<{message_id}@{envelope.mail_from.split('@')[1] if '@' in envelope.mail_from else 'localhost'}>"
if not msg.get('Date'):
msg['Date'] = email.utils.formatdate(localtime=True)
if not msg.get('From'):
msg['From'] = envelope.mail_from
if not msg.get('To'):
msg['To'] = ', '.join(envelope.rcpt_tos)
if not msg.get('MIME-Version'):
msg['MIME-Version'] = '1.0'
if not msg.get('Content-Type'):
msg['Content-Type'] = 'text/plain; charset=utf-8'
if not msg.get('Subject'):
msg['Subject'] = '(No Subject)'
if not msg.get('Content-Transfer-Encoding'):
msg['Content-Transfer-Encoding'] = '7bit'
# Ensure exactly one blank line between headers and body
# The email library will handle this when flattening
from io import StringIO
out = StringIO()
out.write(msg.as_string())
return out.getvalue()
async def handle_DATA(self, server, session, envelope):
"""Handle incoming email data."""
try:
message_id = str(uuid.uuid4())
logger.debug(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
# Ensure required headers are present
content = self._ensure_required_headers(content, envelope, message_id)
# Add custom headers before DKIM signing
if sender_domain:
custom_headers = self.dkim_manager.get_active_custom_headers(sender_domain)
for header_name, header_value in custom_headers:
# Insert header at the top of the message
content = f"{header_name}: {header_value}\r\n" + content
# Relay the email (all modifications done)
signed_content = content
dkim_signed = False
if sender_domain:
# DKIM-sign the final version of the message
signed_content = self.dkim_manager.sign_email(content, sender_domain)
dkim_signed = signed_content != content
if dkim_signed:
logger.debug(f'Email {message_id} signed with DKIM for domain {sender_domain}')
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.debug(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 with enhanced sender validation.
Security Features:
- Validates user can send as the specified address
- Validates IP authorization for domain
- Comprehensive audit logging
"""
logger.debug(f'MAIL FROM: {address}')
# Validate sender authorization
authorized, message = validate_sender_authorization(session, address)
if not authorized:
logger.warning(f'MAIL FROM rejected: {address} - {message}')
return f'550 {message}'
envelope.mail_from = address
logger.info(f'MAIL FROM accepted: {address} - {message}')
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=True, # 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
)