235 lines
9.2 KiB
Python
235 lines
9.2 KiB
Python
"""
|
|
SMTP handler for processing incoming emails.
|
|
"""
|
|
|
|
import uuid
|
|
from datetime import datetime
|
|
from aiosmtpd.smtp import SMTP as AIOSMTP, AuthResult
|
|
from aiosmtpd.controller import Controller
|
|
from email_server.auth import Authenticator, IPAuthenticator
|
|
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 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']
|
|
|
|
def _ensure_required_headers(self, content: str, envelope, message_id: str) -> str:
|
|
"""Ensure all required email headers are present.
|
|
|
|
Args:
|
|
content (str): Email content.
|
|
envelope: SMTP envelope.
|
|
message_id (str): Generated message ID.
|
|
|
|
Returns:
|
|
str: Email content with all required headers.
|
|
"""
|
|
import email.utils
|
|
from datetime import datetime
|
|
|
|
# Parse existing headers
|
|
lines = content.split('\n')
|
|
headers = {}
|
|
body_start = 0
|
|
|
|
# Find where headers end and body begins
|
|
for i, line in enumerate(lines):
|
|
if line.strip() == '':
|
|
body_start = i + 1
|
|
break
|
|
if ':' in line and not line.startswith(' ') and not line.startswith('\t'):
|
|
header_name, header_value = line.split(':', 1)
|
|
headers[header_name.strip().lower()] = header_value.strip()
|
|
|
|
# Extract body
|
|
body = '\n'.join(lines[body_start:]) if body_start < len(lines) else ''
|
|
|
|
# Build required headers
|
|
required_headers = []
|
|
|
|
# Add Message-ID if missing
|
|
if 'message-id' not in headers:
|
|
required_headers.append(f"Message-ID: <{message_id}@{envelope.mail_from.split('@')[1] if '@' in envelope.mail_from else 'localhost'}>")
|
|
|
|
# Add Date if missing
|
|
if 'date' not in headers:
|
|
required_headers.append(f"Date: {email.utils.formatdate(localtime=True)}")
|
|
|
|
# Add From if missing
|
|
if 'from' not in headers:
|
|
required_headers.append(f"From: {envelope.mail_from}")
|
|
|
|
# Add To if missing
|
|
if 'to' not in headers:
|
|
to_list = ', '.join(envelope.rcpt_tos)
|
|
required_headers.append(f"To: {to_list}")
|
|
|
|
# Add MIME-Version if missing
|
|
if 'mime-version' not in headers:
|
|
required_headers.append("MIME-Version: 1.0")
|
|
|
|
# Add Content-Type if missing
|
|
if 'content-type' not in headers:
|
|
required_headers.append("Content-Type: text/plain; charset=utf-8")
|
|
|
|
# Rebuild the email with required headers first, then existing headers, then body
|
|
new_content_lines = required_headers
|
|
|
|
# Add existing headers (excluding the ones we just added)
|
|
for i in range(body_start):
|
|
line = lines[i]
|
|
if ':' in line and not line.startswith(' ') and not line.startswith('\t'):
|
|
header_name = line.split(':', 1)[0].strip().lower()
|
|
if header_name not in ['message-id', 'date', 'from', 'to', 'mime-version', 'content-type']:
|
|
new_content_lines.append(line)
|
|
elif line.startswith(' ') or line.startswith('\t'):
|
|
# Continuation of previous header
|
|
new_content_lines.append(line)
|
|
|
|
# Add empty line between headers and body
|
|
new_content_lines.append('')
|
|
|
|
# Add body
|
|
new_content_lines.append(body)
|
|
|
|
return '\r\n'.join(new_content_lines)
|
|
|
|
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 - 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=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
|
|
)
|