ddns-service/src/ddns_service/cleanup.py

"""TTL cleanup functionality."""

import logging
import threading

from . import now_utc
from .models import DatabaseError, Hostname, User
from datetime import timedelta


def cleanup_expired(app, start_time=None):
    """
    Clean up expired hostnames and return count of cleaned entries.

    Args:
        app: Application instance with dns_service and email_service.
        start_time: Timezone aware datetime object containg the start time of the cleanup thread.

    Returns:
        Number of expired hostnames processed.
    """
    now = now_utc()
    expired_count = 0

    for hostname in Hostname.select(Hostname, User).join(User).where(
            (Hostname.expiry_ttl != 0) &
            ((Hostname.last_ipv4.is_null(False) & Hostname.last_ipv4_update.is_null(False)) |
             (Hostname.last_ipv6.is_null(False) & Hostname.last_ipv6_update.is_null(False)))):

        ipv4_expired = False
        ipv6_expired = False

        if hostname.last_ipv4:
            last_update = max(hostname.last_ipv4_update, start_time) if start_time \
                else hostname.last_ipv4_update
            expiry_time = last_update + timedelta(seconds=hostname.expiry_ttl)
            if now > expiry_time:
                ipv4_expired = True

        if hostname.last_ipv6:
            last_update = max(hostname.last_ipv6_update, start_time) if start_time \
                else hostname.last_ipv6_update
            expiry_time = last_update + timedelta(seconds=hostname.expiry_ttl)
            if now > expiry_time:
                ipv6_expired = True

        if not ipv4_expired and not ipv6_expired:
            continue

        old_ipv4 = hostname.last_ipv4
        old_ipv6 = hostname.last_ipv6
        ipv4_dns_deleted = False
        ipv6_dns_deleted = False

        if ipv4_expired:
            logging.info(
                f"Cleanup: Host expired: hostname={hostname.hostname} zone={hostname.zone} "
                f"ipv4={hostname.last_ipv4}"
            )
            try:
                ipv4_exists = app.dns_service.query_record(
                    hostname.hostname, hostname.zone, "A")
                if ipv4_exists:
                    app.dns_service.delete_record(
                        hostname.hostname, hostname.zone, "A")
                    ipv4_dns_deleted = True
                hostname.last_ipv4 = None
            except Exception as e:
                logging.error(f"DNS error: {e}")
                logging.error(
                    f"Cleanup failed: hostname={hostname.hostname} "
                    f"zone={hostname.zone} type=A"
                )

        if ipv6_expired:
            logging.info(
                f"Cleanup: Host expired: hostname={hostname.hostname} zone={hostname.zone} "
                f"ipv6={hostname.last_ipv6}"
            )
            try:
                ipv6_exists = app.dns_service.query_record(
                    hostname.hostname, hostname.zone, "AAAA")
                if ipv6_exists:
                    app.dns_service.delete_record(
                        hostname.hostname, hostname.zone, "AAAA")
                    ipv6_dns_deleted = True
                hostname.last_ipv6 = None
            except Exception as e:
                logging.error(f"DNS error: {e}")
                logging.error(
                    f"Cleanup failed: hostname={hostname.hostname} "
                    f"zone={hostname.zone} type=AAAA"
                )

        if hostname.last_ipv4 == old_ipv4 and hostname.last_ipv6 == old_ipv6:
            continue

        try:
            hostname.save()
        except DatabaseError as e:
            logging.error(
                f"DB save failed after retries: hostname={hostname.hostname} "
                f"zone={hostname.zone}: {e}"
            )

            # Rollback: re-add DNS records that were deleted
            if ipv4_dns_deleted:
                try:
                    app.dns_service.update_record(
                        hostname.hostname, hostname.zone,
                        old_ipv4, hostname.dns_ttl)
                except Exception as e:
                    logging.error(f"DNS rollback failed (A): {e}")

            if ipv6_dns_deleted:
                try:
                    app.dns_service.update_record(
                        hostname.hostname, hostname.zone,
                        old_ipv6, hostname.dns_ttl)
                except Exception as e:
                    logging.error(f"DNS rollback failed (AAAA): {e}")

            continue

        if app.email_service:
            # Restore old IPs on in-memory model for email template
            hostname.last_ipv4 = old_ipv4
            hostname.last_ipv6 = old_ipv6
            app.email_service.send_expiry_notification(
                hostname.user.email,
                hostname,
                hostname.last_ipv4 == old_ipv4,
                hostname.last_ipv6 == old_ipv6
            )

        expired_count += 1

    return expired_count


class ExpiredRecordsCleanupThread(threading.Thread):
    """Background thread for periodic expired records cleanup."""

    def __init__(self, app):
        """
        Initialize expired records cleanup thread.

        Args:
            app: Application instance.
        """
        super().__init__(daemon=True)
        self.app = app
        self.interval = app.config["dns_service"]["cleanup_interval"]
        self.stop_event = threading.Event()

    def run(self):
        """Run the cleanup loop."""
        logging.info(f"Expired records cleanup thread started: interval={self.interval}s")
        start_time = now_utc()

        while not self.stop_event.wait(self.interval):
            try:
                count = cleanup_expired(self.app, start_time)
                if count > 0:
                    logging.info(f"Expired records cleanup completed: count={count}")
            except Exception as e:
                logging.error(f"Expired records cleanup error: {e}")

    def stop(self):
        """Signal the thread to stop."""
        self.stop_event.set()


class RateLimitCleanupThread(threading.Thread):
    """Background thread for periodic rate limiter cleanup."""

    def __init__(self, app):
        """
        Initialize rate limiter cleanup thread.

        Args:
            app: Application instance.
        """
        super().__init__(daemon=True)
        self.app = app
        self.interval = app.config["rate_limit"]["cleanup_interval"]
        self.stop_event = threading.Event()

    def run(self):
        """Run the cleanup loop."""
        logging.info(f"Rate limit cleanup thread started: interval={self.interval}s")

        while not self.stop_event.wait(self.interval):
            try:
                if self.app.good_limiter:
                    self.app.good_limiter.cleanup()
                if self.app.bad_limiter:
                    self.app.bad_limiter.cleanup()
            except Exception as e:
                logging.error(f"Rate limit cleanup error: {e}")

    def stop(self):
        """Signal the thread to stop."""
        self.stop_event.set()