cleanup.py

#!/usr/bin/env python3
"""
ProtonMail Folder Cleanup Script
Connects via IMAP Bridge and deletes emails older than retention period
"""

import imaplib
import email
from email.header import decode_header, make_header
from datetime import datetime, timedelta, timezone
from email.utils import parsedate_to_datetime
import logging
from typing import Dict, Any, List, Tuple
import os
import json
import re
import socket
import time
from ssl import create_default_context, SSLContext, CERT_NONE

# Configuration
DEFAULT_CONFIG_FILE = '/app/config.json'
CONFIG_FILE = os.environ.get('CONFIG_FILE', DEFAULT_CONFIG_FILE)
LOG_FILE = os.environ.get('LOG_FILE', '/app/logs/cleanup.log')
PREFIX_EXCLUDE = {name.lower() for name in ["Spam", "Trash", "Starred", "Sent", "Drafts"]}
DEFAULT_IMAP_TIMEOUT = 60
DEFAULT_IMAP_RETRIES = 5
DEFAULT_IMAP_RETRY_DELAY = 10

def configure_logging(config: Dict[str, Any]):
    """Configure logging to console and optionally file."""
    root = logging.getLogger()
    root.handlers.clear()

    log_to_file = config.get('log_to_file', True)
    log_file_path = config.get('log_file', LOG_FILE)

    handlers = [logging.StreamHandler()]

    if log_to_file:
        try:
            os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
            handlers.append(logging.FileHandler(log_file_path, mode='a'))
        except OSError as e:
            logging.basicConfig(
                level=logging.INFO,
                format='%(asctime)s - %(levelname)s - %(message)s',
                handlers=handlers
            )
            logging.warning(
                "File logging disabled (could not write to %s): %s",
                log_file_path,
                e,
            )
            return

    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=handlers
    )

def validate_config(config: Dict[str, Any]) -> None:
    """Validate configuration structure and values."""
    required_imap_keys = {'server', 'port', 'email', 'password'}

    if 'imap' not in config:
        raise ValueError("Config missing 'imap' section")

    missing_imap = required_imap_keys - set(config['imap'].keys())
    if missing_imap:
        raise ValueError(f"Config missing IMAP keys: {', '.join(sorted(missing_imap))}")

    port = config['imap'].get('port')
    if not isinstance(port, int) or port <= 0:
        raise ValueError("imap.port must be a positive integer")

    folders = config.get('folders')
    if not isinstance(folders, list) or not folders:
        raise ValueError("Config 'folders' must be a non-empty list")

    for idx, folder in enumerate(folders):
        if not isinstance(folder, dict):
            raise ValueError(f"Folder entry at index {idx} must be an object")
        if 'name' not in folder or not folder['name']:
            raise ValueError(f"Folder entry at index {idx} missing 'name'")
        retention = folder.get('retention_days')
        if retention is not None:
            if not isinstance(retention, int) or retention <= 0:
                raise ValueError(
                    f"Folder '{folder.get('name', idx)}' retention_days must be a positive integer or null"
                )
        if 'dry_run' in folder and not isinstance(folder['dry_run'], bool):
            raise ValueError(
                f"Folder '{folder.get('name', idx)}' dry_run must be boolean when provided"
            )
        if 'sender_rules' in folder:
            validate_sender_rules(folder['sender_rules'], scope=f"folder '{folder['name']}'")
    if 'log_to_file' in config and not isinstance(config['log_to_file'], bool):
        raise ValueError("log_to_file must be boolean when provided")
    if 'log_mailboxes' in config and not isinstance(config['log_mailboxes'], bool):
        raise ValueError("log_mailboxes must be boolean when provided")
    if 'folder_prefix' in config and not isinstance(config['folder_prefix'], str):
        raise ValueError("folder_prefix must be a string when provided")
    if 'sender_rules' in config:
        validate_sender_rules(config['sender_rules'], scope="global")
    if 'imap_timeout_seconds' in config and not isinstance(config['imap_timeout_seconds'], (int, float)):
        raise ValueError("imap_timeout_seconds must be numeric when provided")
    if 'starttls' in config.get('imap', {}) and not isinstance(config['imap']['starttls'], bool):
        raise ValueError("imap.starttls must be boolean when provided")
    if 'starttls_insecure_skip_verify' in config.get('imap', {}) and not isinstance(config['imap']['starttls_insecure_skip_verify'], bool):
        raise ValueError("imap.starttls_insecure_skip_verify must be boolean when provided")
    if 'username' in config.get('imap', {}) and config['imap']['username'] is not None and not isinstance(config['imap']['username'], str):
        raise ValueError("imap.username must be a string when provided")
    if 'imap_connect_retries' in config and not isinstance(config['imap_connect_retries'], int):
        raise ValueError("imap_connect_retries must be an integer when provided")
    if 'imap_connect_retry_delay_seconds' in config and not isinstance(config['imap_connect_retry_delay_seconds'], (int, float)):
        raise ValueError("imap_connect_retry_delay_seconds must be numeric when provided")

def validate_sender_rules(rules: Any, scope: str) -> None:
    if not isinstance(rules, list):
        raise ValueError(f"{scope} sender_rules must be a list")
    for idx, rule in enumerate(rules):
        if not isinstance(rule, dict):
            raise ValueError(f"{scope} sender_rules entry at index {idx} must be an object")
        from_val = rule.get('from_contains')
        subject_val = rule.get('subject_contains')
        from_regex = rule.get('from_regex')
        subject_regex = rule.get('subject_regex')
        if not any([from_val, subject_val, from_regex, subject_regex]):
            raise ValueError(
                f"{scope} sender_rules entry at index {idx} must include from_contains, subject_contains, from_regex, or subject_regex"
            )
        if from_val is not None and not isinstance(from_val, str):
            raise ValueError(f"{scope} sender_rules entry at index {idx} from_contains must be a string when provided")
        if subject_val is not None and not isinstance(subject_val, str):
            raise ValueError(f"{scope} sender_rules entry at index {idx} subject_contains must be a string when provided")
        if from_regex is not None and not isinstance(from_regex, str):
            raise ValueError(f"{scope} sender_rules entry at index {idx} from_regex must be a string when provided")
        if subject_regex is not None and not isinstance(subject_regex, str):
            raise ValueError(f"{scope} sender_rules entry at index {idx} subject_regex must be a string when provided")
        retention = rule.get('retention_days')
        if not isinstance(retention, int) or retention <= 0:
            raise ValueError(
                f"{scope} sender_rules entry at index {idx} retention_days must be a positive integer"
            )
        # Validate regex patterns if present
        for pattern, label in [
            (from_regex, "from_regex"),
            (subject_regex, "subject_regex"),
        ]:
            if pattern:
                try:
                    re.compile(pattern)
                except re.error as e:
                    raise ValueError(f"{scope} sender_rules entry at index {idx} invalid {label}: {e}") from e

def load_config():
    """Load configuration from JSON file"""
    try:
        config_path = CONFIG_FILE
        if not os.path.exists(config_path) and config_path == DEFAULT_CONFIG_FILE:
            local_path = os.path.join(os.getcwd(), 'config.json')
            if os.path.exists(local_path):
                config_path = local_path
                logging.info(f"Default config not found; using local {config_path}")
        with open(config_path, 'r') as f:
            config = json.load(f)
        validate_config(config)
        logging.info(f"Configuration loaded from {config_path}")
        return config
    except Exception as e:
        logging.error(f"Failed to load config: {e}")
        raise

def connect_imap(config):
    """Connect to ProtonMail via IMAP Bridge"""
    imap_config = config['imap']
    username = imap_config.get('username') or imap_config.get('email')
    server = imap_config['server']
    port = imap_config['port']
    use_starttls = imap_config.get('starttls')
    insecure_tls = imap_config.get('starttls_insecure_skip_verify')
    max_retries = config.get('imap_connect_retries', DEFAULT_IMAP_RETRIES)
    retry_delay = config.get('imap_connect_retry_delay_seconds', DEFAULT_IMAP_RETRY_DELAY)

    for attempt in range(1, max_retries + 1):
        try:
            logging.info(
                "Connecting IMAP (attempt %s/%s) to %s:%s starttls=%s insecure=%s user=%s",
                attempt,
                max_retries,
                server,
                port,
                use_starttls,
                insecure_tls,
                username,
            )
            mail = imaplib.IMAP4(server, port)
            if use_starttls:
                context: SSLContext = create_default_context()
                if insecure_tls:
                    context.check_hostname = False
                    context.verify_mode = CERT_NONE
                    logging.warning("STARTTLS certificate verification disabled (insecure)")
                mail.starttls(ssl_context=context)
                logging.info("IMAP STARTTLS negotiated")
            mail.login(username, imap_config['password'])
            logging.info("Successfully connected to ProtonMail Bridge")
            return mail
        except Exception as e:
            logging.error("Failed to connect (attempt %s/%s): %s", attempt, max_retries, e)
            if attempt == max_retries:
                raise
            time.sleep(retry_delay)

def get_cutoff_date(days: Any) -> datetime:
    """Calculate cutoff date based on retention days"""
    if days is None:
        return None
    return datetime.now(timezone.utc) - timedelta(days=days)

def decode_header_value(value: str) -> str:
    """Decode MIME-encoded headers safely to unicode string."""
    if not value:
        return ""
    return str(make_header(decode_header(value))).replace('\n', ' ').replace('\r', ' ').strip()

def shorten(text: str, max_len: int = 120) -> str:
    """Trim text for log output."""
    return text if len(text) <= max_len else text[:max_len - 3] + "..."

def list_mailboxes(mail):
    """List available mailboxes for debugging."""
    status, mailboxes = mail.list()
    if status != 'OK' or not mailboxes:
        logging.warning("Could not list mailboxes (status: %s, data: %s)", status, mailboxes)
        return
    decoded = []
    for entry in mailboxes:
        try:
            decoded_entry = entry.decode('utf-8', errors='replace')
        except Exception:
            decoded_entry = str(entry)
        decoded.append(decoded_entry)
    logging.info("Available mailboxes:\n%s", "\n".join(decoded[:50]))

def quote_mailbox(name: str) -> str:
    """Quote mailbox name for IMAP SELECT."""
    escaped = name.replace('\\', '\\\\').replace('"', r'\"')
    return f'"{escaped}"'

def format_sample_entry(msg_id: bytes, email_date: datetime, cutoff_date: datetime,
                        from_header: str, subject: str, retention_days: int, rule_source: str) -> str:
    """Format a sample log line for a deletable message."""
    expired_days = (cutoff_date - email_date).days
    return (
        f"id={msg_id.decode('utf-8', 'ignore')} "
        f"from='{shorten(from_header)}' "
        f"subject='{shorten(subject)}' "
        f"{expired_days}d expired; retention={retention_days}d"
        f"{' via ' + rule_source if rule_source else ''}; "
        f"sent_at={email_date.isoformat()}"
    )

def resolve_retention(base_retention: int,
                      folder_rules: List[Dict[str, Any]],
                      global_rules: List[Dict[str, Any]],
                      from_header: str,
                      subject_header: str) -> Tuple[int, str]:
    """Return retention days and rule source label."""
    from_lower = from_header.lower()
    subject_lower = subject_header.lower()

    def match_rule(rule: Dict[str, Any]) -> bool:
        from_match = rule.get('from_contains', '').lower() in from_lower if rule.get('from_contains') else False
        subject_match = rule.get('subject_contains', '').lower() in subject_lower if rule.get('subject_contains') else False
        from_regex = rule.get('from_regex')
        subject_regex = rule.get('subject_regex')
        from_regex_match = bool(re.search(from_regex, from_header, re.IGNORECASE)) if from_regex else False
        subject_regex_match = bool(re.search(subject_regex, subject_header, re.IGNORECASE)) if subject_regex else False
        # Match if any field matches (OR semantics across contains/regex)
        return from_match or subject_match or from_regex_match or subject_regex_match

    for rule in folder_rules:
        if match_rule(rule):
            return rule['retention_days'], "folder_sender_rule"
    for rule in global_rules:
        if match_rule(rule):
            return rule['retention_days'], "global_sender_rule"
    return base_retention, ""

def cleanup_folder(mail, folder_name: str, retention_days: int, 
                   dry_run: bool = False, prefix: str = "",
                   folder_sender_rules: List[Dict[str, Any]] = None,
                   global_sender_rules: List[Dict[str, Any]] = None):
    """Delete emails older than retention period in specified folder"""
    try:
        folder_sender_rules = folder_sender_rules or []
        global_sender_rules = global_sender_rules or []
        # Select the folder
        use_prefix = prefix and folder_name.lower() not in PREFIX_EXCLUDE
        full_name = f"{prefix}{folder_name}" if use_prefix else folder_name
        mailbox = quote_mailbox(full_name)
        status, data = mail.select(mailbox)
        if status != 'OK':
            logging.warning(
                "Could not select folder '%s' (response: %s)",
                full_name,
                data,
            )
            return
        
        logging.info(
            f"Processing folder: {folder_name} "
            f"(retention: {retention_days} days, dry_run: {dry_run})"
        )
        
        # Search for all emails
        status, messages = mail.search(None, 'ALL')
        if status != 'OK':
            logging.warning(f"No messages found in {folder_name}")
            return
        
        message_ids = messages[0].split()
        deleted_count = 0
        sample_first = []
        sample_last = []
        processed = 0
        
        for msg_id in message_ids:
            try:
                # Fetch the email date
                status, msg_data = mail.fetch(
                    msg_id, 
                    '(BODY[HEADER.FIELDS (DATE SUBJECT FROM)])'
                )
                if status != 'OK':
                    continue
                
                # Parse the date
                raw_header_bytes = msg_data[0][1]
                msg = email.message_from_bytes(raw_header_bytes)
                date_str = msg.get('Date')
                subject = decode_header_value(msg.get('Subject'))
                from_header = decode_header_value(msg.get('From'))
                
                if not date_str:
                    continue
                
                email_date = parsedate_to_datetime(date_str)
                # Normalize to timezone-aware UTC for safe comparison
                if email_date.tzinfo is None:
                    email_date = email_date.replace(tzinfo=timezone.utc)
                else:
                    email_date = email_date.astimezone(timezone.utc)
                
                msg_retention, rule_source = resolve_retention(
                    retention_days, folder_sender_rules, global_sender_rules, from_header, subject
                )
                cutoff_date = get_cutoff_date(msg_retention)

                # Delete if older than cutoff
                if cutoff_date is None:
                    continue
                if email_date < cutoff_date:
                    sample_entry = format_sample_entry(
                        msg_id, email_date, cutoff_date, from_header, subject, msg_retention, rule_source
                    )
                    if len(sample_first) < 10:
                        sample_first.append(sample_entry)
                    else:
                        if len(sample_last) == 10:
                            sample_last.pop(0)
                        sample_last.append(sample_entry)

                    if not dry_run:
                        mail.store(msg_id, '+FLAGS', '\\Deleted')
                    deleted_count += 1
                processed += 1
                if processed % 200 == 0:
                    logging.info(
                        "Progress: processed %s messages in %s (deleted candidates: %s)",
                        processed,
                        full_name,
                        deleted_count
                    )
                    
            except Exception as e:
                logging.error(
                    f"Error processing message {msg_id}: {e}"
                )
                continue
        
        # Expunge deleted messages
        if deleted_count > 0:
            # Log sample of deletions to help validate targeting
            if deleted_count <= 20:
                samples = sample_first + sample_last
                logging.info(
                    "Sample deletions (%s total) from %s (retention %sd):\n%s",
                    deleted_count,
                    folder_name,
                    retention_days,
                    "\n".join(samples),
                )
            else:
                logging.info(
                    "Sample deletions (first 10) from %s (retention %sd):\n%s",
                    folder_name,
                    retention_days,
                    "\n".join(sample_first)
                )
                logging.info(
                    "Sample deletions (last 10) from %s (retention %sd):\n%s",
                    folder_name,
                    retention_days,
                    "\n".join(sample_last)
                )

            if not dry_run:
                mail.expunge()
                logging.info(
                    f"Deleted {deleted_count} emails from {folder_name}"
                )
            else:
                logging.info(
                    f"[DRY RUN] Would delete {deleted_count} emails "
                    f"from {folder_name}"
                )
        else:
            logging.info(f"No emails to delete in {folder_name}")
            
    except Exception as e:
        logging.error(f"Error cleaning folder {folder_name}: {e}")

def main():
    """Main execution function"""
    try:
        # Basic console logging before full config to surface config errors
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
        config = load_config()
        configure_logging(config)
        logging.info("Starting ProtonMail cleanup job")
        socket.setdefaulttimeout(config.get('imap_timeout_seconds', DEFAULT_IMAP_TIMEOUT))

        mail = connect_imap(config)
        if config.get('log_mailboxes'):
            list_mailboxes(mail)

        dry_run = config.get('dry_run', False)
        if dry_run:
            logging.info("Running in DRY RUN mode - no emails will be deleted")

        # Process each folder with retention policy
        prefix = config.get('folder_prefix', '')
        global_sender_rules = config.get('sender_rules', [])
        for folder_config in config['folders']:
            folder_name = folder_config['name']
            retention_days = folder_config['retention_days']
            folder_dry_run = folder_config.get('dry_run', dry_run)
            folder_sender_rules = folder_config.get('sender_rules', [])
            if folder_dry_run and not dry_run:
                logging.info(f"Folder '{folder_name}' is in dry_run override mode")
            cleanup_folder(
                mail,
                folder_name,
                retention_days,
                folder_dry_run,
                prefix,
                folder_sender_rules,
                global_sender_rules,
            )

        if config.get('log_mailboxes'):
            list_mailboxes(mail)

        try:
            mail.close()
        except imaplib.IMAP4.error as e:
            logging.warning(f"Skipping close(): {e}")
        mail.logout()
        logging.info("Cleanup job completed successfully")
        
    except Exception as e:
        logging.error(f"Cleanup job failed: {e}")
        raise

if __name__ == "__main__":
    main()