import asyncio import errno import functools import logging import os import pty import shlex import time from datetime import UTC, datetime import anyio import socketio from aiohttp import web from .common import Common from .config import Config from .DownloadQueue import DownloadQueue from .encoder import Encoder from .Events import Event, EventBus, Events, error from .ItemDTO import Item from .Presets import Presets from .Utils import is_downloaded, tail_log LOG = logging.getLogger("socket_api") class HttpSocket(Common): """ This class is used to handle WebSocket events. """ config: Config sio: socketio.AsyncServer queue: DownloadQueue subscribers: dict[str, list[str]] = {} """Event subscriber list.""" log_task = None """Task to tail the log file.""" def __init__( self, queue: DownloadQueue | None = None, encoder: Encoder | None = None, config: Config | None = None, sio: socketio.AsyncServer | None = None, ): self.config = config or Config.get_instance() self.queue = queue or DownloadQueue.get_instance() self._notify = EventBus.get_instance() self.sio = sio or socketio.AsyncServer(cors_allowed_origins="*") encoder = encoder or Encoder() def emit(e: Event, _, **kwargs): return self.sio.emit(event=e.event, data=encoder.encode(e.data), **kwargs) self._notify.subscribe("frontend", emit, f"{__class__.__name__}.socket_api") super().__init__(queue=queue, encoder=encoder, config=config) @staticmethod def ws_event(func): # type: ignore """ Decorator to mark a method as a socket event. """ @functools.wraps(func) # type: ignore async def wrapper(*args, **kwargs): return await func(*args, **kwargs) # type: ignore wrapper._ws_event = func.__name__ # type: ignore return wrapper async def on_shutdown(self, _: web.Application): LOG.debug("Shutting down socket server.") def attach(self, app: web.Application): self.sio.attach(app, socketio_path=self.config.url_socketio) for attr_name in dir(self): method = getattr(self, attr_name) if hasattr(method, "_ws_event") and self.sio: self.sio.on(method._ws_event)(method) # type: ignore self._notify.subscribe( Events.ADD_URL, lambda data, _, **kwargs: self.add(item=Item.format(data.data)), # noqa: ARG005 f"{__class__.__name__}.socket_add_url", ) # register the shutdown event. app.on_shutdown.append(self.on_shutdown) @ws_event async def cli_post(self, sid: str, data): if not self.config.console_enabled: await self._notify.emit(Events.ERROR, data=error("Console is disabled."), to=sid) return if not data: await self._notify.emit(Events.CLI_CLOSE, data={"exitcode": 0}, to=sid) return try: LOG.info(f"Cli command from client '{sid}'. '{data}'") args = ["yt-dlp", *shlex.split(data)] _env = os.environ.copy() _env.update( { "TERM": "xterm-256color", "LANG": "en_US.UTF-8", "SHELL": "/bin/bash", "LC_ALL": "en_US.UTF-8", "PWD": self.config.download_path, "FORCE_COLOR": "1", "PYTHONUNBUFFERED": "1", } ) master_fd, slave_fd = pty.openpty() proc = await asyncio.create_subprocess_exec( *args, cwd=self.config.download_path, stdin=asyncio.subprocess.DEVNULL, stdout=slave_fd, stderr=slave_fd, env=_env, ) try: os.close(slave_fd) except Exception as e: LOG.error(f"Error closing PTY. '{e!s}'.") async def read_pty(sid: str): loop = asyncio.get_running_loop() buffer = b"" while True: try: chunk = await loop.run_in_executor(None, lambda: os.read(master_fd, 1024)) except OSError as e: if e.errno == errno.EIO: break raise if not chunk: # No more output if buffer: # Emit any remaining partial line await self._notify.emit( Events.CLI_OUTPUT, data={"type": "stdout", "line": buffer.decode("utf-8", errors="replace")}, to=sid, ) break buffer += chunk *lines, buffer = buffer.split(b"\n") for line in lines: await self._notify.emit( Events.CLI_OUTPUT, data={"type": "stdout", "line": line.decode("utf-8", errors="replace")}, to=sid, ) try: os.close(master_fd) except Exception as e: LOG.error(f"Error closing PTY. '{e!s}'.") # Start reading output from PTY read_task = asyncio.create_task(read_pty(sid=sid)) # Wait until process finishes returncode = await proc.wait() # Ensure reading is done await read_task await self._notify.emit(Events.CLI_CLOSE, data={"exitcode": returncode}, to=sid) except Exception as e: LOG.error(f"CLI execute exception was thrown for client '{sid}'.") LOG.exception(e) await self._notify.emit(Events.CLI_OUTPUT, data={"type": "stderr", "line": str(e)}, to=sid) await self._notify.emit(Events.CLI_CLOSE, data={"exitcode": -1}, to=sid) @ws_event async def add_url(self, sid: str, data: dict): url: str | None = data.get("url") if not url: await self._notify.emit(Events.ERROR, data=error("No URL provided.", data={"unlock": True}), to=sid) return try: await self._notify.emit( event=Events.STATUS, data=await self.add(item=Item.format(data)), to=sid, ) except ValueError as e: LOG.exception(e) await self._notify.emit(Events.ERROR, data=error(str(e)), to=sid) return @ws_event async def item_cancel(self, sid: str, id: str): if not id: await self._notify.emit(Events.ERROR, data=error("Invalid request."), to=sid) return status: dict[str, str] = {} status = await self.queue.cancel([id]) status.update({"identifier": id}) await self._notify.emit(Events.ITEM_CANCEL, data=status) @ws_event async def item_delete(self, sid: str, data: dict): if not data: await self._notify.emit(Events.ERROR, data=error("Invalid request."), to=sid) return id: str | None = data.get("id") if not id: await self._notify.emit(Events.ERROR, data=error("Invalid request."), to=sid) return status: dict[str, str] = {} status = await self.queue.clear([id], remove_file=bool(data.get("remove_file", False))) status.update({"identifier": id}) await self._notify.emit(Events.ITEM_DELETE, data=status) @ws_event async def archive_item(self, _: str, data: dict): if not isinstance(data, dict) or "url" not in data: return from .YTDLPOpts import YTDLPOpts params = YTDLPOpts.get_instance() if "preset" in data and isinstance(data["preset"], str): params.preset(name=data["preset"]) if "cli" in data and isinstance(data["cli"], str) and len(data["cli"]) > 1: params.add_cli(data["cli"], from_user=True) params = params.get_all() file: str = params.get("download_archive", None) if not file: return exists, idDict = is_downloaded(file, data["url"]) if exists or "archive_id" not in idDict or idDict["archive_id"] is None: return async with await anyio.open_file(file, "a") as f: await f.write(f"{idDict['archive_id']}\n") manual_archive = self.config.manual_archive if manual_archive: previouslyArchived = False if os.path.exists(manual_archive): async with await anyio.open_file(manual_archive) as f: async for line in f: if idDict["archive_id"] in line: previouslyArchived = True break if not previouslyArchived: async with await anyio.open_file(manual_archive, "a") as f: await f.write(f"{idDict['archive_id']} - at: {datetime.now(UTC).isoformat()}\n") LOG.info(f"Archiving url '{data['url']}' with id '{idDict['archive_id']}'.") else: LOG.info(f"URL '{data['url']}' with id '{idDict['archive_id']}' already archived.") @ws_event async def connect(self, sid: str, _=None): data = { **self.queue.get(), "config": self.config.frontend(), "presets": Presets.get_instance().get_all(), "paused": self.queue.is_paused(), } # get download folder listing. downloadPath: str = self.config.download_path data["folders"] = [name for name in os.listdir(downloadPath) if os.path.isdir(os.path.join(downloadPath, name))] await self._notify.emit(Events.INITIAL_DATA, data=data, to=sid) @ws_event async def pause(self, *_, **__): self.queue.pause() await self._notify.emit(Events.PAUSED, data={"paused": True, "at": time.time()}) @ws_event async def resume(self, *_, **__): self.queue.resume() await self._notify.emit(Events.PAUSED, data={"paused": False, "at": time.time()}) @ws_event async def subscribe(self, sid: str, event: str): """ Subscribe to a specific event. Args: sid (str): The session ID of the client. event (str): The event to subscribe to. """ if not isinstance(event, str): await self._notify.emit(Events.ERROR, data=error("Invalid event."), to=sid) return if event not in self.subscribers: self.subscribers[event] = [] if sid not in self.subscribers[event]: self.subscribers[event].append(sid) LOG.debug(f"Client '{sid}' subscribed to event '{event}'.") await self.sio.emit(Events.SUBSCRIBED, data={"event": event}, to=sid) async def emit_logs(data: dict): await self.subscribe_emit(event=event, data=data) if "log_lines" == event and self.log_task is None: LOG.debug("Starting log tailing task.") self.log_task = asyncio.create_task( tail_log( file=os.path.join(self.config.config_path, "logs", "app.log"), emitter=emit_logs, ), name="tail_log", ) @ws_event async def unsubscribe(self, sid: str, event: str): """ Unsubscribe from a specific event. Args: sid (str): The session ID of the client. event (str): The event to unsubscribe from. """ if event not in self.subscribers: return if sid not in self.subscribers[event]: return self.subscribers[event].remove(sid) await self.sio.emit(Events.UNSUBSCRIBED, data={"event": event}, to=sid) LOG.debug(f"Client '{sid}' unsubscribed from event '{event}'.") if "log_lines" != event or not self.log_task or "log_lines" not in self.subscribers: return if len(self.subscribers["log_lines"]) < 1: try: LOG.debug("Stopping log tailing task.") self.log_task.cancel() self.log_task = None except asyncio.CancelledError: pass @ws_event async def disconnect(self, sid: str): """ Handle client disconnection. Args: sid (str): The session ID of the client. """ LOG.debug(f"Client '{sid}' disconnected.") for event in self.subscribers: if sid in self.subscribers[event]: await self.unsubscribe(sid=sid, event=event) async def subscribe_emit(self, event: str, data: dict): if event not in self.subscribers or len(self.subscribers[event]) < 1: return for sid in self.subscribers[event]: await self.sio.emit(event=event, data=data, to=sid)