ge-tool/backend/services/aria2/daemon.py

"""
aria2c RPC daemon management

Starts/stops the aria2c daemon with optimized settings for NAS downloads
"""
import os
import subprocess
import time
import logging
from pathlib import Path
from typing import Optional
import psutil

logger = logging.getLogger(__name__)

# Global daemon process
_aria2_process: Optional[subprocess.Popen] = None


def get_aria2_executable() -> str:
    """Get path to aria2c.exe in project folder"""
    project_root = Path(__file__).parent.parent.parent.parent
    aria2_exe = project_root / "aria2" / "aria2c.exe"

    if not aria2_exe.exists():
        raise FileNotFoundError(
            f"aria2c.exe not found at {aria2_exe}. "
            f"Please ensure aria2 folder exists in project root."
        )

    return str(aria2_exe)


def is_aria2_running(port: int = 6800) -> bool:
    """
    Check if aria2c RPC server is already running on specified port

    Args:
        port: RPC port to check (default: 6800)

    Returns:
        True if aria2c is running on the port
    """
    try:
        for conn in psutil.net_connections():
            # Check if laddr exists and has port attribute
            if hasattr(conn, 'laddr') and conn.laddr and hasattr(conn.laddr, 'port'):
                if conn.laddr.port == port and conn.status == 'LISTEN':
                    return True
        return False
    except Exception as e:
        logger.warning(f"Could not check if aria2 is running: {e}")
        return False


def start_aria2_daemon(
    port: int = 6800,
    secret: str = "dkidownload_secret_2025",
    max_concurrent: int = 10,
    max_connections: Optional[int] = None,
    split: Optional[int] = None,
    min_split_size: str = "1M",
    download_dir: Optional[str] = None
) -> bool:
    """
    Start aria2c RPC daemon with optimized settings

    Args:
        port: RPC port (default: 6800)
        secret: RPC secret token for authentication
        max_concurrent: Max concurrent downloads
        max_connections: Max connections per server (default: from env ARIA2_MAX_CONNECTIONS_PER_FILE or 16)
        split: Number of connections per file (default: same as max_connections)
        min_split_size: Minimum size to split (e.g., "1M")
        download_dir: Temporary download directory

    Returns:
        True if started successfully
    """
    global _aria2_process

    # Load max_connections from environment if not provided
    if max_connections is None:
        max_connections = int(
            os.getenv('ARIA2_MAX_CONNECTIONS_PER_FILE', '16'))

    # Default split to max_connections if not specified
    if split is None:
        split = max_connections

    # Check if already running
    if is_aria2_running(port):
        logger.debug(f"aria2c already running on port {port}")
        return True

    try:
        aria2_exe = get_aria2_executable()

        # Default download dir: project_root/aria2/downloads
        if download_dir is None:
            project_root = Path(__file__).parent.parent.parent.parent
            download_dir = str(project_root / "aria2" / "downloads")

        # Create download dir if not exists
        os.makedirs(download_dir, exist_ok=True)

        # aria2c command with optimized settings
        cmd = [
            aria2_exe,
            '--enable-rpc',
            '--rpc-listen-all=false',  # Only localhost for security
            f'--rpc-listen-port={port}',
            f'--rpc-secret={secret}',
            f'--max-concurrent-downloads={max_concurrent}',
            f'--max-connection-per-server={max_connections}',
            f'--split={split}',
            f'--min-split-size={min_split_size}',
            '--continue=true',  # Resume support
            '--auto-file-renaming=false',  # Don't auto-rename
            f'--dir={download_dir}',
            '--log=-',  # Log to stdout
            '--log-level=notice',
            '--console-log-level=warn',
            '--summary-interval=0',  # Disable periodic summary
            '--disable-ipv6=true',  # Faster connection
            '--check-certificate=false',  # NAS uses self-signed cert
        ]

        logger.debug(f"Starting aria2c daemon on port {port}...")
        logger.debug(f"Download directory: {download_dir}")

        # Start process (detached, no window on Windows)
        _aria2_process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0
        )

        # Wait a bit and check if started
        time.sleep(1)

        if _aria2_process.poll() is not None:
            # Process died immediately
            stdout, stderr = _aria2_process.communicate()
            logger.error(f"aria2c failed to start: {stderr.decode()}")
            return False

        # Verify it's listening
        if is_aria2_running(port):
            logger.debug(
                f"✅ aria2c daemon started successfully on port {port}")
            return True
        else:
            logger.error("aria2c started but not listening on port")
            return False

    except Exception as e:
        logger.error(f"Failed to start aria2c: {e}", exc_info=True)
        return False


def stop_aria2_daemon() -> bool:
    """
    Stop the aria2c daemon gracefully

    Returns:
        True if stopped successfully
    """
    global _aria2_process

    if _aria2_process is None:
        logger.debug("No aria2c process to stop")
        return True

    try:
        logger.debug("Stopping aria2c daemon...")
        _aria2_process.terminate()

        # Wait up to 5 seconds for graceful shutdown
        try:
            _aria2_process.wait(timeout=5)
            logger.debug("✅ aria2c daemon stopped")
        except subprocess.TimeoutExpired:
            logger.warning("aria2c didn't stop gracefully, forcing...")
            _aria2_process.kill()
            _aria2_process.wait()
            logger.debug("✅ aria2c daemon force-killed")

        _aria2_process = None
        return True

    except Exception as e:
        logger.error(f"Error stopping aria2c: {e}")
        return False


def restart_aria2_daemon(**kwargs) -> bool:
    """
    Restart aria2c daemon with new settings

    Args:
        **kwargs: Arguments to pass to start_aria2_daemon()

    Returns:
        True if restarted successfully
    """
    stop_aria2_daemon()
    time.sleep(1)
    return start_aria2_daemon(**kwargs)