299 lines
10 KiB
Python
299 lines
10 KiB
Python
import asyncio
|
|
import errno
|
|
import functools
|
|
import json
|
|
import logging
|
|
import os
|
|
import pty
|
|
import shlex
|
|
import time
|
|
from datetime import datetime
|
|
|
|
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 .Utils import isDownloaded, arg_converter
|
|
|
|
LOG = logging.getLogger("socket_api")
|
|
|
|
|
|
class HttpSocket(common):
|
|
"""
|
|
This class is used to handle WebSocket events.
|
|
"""
|
|
|
|
config: Config
|
|
sio: socketio.AsyncServer
|
|
|
|
def __init__(self, queue: DownloadQueue, emitter: Emitter, encoder: Encoder):
|
|
super().__init__(queue=queue, encoder=encoder)
|
|
|
|
self.sio = socketio.AsyncServer(cors_allowed_origins="*")
|
|
emitter.add_emitter(lambda event, data, **kwargs: self.sio.emit(event, encoder.encode(data), **kwargs))
|
|
|
|
self.config = Config.get_instance()
|
|
|
|
self.queue = queue
|
|
self.emitter = emitter
|
|
|
|
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
|
|
|
|
def attach(self, app: web.Application):
|
|
self.sio.attach(app, socketio_path=self.config.url_prefix + "socket.io")
|
|
|
|
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
|
|
|
|
@ws_event # type: ignore
|
|
async def cli_post(self, sid: str, data):
|
|
if not data:
|
|
await self.emitter.emit("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. '{str(e)}'.")
|
|
|
|
async def read_pty():
|
|
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
|
|
else:
|
|
raise
|
|
|
|
if not chunk:
|
|
# No more output
|
|
if buffer:
|
|
# Emit any remaining partial line
|
|
await self.emitter.emit(
|
|
"cli_output", {"type": "stdout", "line": buffer.decode("utf-8", errors="replace")}
|
|
)
|
|
break
|
|
buffer += chunk
|
|
*lines, buffer = buffer.split(b"\n")
|
|
for line in lines:
|
|
await self.emitter.emit(
|
|
"cli_output", {"type": "stdout", "line": line.decode("utf-8", errors="replace")}
|
|
)
|
|
try:
|
|
os.close(master_fd)
|
|
except Exception as e:
|
|
LOG.error(f"Error closing PTY. '{str(e)}'.")
|
|
|
|
# Start reading output from PTY
|
|
read_task = asyncio.create_task(read_pty())
|
|
|
|
# Wait until process finishes
|
|
returncode = await proc.wait()
|
|
|
|
# Ensure reading is done
|
|
await read_task
|
|
|
|
await self.emitter.emit("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(
|
|
"cli_out1put",
|
|
{
|
|
"type": "stderr",
|
|
"line": str(e),
|
|
},
|
|
to=sid,
|
|
)
|
|
await self.emitter.emit("cli_close", {"exitcode": -1}, to=sid)
|
|
|
|
@ws_event # type: ignore
|
|
async def add_url(self, sid: str, data: dict):
|
|
url: str | None = data.get("url")
|
|
|
|
if not url:
|
|
await self.emitter.warning("No URL provided.", to=sid)
|
|
return
|
|
|
|
preset: str = str(data.get("preset", self.config.default_preset))
|
|
folder: str = str(data.get("folder")) if data.get("folder") else ""
|
|
ytdlp_cookies: str = str(data.get("ytdlp_cookies")) if data.get("ytdlp_cookies") else ""
|
|
output_template: str = str(data.get("output_template")) if data.get("output_template") else ""
|
|
|
|
ytdlp_config = data.get("ytdlp_config")
|
|
if isinstance(ytdlp_config, str) and ytdlp_config:
|
|
try:
|
|
ytdlp_config = json.loads(ytdlp_config)
|
|
except Exception as e:
|
|
await self.emitter.warning(f"Failed to parse json yt-dlp config. {str(e)}", to=sid)
|
|
return
|
|
|
|
status = await self.add(
|
|
url=url,
|
|
preset=preset,
|
|
folder=folder,
|
|
ytdlp_cookies=ytdlp_cookies,
|
|
ytdlp_config=ytdlp_config if isinstance(ytdlp_config, dict) else {},
|
|
output_template=output_template,
|
|
)
|
|
|
|
await self.emitter.emit("status", status, to=sid)
|
|
|
|
@ws_event # type: ignore
|
|
async def item_cancel(self, sid: str, id: str):
|
|
if not id:
|
|
await self.emitter.warning("Invalid request.", to=sid)
|
|
return
|
|
|
|
status: dict[str, str] = {}
|
|
status = await self.queue.cancel([id])
|
|
status.update({"identifier": id})
|
|
|
|
await self.emitter.emit("item_cancel", status)
|
|
|
|
@ws_event # type: ignore
|
|
async def item_delete(self, sid: str, data: dict):
|
|
if not data:
|
|
await self.emitter.warning("Invalid request.", to=sid)
|
|
return
|
|
|
|
id: str | None = data.get("id")
|
|
if not id:
|
|
await self.emitter.warning("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("item_delete", status)
|
|
|
|
@ws_event # type: ignore
|
|
async def archive_item(self, sid: 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 = isDownloaded(file, data["url"])
|
|
if exists or "archive_id" not in idDict or idDict["archive_id"] is None:
|
|
return
|
|
|
|
with open(file, "a") as f:
|
|
f.write(f"{idDict['archive_id']}\n")
|
|
|
|
manual_archive = self.config.manual_archive
|
|
if manual_archive:
|
|
previouslyArchived = False
|
|
if os.path.exists(manual_archive):
|
|
with open(manual_archive, "r") as f:
|
|
for line in f.readlines():
|
|
if idDict["archive_id"] in line:
|
|
previouslyArchived = True
|
|
break
|
|
|
|
if not previouslyArchived:
|
|
with open(manual_archive, "a") as f:
|
|
f.write(f"{idDict['archive_id']} - at: {datetime.now().isoformat()}\n")
|
|
|
|
LOG.info(f"Archiving url '{data['url']}' with id '{idDict['archive_id']}'.")
|
|
|
|
@ws_event # type: ignore
|
|
async def connect(self, sid: str, _=None):
|
|
data = {
|
|
**self.queue.get(),
|
|
"config": self.config.frontend(),
|
|
"tasks": self.config.tasks,
|
|
"presets": self.config.presets,
|
|
"paused": self.queue.isPaused(),
|
|
}
|
|
|
|
# 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("initial_data", data, to=sid)
|
|
|
|
@ws_event # type: ignore
|
|
async def pause(self, sid: str, _=None):
|
|
self.queue.pause()
|
|
await self.emitter.emit("paused", {"paused": True, "at": time.time()})
|
|
|
|
@ws_event # type: ignore
|
|
async def resume(self, sid: str, _=None):
|
|
self.queue.resume()
|
|
await self.emitter.emit("paused", {"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.warning("Invalid request or no options were given.", to=sid)
|
|
return
|
|
|
|
args: str | None = data.get("args")
|
|
|
|
if not args:
|
|
await self.emitter.warning("no options were given.", to=sid)
|
|
return
|
|
|
|
try:
|
|
await self.emitter.emit("ytdlp_convert", arg_converter(args), to=sid)
|
|
return
|
|
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. '{e}'.", to=sid)
|