Merge pull request #1 from ghostersk/testing

Testing

.gitignore

@@ -12,6 +12,8 @@ test.txt
 custom_test.py
 custom_test.sh
 .github/
+*.backup
+mytest_*
 
 # C extensions
 *.so

alembic.ini

@@ -0,0 +1,141 @@
+# A generic, single database configuration.
+
+[alembic]
+script_location = alembic
+# path to migration scripts.
+# this is typically a path given in POSIX (e.g. forward slashes)
+# format, relative to the token %(here)s which refers to the location of this
+# ini file
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.  for multiple paths, the path separator
+# is defined by "path_separator" below.
+prepend_sys_path = .
+
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python>=3.9 or backports.zoneinfo library and tzdata library.
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to ZoneInfo()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to <script_location>/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "path_separator"
+# below.
+# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions
+
+# path_separator; This indicates what character is used to split lists of file
+# paths, including version_locations and prepend_sys_path within configparser
+# files such as alembic.ini.
+# The default rendered in new alembic.ini files is "os", which uses os.pathsep
+# to provide os-dependent path splitting.
+#
+# Note that in order to support legacy alembic.ini files, this default does NOT
+# take place if path_separator is not present in alembic.ini.  If this
+# option is omitted entirely, fallback logic is as follows:
+#
+# 1. Parsing of the version_locations option falls back to using the legacy
+#    "version_path_separator" key, which if absent then falls back to the legacy
+#    behavior of splitting on spaces and/or commas.
+# 2. Parsing of the prepend_sys_path option falls back to the legacy
+#    behavior of splitting on spaces, commas, or colons.
+#
+# Valid values for path_separator are:
+#
+# path_separator = :
+# path_separator = ;
+# path_separator = space
+# path_separator = newline
+#
+# Use os.pathsep. Default configuration used for new projects.
+path_separator = os
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+# database URL.  This is consumed by the user-maintained env.py script only.
+# other means of configuring database URLs may be customized within the env.py
+# file.
+sqlalchemy.url = sqlite:///email_server/server_data/smtp_server.db
+
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
+# hooks = ruff
+# ruff.type = exec
+# ruff.executable = %(here)s/.venv/bin/ruff
+# ruff.options = check --fix REVISION_SCRIPT_FILENAME
+
+# Logging configuration.  This is also consumed by the user-maintained
+# env.py script only.
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARNING
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARNING
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

alembic/README

@@ -0,0 +1 @@
+Generic single-database configuration.

alembic/env.py

@@ -0,0 +1,79 @@
+from logging.config import fileConfig
+
+from sqlalchemy import engine_from_config
+from sqlalchemy import pool
+
+from alembic import context
+from email_server.models import Base
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+target_metadata = Base.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    with connectable.connect() as connection:
+        context.configure(
+            connection=connection, target_metadata=target_metadata
+        )
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

alembic/script.py.mako

@@ -0,0 +1,28 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    ${downgrades if downgrades else "pass"}

alembic/versions/53036910f343_initial_migration.py

@@ -0,0 +1,133 @@
+"""Initial migration
+
+Revision ID: 53036910f343
+Revises: 
+Create Date: 2025-06-01 11:14:39.589608
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '53036910f343'
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    # The following columns may already exist, so we skip adding them if present.
+    # op.add_column('esrv_auth_logs', sa.Column('auth_type', sa.String(), nullable=False, server_default='user'))
+    # op.add_column('esrv_auth_logs', sa.Column('identifier', sa.String(), nullable=False, server_default='unknown'))
+    # op.add_column('esrv_auth_logs', sa.Column('ip_address', sa.String(), nullable=True))
+    # op.add_column('esrv_auth_logs', sa.Column('created_at', sa.DateTime(), nullable=True))
+    # SQLite does not support ALTER COLUMN, so we skip changing 'message' type/nullability here.
+    # If you need to change the type/nullability, do a manual migration as described in Alembic docs.
+    # op.drop_column('esrv_auth_logs', 'peer')
+    # op.drop_column('esrv_auth_logs', 'timestamp')
+    # op.drop_column('esrv_auth_logs', 'username')
+    # op.alter_column('esrv_custom_headers', 'created_at',
+    #            existing_type=sa.TIMESTAMP(),
+    #            type_=sa.DateTime(),
+    #            existing_nullable=True,
+    #            existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.drop_index(op.f('idx_esrv_custom_headers_domain'), table_name='esrv_custom_headers')
+    op.drop_index(op.f('idx_esrv_dkim_keys_domain'), table_name='esrv_dkim_keys')
+    # op.alter_column('esrv_domains', 'created_at',
+    #            existing_type=sa.TIMESTAMP(),
+    #            type_=sa.DateTime(),
+    #            existing_nullable=True,
+    #            existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.drop_column('esrv_domains', 'requires_auth')
+    op.add_column('esrv_email_logs', sa.Column('from_address', sa.String(), nullable=False, server_default='unknown'))
+    op.add_column('esrv_email_logs', sa.Column('to_address', sa.String(), nullable=False, server_default='unknown'))
+    op.add_column('esrv_email_logs', sa.Column('subject', sa.Text(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('message', sa.Text(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('created_at', sa.DateTime(), nullable=True))
+    op.drop_index(op.f('idx_esrv_email_logs_mail_from'), table_name='esrv_email_logs')
+    op.drop_index(op.f('idx_esrv_email_logs_timestamp'), table_name='esrv_email_logs')
+    # op.drop_column('esrv_email_logs', 'rcpt_tos')
+    # op.drop_column('esrv_email_logs', 'content')
+    # op.drop_column('esrv_email_logs', 'dkim_signed')
+    # op.drop_column('esrv_email_logs', 'peer')
+    # op.drop_column('esrv_email_logs', 'message_id')
+    # op.drop_column('esrv_email_logs', 'mail_from')
+    # op.drop_column('esrv_email_logs', 'timestamp')
+    # op.alter_column('esrv_users', 'created_at',
+    #            existing_type=sa.TIMESTAMP(),
+    #            type_=sa.DateTime(),
+    #            existing_nullable=True,
+    #            existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.drop_index(op.f('idx_esrv_users_domain'), table_name='esrv_users')
+    op.drop_index(op.f('idx_esrv_users_email'), table_name='esrv_users')
+    # op.alter_column('esrv_whitelisted_ips', 'created_at',
+    #            existing_type=sa.TIMESTAMP(),
+    #            type_=sa.DateTime(),
+    #            existing_nullable=True,
+    #            existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.drop_index(op.f('idx_esrv_whitelisted_ips_domain'), table_name='esrv_whitelisted_ips')
+    op.drop_index(op.f('idx_esrv_whitelisted_ips_ip'), table_name='esrv_whitelisted_ips')
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_index(op.f('idx_esrv_whitelisted_ips_ip'), 'esrv_whitelisted_ips', ['ip_address'], unique=False)
+    op.create_index(op.f('idx_esrv_whitelisted_ips_domain'), 'esrv_whitelisted_ips', ['domain_id'], unique=False)
+    op.alter_column('esrv_whitelisted_ips', 'created_at',
+               existing_type=sa.DateTime(),
+               type_=sa.TIMESTAMP(),
+               existing_nullable=True,
+               existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.create_index(op.f('idx_esrv_users_email'), 'esrv_users', ['email'], unique=False)
+    op.create_index(op.f('idx_esrv_users_domain'), 'esrv_users', ['domain_id'], unique=False)
+    op.alter_column('esrv_users', 'created_at',
+               existing_type=sa.DateTime(),
+               type_=sa.TIMESTAMP(),
+               existing_nullable=True,
+               existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.add_column('esrv_email_logs', sa.Column('timestamp', sa.DATETIME(), nullable=False))
+    op.add_column('esrv_email_logs', sa.Column('mail_from', sa.VARCHAR(), nullable=False))
+    op.add_column('esrv_email_logs', sa.Column('message_id', sa.VARCHAR(), nullable=False))
+    op.add_column('esrv_email_logs', sa.Column('peer', sa.VARCHAR(), nullable=False))
+    op.add_column('esrv_email_logs', sa.Column('dkim_signed', sa.BOOLEAN(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('content', sa.TEXT(), nullable=False))
+    op.add_column('esrv_email_logs', sa.Column('rcpt_tos', sa.VARCHAR(), nullable=False))
+    op.create_index(op.f('idx_esrv_email_logs_timestamp'), 'esrv_email_logs', ['timestamp'], unique=False)
+    op.create_index(op.f('idx_esrv_email_logs_mail_from'), 'esrv_email_logs', ['mail_from'], unique=False)
+    op.drop_column('esrv_email_logs', 'created_at')
+    op.drop_column('esrv_email_logs', 'message')
+    op.drop_column('esrv_email_logs', 'subject')
+    op.drop_column('esrv_email_logs', 'to_address')
+    op.drop_column('esrv_email_logs', 'from_address')
+    op.add_column('esrv_domains', sa.Column('requires_auth', sa.BOOLEAN(), nullable=True))
+    op.alter_column('esrv_domains', 'created_at',
+               existing_type=sa.DateTime(),
+               type_=sa.TIMESTAMP(),
+               existing_nullable=True,
+               existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.create_index(op.f('idx_esrv_dkim_keys_domain'), 'esrv_dkim_keys', ['domain_id'], unique=False)
+    op.create_index(op.f('idx_esrv_custom_headers_domain'), 'esrv_custom_headers', ['domain_id'], unique=False)
+    op.alter_column('esrv_custom_headers', 'created_at',
+               existing_type=sa.DateTime(),
+               type_=sa.TIMESTAMP(),
+               existing_nullable=True,
+               existing_server_default=sa.text("'2025-05-31 00:00:00'"))
+    op.add_column('esrv_auth_logs', sa.Column('username', sa.VARCHAR(), nullable=True))
+    op.add_column('esrv_auth_logs', sa.Column('timestamp', sa.DATETIME(), nullable=False))
+    op.add_column('esrv_auth_logs', sa.Column('peer', sa.VARCHAR(), nullable=False))
+    op.alter_column('esrv_auth_logs', 'message',
+               existing_type=sa.Text(),
+               type_=sa.VARCHAR(),
+               nullable=False)
+    op.drop_column('esrv_auth_logs', 'created_at')
+    op.drop_column('esrv_auth_logs', 'ip_address')
+    op.drop_column('esrv_auth_logs', 'identifier')
+    op.drop_column('esrv_auth_logs', 'auth_type')
+    # ### end Alembic commands ###

alembic/versions/d02f993649e8_add_legacy_emaillog_columns_for_.py

@@ -0,0 +1,42 @@
+"""Add legacy EmailLog columns for backward compatibility
+
+Revision ID: d02f993649e8
+Revises: 53036910f343
+Create Date: 2025-06-01 11:50:54.362830
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'd02f993649e8'
+down_revision: Union[str, None] = '53036910f343'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # Add legacy columns for backward compatibility with existing EmailRelay code
+    op.add_column('esrv_email_logs', sa.Column('message_id', sa.String(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('timestamp', sa.DateTime(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('peer', sa.String(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('mail_from', sa.String(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('rcpt_tos', sa.String(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('content', sa.Text(), nullable=True))
+    op.add_column('esrv_email_logs', sa.Column('dkim_signed', sa.Boolean(), nullable=True, default=False))
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # Remove the legacy columns
+    op.drop_column('esrv_email_logs', 'dkim_signed')
+    op.drop_column('esrv_email_logs', 'content')
+    op.drop_column('esrv_email_logs', 'rcpt_tos')
+    op.drop_column('esrv_email_logs', 'mail_from')
+    op.drop_column('esrv_email_logs', 'peer')
+    op.drop_column('esrv_email_logs', 'timestamp')
+    op.drop_column('esrv_email_logs', 'message_id')

email_server/auth.py

@@ -1,16 +1,33 @@
 """
-Authentication modules for the SMTP server.
+Enhanced authentication modules for the SMTP server using ESRV schema.
+
+Security Features:
+- Users can only send as their own email or domain emails (if permitted)
+- IP authentication is domain-specific  
+- Comprehensive audit logging
+- Enhanced validation and error handling
 """
 
 from datetime import datetime
 from aiosmtpd.smtp import AuthResult, LoginPassword
-from email_server.models import Session, User, Domain, WhitelistedIP, AuthLog, check_password
+from email_server.models import (
+    Session, User, Domain, WhitelistedIP, 
+    check_password, log_auth_attempt, get_user_by_email, 
+    get_whitelisted_ip, get_domain_by_name
+)
 from email_server.tool_box import get_logger
 
 logger = get_logger()
 
-class Authenticator:
-    """Username/password authenticator."""
+class EnhancedAuthenticator:
+    """
+    Enhanced username/password authenticator with sender validation.
+    
+    Features:
+    - Validates user credentials
+    - Stores authenticated user info for sender validation
+    - Comprehensive audit logging
+    """
     
     def __call__(self, server, session, envelope, mechanism, auth_data):
         if not isinstance(auth_data, LoginPassword):
@@ -25,82 +42,174 @@ class Authenticator:
             username = username.decode('utf-8')
         if isinstance(password, bytes):
             password = password.decode('utf-8')
-            
-        session_db = Session()
+        
+        peer_ip = session.peer[0]
+        logger.debug(f'Authentication attempt: {username} from {peer_ip}')
         
         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()
+            user = get_user_by_email(username)
             
             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()
+                # Store authenticated user info in session for later validation
+                session.authenticated_user = user
+                session.auth_type = 'user'
                 
-                logger.debug(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
+                # Log successful authentication
+                log_auth_attempt(
+                    auth_type='user',
+                    identifier=username,
+                    ip_address=peer_ip,
+                    success=True,
+                    message=f'Successful user authentication'
+                )
+                
+                logger.info(f'Authenticated user: {username} (ID: {user.id}, can_send_as_domain: {user.can_send_as_domain})')
                 return AuthResult(success=True, handled=True)
             else:
-                auth_log = AuthLog(
-                    timestamp=datetime.now(),
-                    peer=str(session.peer),
-                    username=username,
+                # Log failed authentication
+                log_auth_attempt(
+                    auth_type='user',
+                    identifier=username,
+                    ip_address=peer_ip,
                     success=False,
-                    message=f'Failed login for {username}: invalid credentials'
+                    message=f'Invalid credentials for {username}'
                 )
-                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}')
+            logger.error(f'Authentication error for {username}: {e}')
+            log_auth_attempt(
+                auth_type='user',
+                identifier=username,
+                ip_address=peer_ip,
+                success=False,
+                message=f'Authentication error: {str(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."""
+class EnhancedIPAuthenticator:
+    """
+    Enhanced IP-based authenticator with domain-specific authorization.
     
-    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.debug(f'Authenticated via whitelist: IP {peer_ip} for {domain.domain_name}')
-                    return AuthResult(success=True, handled=True, message='Authenticated via whitelist')
+    Features:
+    - Domain-specific IP authentication
+    - Only allows sending for authorized domain
+    - Comprehensive audit logging
+    """
+    
+    def can_authenticate_for_domain(self, ip_address: str, domain_name: str) -> tuple[bool, str]:
+        """
+        Check if IP can authenticate for a specific domain.
+        
+        Args:
+            ip_address: Client IP address
+            domain_name: Domain to check authorization for
             
-            return AuthResult(success=False, handled=True, message='IP not whitelisted')
+        Returns:
+            (success, message) tuple
+        """
+        try:
+            whitelisted_ip = get_whitelisted_ip(ip_address, domain_name)
+            if whitelisted_ip:
+                return True, f"IP {ip_address} authorized for domain {domain_name}"
+            else:
+                return False, f"IP {ip_address} not authorized for domain {domain_name}"
         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()
+            logger.error(f"Error checking IP authorization: {e}")
+            return False, f"Error checking IP authorization: {str(e)}"
+
+def validate_sender_authorization(session, mail_from: str) -> tuple[bool, str]:
+    """
+    Validate if the authenticated entity can send as the specified from address.
+    
+    Args:
+        session: SMTP session with authentication info
+        mail_from: The MAIL FROM address
+        
+    Returns:
+        (authorized, message) tuple
+    """
+    if not mail_from:
+        return False, "No sender address provided"
+    
+    # Extract domain from mail_from
+    try:
+        from_domain = mail_from.split('@')[1].lower() if '@' in mail_from else ''
+        if not from_domain:
+            return False, "Invalid sender address format"
+    except (IndexError, AttributeError):
+        return False, "Invalid sender address format"
+    
+    peer_ip = session.peer[0]
+    
+    # Check user authentication
+    if hasattr(session, 'authenticated_user') and session.authenticated_user:
+        user = session.authenticated_user
+        
+        if user.can_send_as(mail_from):
+            logger.info(f"User {user.email} authorized to send as {mail_from}")
+            return True, f"User authorized to send as {mail_from}"
+        else:
+            message = f"User {user.email} not authorized to send as {mail_from}"
+            logger.warning(message)
+            log_auth_attempt(
+                auth_type='sender_validation',
+                identifier=f"{user.email} -> {mail_from}",
+                ip_address=peer_ip,
+                success=False,
+                message=message
+            )
+            return False, message
+    
+    # Check IP authentication for domain
+    authenticator = EnhancedIPAuthenticator()
+    can_auth, auth_message = authenticator.can_authenticate_for_domain(peer_ip, from_domain)
+    
+    if can_auth:
+        # Store IP auth info in session
+        session.auth_type = 'ip'
+        session.authorized_domain = from_domain
+        
+        log_auth_attempt(
+            auth_type='ip',
+            identifier=f"{peer_ip} -> {from_domain}",
+            ip_address=peer_ip,
+            success=True,
+            message=f"IP authorized for domain {from_domain}"
+        )
+        
+        logger.info(f"IP {peer_ip} authorized to send for domain {from_domain}")
+        return True, f"IP authorized for domain {from_domain}"
+    else:
+        log_auth_attempt(
+            auth_type='ip',
+            identifier=f"{peer_ip} -> {from_domain}",
+            ip_address=peer_ip,
+            success=False,
+            message=auth_message
+        )
+        
+        logger.warning(f"IP {peer_ip} not authorized for domain {from_domain}: {auth_message}")
+        return False, f"Not authorized to send for domain {from_domain}"
+
+def get_authenticated_domain_id(session) -> int:
+    """
+    Get the domain ID for the authenticated entity.
+    
+    Args:
+        session: SMTP session with authentication info
+        
+    Returns:
+        Domain ID or None if not authenticated
+    """
+    if hasattr(session, 'authenticated_user') and session.authenticated_user:
+        return session.authenticated_user.domain_id
+    
+    if hasattr(session, 'authorized_domain') and session.authorized_domain:
+        domain = get_domain_by_name(session.authorized_domain)
+        return domain.id if domain else None
+    
+    return None

email_server/cli_tools.py

@@ -9,7 +9,7 @@ from email_server.tool_box import get_logger
 
 logger = get_logger()
 
-def add_domain(domain_name, requires_auth=True):
+def add_domain(domain_name):
     """Add a new domain to the database."""
     session = Session()
     try:
@@ -18,7 +18,7 @@ def add_domain(domain_name, requires_auth=True):
             print(f"Domain {domain_name} already exists")
             return False
         
-        domain = Domain(domain_name=domain_name, requires_auth=requires_auth)
+        domain = Domain(domain_name=domain_name)
         session.add(domain)
         session.commit()
         print(f"Added domain: {domain_name}")

email_server/dkim_manager.py

@@ -13,7 +13,6 @@ 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()

email_server/models.py

@@ -1,9 +1,15 @@
 """
-Database models for the SMTP server.
+Database models for the SMTP server using ESRV schema.
+
+Enhanced security features:
+- Users can only send as their own email or domain emails (if permitted)
+- IP authentication is domain-specific
+- All tables use 'esrv_' prefix for namespace isolation
 """
 
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, Boolean
-from sqlalchemy.orm import declarative_base, sessionmaker
+from sqlalchemy.orm import declarative_base, sessionmaker, relationship
+from sqlalchemy.sql import func
 from datetime import datetime
 import bcrypt
 from email_server.settings_loader import load_settings
@@ -22,27 +28,109 @@ Session = sessionmaker(bind=engine)
 logger = get_logger()
 
 class Domain(Base):
-    __tablename__ = 'domains'
+    """Domain model with enhanced security features."""
+    __tablename__ = 'esrv_domains'
+    
     id = Column(Integer, primary_key=True)
     domain_name = Column(String, unique=True, nullable=False)
-    requires_auth = Column(Boolean, default=True)
+    is_active = Column(Boolean, default=True)
+    created_at = Column(DateTime, default=func.now())
+    
+    def __repr__(self):
+        return f"<Domain(id={self.id}, domain_name='{self.domain_name}', active={self.is_active})>"
 
 class User(Base):
-    __tablename__ = 'users'
+    """
+    User model with enhanced authentication controls.
+    
+    Security features:
+    - can_send_as_domain: If True, user can send as any email from their domain
+    - If False, user can only send as their own email address
+    """
+    __tablename__ = 'esrv_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)
+    can_send_as_domain = Column(Boolean, default=False)
+    is_active = Column(Boolean, default=True)
+    created_at = Column(DateTime, default=func.now())
+    
+    def can_send_as(self, from_address: str) -> bool:
+        """
+        Check if this user can send emails as the given from_address.
+        
+        Args:
+            from_address: The email address the user wants to send from
+            
+        Returns:
+            True if user is allowed to send as this address
+        """
+        # User can always send as their own email
+        if from_address.lower() == self.email.lower():
+            return True
+            
+        # If user has domain privileges, check if from_address is from same domain
+        if self.can_send_as_domain:
+            user_domain = self.email.split('@')[1].lower()
+            from_domain = from_address.split('@')[1].lower() if '@' in from_address else ''
+            return user_domain == from_domain
+            
+        return False
+    
+    def __repr__(self):
+        return f"<User(id={self.id}, email='{self.email}', domain_id={self.domain_id}, can_send_as_domain={self.can_send_as_domain})>"
 
 class WhitelistedIP(Base):
-    __tablename__ = 'whitelisted_ips'
+    """
+    IP whitelist model with domain-specific authentication.
+    
+    Security feature:
+    - IPs can only send emails for their specific authorized domain
+    """
+    __tablename__ = 'esrv_whitelisted_ips'
+    
     id = Column(Integer, primary_key=True)
-    ip_address = Column(String, unique=True, nullable=False)
+    ip_address = Column(String, nullable=False)
     domain_id = Column(Integer, nullable=False)
+    is_active = Column(Boolean, default=True)
+    created_at = Column(DateTime, default=func.now())
+    
+    def can_send_for_domain(self, domain_name: str) -> bool:
+        """
+        Check if this IP can send emails for the given domain.
+        
+        Args:
+            domain_name: The domain name to check
+            
+        Returns:
+            True if IP is authorized for this domain
+        """
+        if not self.is_active:
+            return False
+            
+        # Need to check against the actual domain
+        session = Session()
+        try:
+            domain = session.query(Domain).filter_by(
+                domain_name=domain_name.lower(),
+                is_active=True
+            ).first()
+            return domain and domain.id == self.domain_id
+        finally:
+            session.close()
+    
+    def __repr__(self):
+        return f"<WhitelistedIP(id={self.id}, ip='{self.ip_address}', domain_id={self.domain_id})>"
 
 class EmailLog(Base):
-    __tablename__ = 'email_logs'
+    """Email log model for tracking sent emails."""
+    __tablename__ = 'esrv_email_logs'
+    
     id = Column(Integer, primary_key=True)
+    
+    # Legacy columns (from original schema)
     message_id = Column(String, unique=True, nullable=False)
     timestamp = Column(DateTime, nullable=False)
     peer = Column(String, nullable=False)
@@ -51,43 +139,179 @@ class EmailLog(Base):
     content = Column(Text, nullable=False)
     status = Column(String, nullable=False)
     dkim_signed = Column(Boolean, default=False)
+    
+    # New columns (added later)
+    from_address = Column(String, nullable=False, server_default='unknown')
+    to_address = Column(String, nullable=False, server_default='unknown')
+    subject = Column(Text, nullable=True)
+    message = Column(Text, nullable=True)
+    created_at = Column(DateTime, default=func.now())
+    
+    def __repr__(self):
+        return f"<EmailLog(id={self.id}, message_id='{self.message_id}', from='{self.mail_from}', to='{self.rcpt_tos}', status='{self.status}')>"
 
 class AuthLog(Base):
-    __tablename__ = 'auth_logs'
+    """Authentication log model for security auditing."""
+    __tablename__ = 'esrv_auth_logs'
+    
     id = Column(Integer, primary_key=True)
-    timestamp = Column(DateTime, nullable=False)
-    peer = Column(String, nullable=False)
-    username = Column(String)
+    auth_type = Column(String, nullable=False)  # 'user' or 'ip'
+    identifier = Column(String, nullable=False)  # email or IP address
+    ip_address = Column(String)
     success = Column(Boolean, nullable=False)
-    message = Column(String, nullable=False)
+    message = Column(Text)
+    created_at = Column(DateTime, default=func.now())
+    
+    def __repr__(self):
+        return f"<AuthLog(id={self.id}, type='{self.auth_type}', identifier='{self.identifier}', success={self.success})>"
 
 class DKIMKey(Base):
-    __tablename__ = 'dkim_keys'
+    """DKIM key model for email signing."""
+    __tablename__ = 'esrv_dkim_keys'
+    
     id = Column(Integer, primary_key=True)
     domain_id = Column(Integer, nullable=False)
-    selector = Column(String, nullable=False)
+    selector = Column(String, nullable=False, default='default')
     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)
+    created_at = Column(DateTime, default=func.now())
+    
+    def __repr__(self):
+        return f"<DKIMKey(id={self.id}, domain_id={self.domain_id}, selector='{self.selector}', active={self.is_active})>"
 
 class CustomHeader(Base):
-    """Model for storing custom headers per domain."""
-    __tablename__ = 'custom_headers'
+    """Custom header model for domain-specific email headers."""
+    __tablename__ = 'esrv_custom_headers'
+    
     id = Column(Integer, primary_key=True)
     domain_id = Column(Integer, nullable=False)
     header_name = Column(String, nullable=False)
     header_value = Column(String, nullable=False)
     is_active = Column(Boolean, default=True)
+    created_at = Column(DateTime, default=func.now())
+    
+    def __repr__(self):
+        return f"<CustomHeader(id={self.id}, domain_id={self.domain_id}, header='{self.header_name}: {self.header_value}', active={self.is_active})>"
+
 
 def create_tables():
-    """Create all database tables."""
+    """Create all database tables using ESRV schema."""
     Base.metadata.create_all(engine)
+    logger.info("Created ESRV database tables")
 
-def hash_password(password):
+def hash_password(password: str) -> str:
     """Hash a password using bcrypt."""
     return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
 
-def check_password(password, hashed):
+def check_password(password: str, hashed: str) -> bool:
     """Check a password against its hash."""
     return bcrypt.checkpw(password.encode('utf-8'), hashed.encode('utf-8'))
+
+def log_auth_attempt(auth_type: str, identifier: str, ip_address: str, 
+                     success: bool, message: str) -> None:
+    """
+    Log an authentication attempt for security auditing.
+    
+    Args:
+        auth_type: Type of auth ('user' or 'ip')
+        identifier: User email or IP address
+        ip_address: Client IP address
+        success: Whether auth was successful
+        message: Additional details
+    """
+    session = Session()
+    try:
+        auth_log = AuthLog(
+            auth_type=auth_type,
+            identifier=identifier,
+            ip_address=ip_address,
+            success=success,
+            message=message
+        )
+        session.add(auth_log)
+        session.commit()
+        logger.info(f"Auth log: {auth_type} {identifier} from {ip_address} - {'SUCCESS' if success else 'FAILED'}")
+    except Exception as e:
+        session.rollback()
+        logger.error(f"Failed to log auth attempt: {e}")
+    finally:
+        session.close()
+
+def log_email(from_address: str, to_address: str, subject: str, 
+              status: str, message: str = None) -> None:
+    """
+    Log an email send attempt.
+    
+    Args:
+        from_address: Sender email
+        to_address: Recipient email  
+        subject: Email subject
+        status: Send status ('sent', 'failed', etc.)
+        message: Additional details
+    """
+    session = Session()
+    try:
+        email_log = EmailLog(
+            from_address=from_address,
+            to_address=to_address,
+            subject=subject,
+            status=status,
+            message=message
+        )
+        session.add(email_log)
+        session.commit()
+        logger.info(f"Email log: {from_address} -> {to_address} - {status}")
+    except Exception as e:
+        session.rollback()
+        logger.error(f"Failed to log email: {e}")
+    finally:
+        session.close()
+
+def get_user_by_email(email: str):
+    """Get user by email address."""
+    session = Session()
+    try:
+        return session.query(User).filter_by(email=email.lower(), is_active=True).first()
+    finally:
+        session.close()
+
+def get_domain_by_name(domain_name: str):
+    """Get domain by name."""
+    session = Session()
+    try:
+        return session.query(Domain).filter_by(domain_name=domain_name.lower(), is_active=True).first()
+    finally:
+        session.close()
+
+def get_whitelisted_ip(ip_address: str, domain_name: str = None):
+    """
+    Get whitelisted IP, optionally filtered by domain.
+    
+    Args:
+        ip_address: IP address to check
+        domain_name: Optional domain name to restrict to
+        
+    Returns:
+        WhitelistedIP object if found and authorized for domain
+    """
+    session = Session()
+    try:
+        query = session.query(WhitelistedIP).filter_by(
+            ip_address=ip_address,
+            is_active=True
+        )
+        
+        if domain_name:
+            # Join with domain to check authorization
+            domain = session.query(Domain).filter_by(
+                domain_name=domain_name.lower(),
+                is_active=True
+            ).first()
+            if not domain:
+                return None
+            query = query.filter_by(domain_id=domain.id)
+            
+        return query.first()
+    finally:
+        session.close()

email_server/server_runner.py

@@ -9,7 +9,7 @@ from email_server.tool_box import get_logger
 
 # Import our modules
 from email_server.models import create_tables
-from email_server.smtp_handler import CustomSMTPHandler, PlainController
+from email_server.smtp_handler import EnhancedCustomSMTPHandler, PlainController
 from email_server.tls_utils import generate_self_signed_cert, create_ssl_context
 from email_server.dkim_manager import DKIMManager
 from aiosmtpd.controller import Controller
@@ -52,7 +52,7 @@ async def start_server():
         # 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)
+            domain = Domain(domain_name='example.com')
             session.add(domain)
             session.commit()
             logger.debug("Added example.com domain")
@@ -95,7 +95,7 @@ async def start_server():
         return
     
     # Start plain SMTP server (with IP whitelist fallback)
-    handler_plain = CustomSMTPHandler()
+    handler_plain = EnhancedCustomSMTPHandler()
     controller_plain = PlainController(
         handler_plain,
         hostname=BIND_IP,
@@ -106,7 +106,7 @@ async def start_server():
     logger.debug(f'Starting plain SMTP server on {HOSTNAME}:{SMTP_PORT}...')
     
     # Start TLS SMTP server using closure pattern like the original
-    handler_tls = CustomSMTPHandler()
+    handler_tls = EnhancedCustomSMTPHandler()
     
     # Define TLS controller class with ssl_context in closure (like original)
     class TLSController(Controller):

email_server/settings_loader.py

@@ -31,7 +31,6 @@ DEFAULTS = {
         'TLS_KEY_FILE': 'email_server/ssl_certs/server.key',
     },
     'DKIM': {
-        'DKIM_SELECTOR': 'default',
         'DKIM_KEY_SIZE': '2048',
     },
 }

email_server/smtp_handler.py

@@ -1,24 +1,38 @@
 """
-SMTP handler for processing incoming emails.
+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 Authenticator, IPAuthenticator
+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 CombinedAuthenticator:
-    """Combined authenticator that tries username/password first, then falls back to IP whitelist."""
+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 = Authenticator()
-        self.ip_auth = IPAuthenticator()
+        self.user_auth = EnhancedAuthenticator()
+        self.ip_auth = EnhancedIPAuthenticator()
     
     def __call__(self, server, session, envelope, mechanism, auth_data):
         from aiosmtpd.smtp import LoginPassword
@@ -31,21 +45,65 @@ class CombinedAuthenticator:
             # 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)
+        # 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 CustomSMTPHandler:
-    """Custom SMTP handler for processing emails."""
+class EnhancedCustomSMTPHandler:
+    """Enhanced custom SMTP handler with security controls."""
     
     def __init__(self):
-        self.authenticator = Authenticator()
-        self.ip_authenticator = IPAuthenticator()
-        self.combined_authenticator = CombinedAuthenticator()
+        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:
@@ -61,6 +119,9 @@ class CustomSMTPHandler:
             # 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)
@@ -114,9 +175,25 @@ class CustomSMTPHandler:
         return '250 OK'
 
     async def handle_MAIL(self, server, session, envelope, address, mail_options):
-        """Handle MAIL FROM command - validate sender."""
+        """
+        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):

main.py

@@ -0,0 +1,17 @@
+from email_server.server_runner import start_server
+from email_server.tool_box import get_logger
+import asyncio
+import sys
+
+logger = get_logger()
+
+if __name__ == '__main__':
+    try:
+        logger.info('Server started')
+        asyncio.run(start_server())
+    except KeyboardInterrupt:
+        logger.info('Server interrupted by user')
+        sys.exit(0)
+    except Exception as e:
+        logger.error(f'Server error: {e}')
+        sys.exit(1)

start_server.py

@@ -1,17 +0,0 @@
-from email_server.server_runner import start_server
-from email_server.tool_box import get_logger
-import asyncio
-import sys
-
-logger = get_logger()
-
-if __name__ == '__main__':
-    try:
-        logger.info('Server started')
-        asyncio.run(start_server())
-    except KeyboardInterrupt:
-        logger.info('Server interrupted by user')
-        sys.exit(0)
-    except Exception as e:
-        logger.error(f'Server error: {e}')
-        sys.exit(1)