310 lines
10 KiB
Python
310 lines
10 KiB
Python
import asyncio
|
|
import errno
|
|
import functools
|
|
import logging
|
|
import os
|
|
import pty
|
|
import shlex
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
import anyio
|
|
import socketio
|
|
from aiohttp import web
|
|
|
|
from .common import Common
|
|
from .config import Config
|
|
from .DownloadQueue import DownloadQueue
|
|
from .Emitter import Emitter
|
|
from .encoder import Encoder
|
|
from .Events import EventBus, Events
|
|
from .Presets import Presets
|
|
from .Utils import arg_converter, is_downloaded
|
|
|
|
LOG = logging.getLogger("socket_api")
|
|
|
|
|
|
class HttpSocket(Common):
|
|
"""
|
|
This class is used to handle WebSocket events.
|
|
"""
|
|
|
|
config: Config
|
|
sio: socketio.AsyncServer
|
|
queue: DownloadQueue
|
|
emitter: Emitter
|
|
|
|
def __init__(
|
|
self,
|
|
queue: DownloadQueue | None = None,
|
|
emitter: Emitter | 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.emitter = emitter or Emitter.get_instance()
|
|
self.sio = sio or socketio.AsyncServer(cors_allowed_origins="*")
|
|
encoder = encoder or Encoder()
|
|
|
|
def emit(event: str, data: Any, **kwargs):
|
|
return self.sio.emit(event=event, data=encoder.encode(data), **kwargs)
|
|
|
|
self.emitter.add_emitter([emit], local=False)
|
|
|
|
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
|
|
|
|
EventBus.get_instance().subscribe(
|
|
Events.ADD_URL, lambda data, _: self.add(**data.data), "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.emitter.error("Console is disabled.", to=sid)
|
|
return
|
|
|
|
if not data:
|
|
await self.emitter.emit(Events.CLI_CLOSE, {"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.emitter.emit(
|
|
Events.CLI_OUTPUT,
|
|
{"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.emitter.emit(
|
|
Events.CLI_OUTPUT,
|
|
{"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.emitter.emit(Events.CLI_CLOSE, {"exitcode": returncode}, to=sid)
|
|
except Exception as e:
|
|
LOG.error(f"CLI execute exception was thrown for client '{sid}'.")
|
|
LOG.exception(e)
|
|
await self.emitter.emit(Events.CLI_OUTPUT, {"type": "stderr", "line": str(e)}, to=sid)
|
|
await self.emitter.emit(Events.CLI_CLOSE, {"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.emitter.error("No URL provided.", data={"unlock": True}, to=sid)
|
|
return
|
|
|
|
try:
|
|
item = self.format_item(data)
|
|
except ValueError as e:
|
|
await self.emitter.error(str(e), to=sid)
|
|
return
|
|
|
|
status = await self.add(**item)
|
|
await self.emitter.emit(event=Events.STATUS, data=status, to=sid)
|
|
|
|
@ws_event
|
|
async def item_cancel(self, sid: str, id: str):
|
|
if not id:
|
|
await self.emitter.error("Invalid request.", to=sid)
|
|
return
|
|
|
|
status: dict[str, str] = {}
|
|
status = await self.queue.cancel([id])
|
|
status.update({"identifier": id})
|
|
|
|
await self.emitter.emit(event=Events.ITEM_CANCEL, data=status)
|
|
|
|
@ws_event
|
|
async def item_delete(self, sid: str, data: dict):
|
|
if not data:
|
|
await self.emitter.error("Invalid request.", to=sid)
|
|
return
|
|
|
|
id: str | None = data.get("id")
|
|
if not id:
|
|
await self.emitter.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.emitter.emit(event=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 or not self.config.keep_archive:
|
|
return
|
|
|
|
if not isinstance(self.config.ytdl_options, dict):
|
|
self.config.ytdl_options = {}
|
|
|
|
file: str = self.config.ytdl_options.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']}'.")
|
|
|
|
@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.emitter.emit(event=Events.INITIAL_DATA, data=data, to=sid)
|
|
|
|
@ws_event
|
|
async def pause(self, *_, **__):
|
|
self.queue.pause()
|
|
await self.emitter.emit(event=Events.PAUSED, data={"paused": True, "at": time.time()})
|
|
|
|
@ws_event
|
|
async def resume(self, *_, **__):
|
|
self.queue.resume()
|
|
await self.emitter.emit(event=Events.PAUSED, data={"paused": False, "at": time.time()})
|
|
|
|
@ws_event
|
|
async def ytdlp_convert(self, sid: str, data: dict):
|
|
if not isinstance(data, dict) or "args" not in data:
|
|
await self.emitter.error("Invalid request or no options were given.", to=sid)
|
|
return
|
|
|
|
args: str | None = data.get("args")
|
|
|
|
if not args:
|
|
await self.emitter.error("no options were given.", to=sid)
|
|
return
|
|
|
|
try:
|
|
await self.emitter.emit(event=Events.YTDLP_CONVERT, data=arg_converter(args), to=sid)
|
|
except Exception as e:
|
|
err = str(e).strip()
|
|
err = err.split("\n")[-1] if "\n" in err else err
|
|
LOG.error(f"Failed to convert args. '{err}'.")
|
|
await self.emitter.error(f"Failed to convert options. '{err}'.", to=sid)
|
|
|
|
return
|