ytptube/app/routes/api/system.py

import asyncio
import errno
import logging
import os
import shlex
import time
from typing import TYPE_CHECKING

from aiohttp import web
from aiohttp.web import Request, Response
from aiohttp.web_runner import GracefulExit

from app.library.config import Config
from app.library.downloads import DownloadQueue
from app.library.encoder import Encoder
from app.library.Events import EventBus, Events
from app.library.router import route
from app.library.UpdateChecker import UpdateChecker

if TYPE_CHECKING:
    from asyncio import Task
    from asyncio.events import AbstractEventLoop
    from asyncio.subprocess import Process

LOG: logging.Logger = logging.getLogger(__name__)


@route("POST", "api/system/pause", "system.pause")
async def downloads_pause(queue: DownloadQueue, encoder: Encoder, notify: EventBus) -> Response:
    """
    Pause non-active downloads.

    Args:
        request (Request): The request object.
        queue (DownloadQueue): The download queue instance.
        encoder (Encoder): The encoder instance.
        notify (EventBus): The event bus instance.

    Returns:
        Response: The response object.

    """
    if queue.is_paused():
        return web.json_response(
            {"message": "Non-active downloads are already paused."},
            status=web.HTTPNotAcceptable.status_code,
            dumps=encoder.encode,
        )

    queue.pause()

    msg = "Non-active downloads have been paused."
    notify.emit(
        Events.PAUSED,
        data={"paused": True, "at": time.time()},
        title="Downloads Paused",
        message=msg,
    )

    return web.json_response(data={"message": msg}, status=web.HTTPOk.status_code, dumps=encoder.encode)


@route("POST", "api/system/resume", "system.resume")
async def downloads_resume(queue: DownloadQueue, encoder: Encoder, notify: EventBus) -> Response:
    """
    Resume non-active downloads.

    Args:
        request (Request): The request object.
        queue (DownloadQueue): The download queue instance.
        encoder (Encoder): The encoder instance.
        notify (EventBus): The event bus instance.

    Returns:
        Response: The response object.

    """
    if not queue.is_paused():
        return web.json_response(
            {"message": "Non-active downloads are not paused."},
            status=web.HTTPNotAcceptable.status_code,
            dumps=encoder.encode,
        )

    queue.resume()

    msg = "Resumed all downloads."
    notify.emit(
        Events.RESUMED,
        data={"paused": False, "at": time.time()},
        title="Downloads Resumed",
        message=msg,
    )

    return web.json_response(data={"message": msg}, status=web.HTTPOk.status_code, dumps=encoder.encode)


@route("POST", "api/system/shutdown", "system.shutdown")
async def shutdown_system(request: Request, config: Config, encoder: Encoder, notify: EventBus) -> Response:
    """
    Initiate application shutdown.

    Args:
        request (Request): The request object.
        config (Config): The config instance.
        encoder (Encoder): The encoder instance.
        notify (EventBus): The event bus instance.

    Returns:
        Response: The response object.

    """
    if config.is_native is not True:
        return web.json_response(
            {"error": "Shutdown is only available in the native mode."},
            status=web.HTTPBadRequest.status_code,
        )

    app = request.app

    async def do_shutdown():
        notify.emit(
            Events.SHUTDOWN,
            data={"app": app},
            title="Application Shutdown",
            message="Shutdown initiated by user request.",
        )
        await asyncio.sleep(0.5)
        await app.shutdown()
        await app.cleanup()
        raise GracefulExit

    # Schedule shutdown after response
    asyncio.create_task(do_shutdown())
    LOG.info("Shutdown initiated by user request. Stopping the server...")
    return web.json_response(
        data={
            "message": "The application shutting down.",
        },
        status=web.HTTPOk.status_code,
        dumps=encoder.encode,
    )


@route("POST", "api/system/check-updates", "system.check_updates")
async def check_updates(config: Config, encoder: Encoder, update_checker: UpdateChecker) -> Response:
    """
    Manually trigger update check.

    Args:
        config (Config): The config instance.
        encoder (Encoder): The encoder instance.
        update_checker (UpdateChecker): The update checker instance.

    Returns:
        Response: The response object.

    """
    if not config.check_for_updates:
        return web.json_response(
            {"error": "Update checking is disabled in configuration."},
            status=web.HTTPBadRequest.status_code,
            dumps=encoder.encode,
        )

    # If update already found, return cached result
    if config.new_version:
        return web.json_response(
            data={
                "status": "update_available",
                "current_version": config.app_version,
                "new_version": config.new_version,
            },
            status=web.HTTPOk.status_code,
            dumps=encoder.encode,
        )

    # Run check and await result
    status, new_version = await update_checker.check_for_updates()

    return web.json_response(
        data={
            "status": status,
            "current_version": config.app_version,
            "new_version": new_version,
        },
        status=web.HTTPOk.status_code,
        dumps=encoder.encode,
    )


LOG: logging.Logger = logging.getLogger(__name__)


@route("POST", "api/system/terminal", "system.terminal")
async def stream_terminal(request: Request, config: Config, encoder: Encoder) -> Response | web.StreamResponse:
    if not config.console_enabled:
        return web.json_response(
            {"error": "Console feature is disabled."},
            status=web.HTTPForbidden.status_code,
        )

    if not request.can_read_body:
        return web.json_response(
            {"error": "Request body is required."},
            status=web.HTTPBadRequest.status_code,
        )

    payload = await request.json()
    if not isinstance(payload, dict):
        return web.json_response(
            {"error": "Invalid request payload."},
            status=web.HTTPBadRequest.status_code,
        )

    raw_command = payload.get("command")
    if not isinstance(raw_command, str) or not raw_command.strip():
        return web.json_response(
            {"error": "Command is required."},
            status=web.HTTPBadRequest.status_code,
        )

    response = web.StreamResponse(
        status=web.HTTPOk.status_code,
        headers={
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        },
    )
    await response.prepare(request)

    async def emit_event(event: str, data: dict) -> None:
        if request.transport is None or request.transport.is_closing():
            return
        payload: str = f"event: {event}\ndata: {encoder.encode(data)}\n\n"
        await response.write(payload.encode("utf-8"))

    returncode: int = -1
    try:
        LOG.info("Cli command from client. '%s'", raw_command)

        args: list[str] = ["yt-dlp", *shlex.split(raw_command, posix=os.name != "nt")]
        env_vars: dict[str, str] = os.environ.copy()
        env_vars.update(
            {
                "PWD": config.download_path,
                "FORCE_COLOR": "1",
                "PYTHONUNBUFFERED": "1",
            }
        )

        if "nt" != os.name:
            env_vars.update(
                {
                    "TERM": "xterm-256color",
                    "LANG": "en_US.UTF-8",
                    "LC_ALL": "en_US.UTF-8",
                    "SHELL": "/bin/bash",
                }
            )

        try:
            import pty

            master_fd, slave_fd = pty.openpty()
            stdin_arg = asyncio.subprocess.DEVNULL
            stdout_arg = stderr_arg = slave_fd
            use_pty = True
        except ImportError:
            use_pty = False
            master_fd = slave_fd = None
            stdin_arg = asyncio.subprocess.DEVNULL
            stdout_arg = asyncio.subprocess.PIPE
            stderr_arg = asyncio.subprocess.STDOUT

        creationflags = 0
        if os.name == "nt":
            import subprocess

            creationflags = subprocess.CREATE_NO_WINDOW

        proc: Process = await asyncio.create_subprocess_exec(
            *args,
            cwd=config.download_path,
            stdin=stdin_arg,
            stdout=stdout_arg,
            stderr=stderr_arg,
            env=env_vars,
            creationflags=creationflags,
        )

        if use_pty:
            assert slave_fd is not None
            try:
                os.close(slave_fd)
            except Exception as e:
                LOG.error("Error closing PTY. '%s'.", str(e))

        async def reader() -> None:
            if use_pty is False:
                assert proc.stdout is not None
                async for raw_line in proc.stdout:
                    line = raw_line.rstrip(b"\n")
                    await emit_event("output", {"type": "stdout", "line": line.decode("utf-8", errors="replace")})
                return

            assert master_fd is not None
            read_fd = master_fd
            loop: AbstractEventLoop = asyncio.get_running_loop()
            buffer: bytes = b""
            while True:
                try:
                    chunk: bytes = await loop.run_in_executor(None, lambda: os.read(read_fd, 1024))
                except OSError as e:
                    if e.errno == errno.EIO:
                        break
                    raise

                if not chunk:
                    if buffer:
                        await emit_event(
                            "output",
                            {"type": "stdout", "line": buffer.decode("utf-8", errors="replace")},
                        )
                    break

                buffer += chunk
                *lines, buffer = buffer.split(b"\n")

                for line in lines:
                    await emit_event(
                        "output",
                        {"type": "stdout", "line": line.decode("utf-8", errors="replace")},
                    )
            if master_fd is None:
                return
            try:
                os.close(master_fd)
            except Exception as e:
                LOG.error("Error closing PTY. '%s'.", str(e))

        read_task: Task = asyncio.create_task(reader(), name="cli_reader")

        returncode = await proc.wait()
        await read_task
    except Exception as e:
        LOG.error("CLI execute exception was thrown.")
        LOG.exception(e)
        await emit_event("output", {"type": "stderr", "line": str(e)})
    finally:
        await emit_event("close", {"exitcode": returncode})
        await response.write_eof()

    return response