From 1077fc1d5f99e86eb0ac1872773e4c9fa27ee9e1 Mon Sep 17 00:00:00 2001 From: nikita Date: Mon, 7 Sep 2026 03:23:03 +0300 Subject: [PATCH] Added basic downloading ability --- .gitignore | 1 + README.md | 8 +++++--- downloader.py | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++ logic.py | 39 +++++++++++++++++++++++++++++++++------ 4 files changed, 89 insertions(+), 9 deletions(-) create mode 100644 downloader.py diff --git a/.gitignore b/.gitignore index 150a104..42d3837 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ .venv/ session_storage/ __pycache__/ +temp/ *.swp *.swo *.vscode diff --git a/README.md b/README.md index 9b61260..2173707 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ Matrix бот для скачивания видео/аудио с *YouTube, Pin тестировать библиотеку [`mab`](https://git.tyukalov.su/nikita/mab), а точнее её возможности загрузки вложений в Matrix. -> В системе должны быть установлены `ffmpeg` и `ffprobe`. +> В системе должны быть установлены `ffmpeg`, `ffprobe` и `yt-dlp`. ## Работа с ботом @@ -23,6 +23,8 @@ Matrix бот для скачивания видео/аудио с *YouTube, Pin бот отправит медиафайлы из папки `test_media` - Если сообщение начинается не с `http://` или `https://`, то бот даст справку - Если сообщение начинается с `http://` или `https://`, то бот передаст текст - сообщения в качестве ссылки для `yt-dlp` + сообщения в качестве ссылки `yt-dlp`, скачает видео в максимально возможном + качестве, и отправит его в Matrix -**Дальнейший интерфейс ещё не придуман** \ No newline at end of file +**Интерфейс планируется добработать, в том числе добавить возможность скачивать +аудио, копировать cookies из браузера, выбирать качество.** \ No newline at end of file diff --git a/downloader.py b/downloader.py new file mode 100644 index 0000000..0de7e54 --- /dev/null +++ b/downloader.py @@ -0,0 +1,50 @@ +import asyncio +import tempfile +import signal +import os +from pathlib import Path +from dataclasses import dataclass + +class Downloader: + """Class that performs content download.""" + + # + # PUBLIC + # + def __init__(self, url: str): + self._url = url + + async def download_video(self, directory: str | Path) -> Path: + directory = Path(directory).resolve() + directory.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory(dir=directory, prefix=".ytdlp-") as tmp: + proc = await asyncio.create_subprocess_exec( + "yt-dlp", + "-f", "bv*+ba/b", + "-P", f"home:{directory}", + "-P", f"temp:{tmp}", + "--no-playlist", + "--print", "after_move:filepath", + self._url, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=True, + ) + + try: + stdout, stderr = await proc.communicate() + + except asyncio.CancelledError: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + + await asyncio.shield(proc.wait()) + raise + + if proc.returncode != 0: + raise RuntimeError(stderr.decode(errors="replace")) + + return Path(stdout.decode().strip()) \ No newline at end of file diff --git a/logic.py b/logic.py index 866c079..6741a0b 100644 --- a/logic.py +++ b/logic.py @@ -3,6 +3,8 @@ import asyncio import glob import html +import os +import shutil import traceback from mab import MatrixBot @@ -11,10 +13,35 @@ from mab.filters.text import TextFilter, TextStartsWithFilter, TextEndsWithFilte from nio import MatrixRoom, RoomMessageText +from downloader import Downloader + # # PRIVATE DATA # _bot: MatrixBot +_tasks: set[asyncio.Task] = set() + +# +# PRIVATE +# +async def worker_download(room: MatrixRoom, text: RoomMessageText) -> None: + # respond + response = ( + f"Скачивание видео по " + f"вашей ссылке" + ) + await _bot.send_text_to_room(room.room_id, response, True) + # get formats + d = Downloader(text.body) + result = await d.download_video("temp") + await _bot.send_text_to_room(room.room_id, "Скачано, загрузка в Matrix...", True) + # upload + try: + await _bot.send_video_to_room(room.room_id, result, caption=f"Скачано из {text.body}") + except: + await _bot.send_text_to_room(room.room_id, traceback.format_exc(), False) + finally: + os.unlink(result) # # CALLBACKS @@ -33,9 +60,9 @@ async def command_http_https(room: MatrixRoom, text: RoomMessageText) -> None: """The message starts with `http://` or `https://`""" if text.sender == _bot.get_client().user_id: return - resp = "Сообщение начинается с http:// или с " - resp += "https://. Ещё не реализовано." - await _bot.send_text_to_room(room.room_id, resp, True) + t = asyncio.create_task(worker_download(room, text)) + _tasks.add(t) + t.add_done_callback(_tasks.discard) async def command_test(room: MatrixRoom, text: RoomMessageText) -> None: """The message ends with ` TesT` (case matters)""" @@ -83,14 +110,14 @@ async def command_test(room: MatrixRoom, text: RoomMessageText) -> None: await _bot.send_text_to_room(room.room_id, traceback.format_exc(), False) await asyncio.sleep(0.5) - - # # PUBLIC # async def initialize(bot: MatrixBot) -> None: """This function initializes the bot. It does not start it.""" global _bot + shutil.rmtree("temp", ignore_errors=True) + os.makedirs("temp", exist_ok=True) _bot = bot bot.add_event_callback(command_http_https, TextStartsWithFilter(["https://", "http://"])) bot.add_event_callback(command_test, TextEndsWithFilter(" TesT", any_case=False)) @@ -98,4 +125,4 @@ async def initialize(bot: MatrixBot) -> None: async def deinitizalize(bot: MatrixBot) -> None: """This function deinitializes the bot. It does not stop it.""" - pass \ No newline at end of file + shutil.rmtree("temp", ignore_errors=True) \ No newline at end of file