132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
"""Bot's callbacks implementation"""
|
||
|
||
import asyncio
|
||
import glob
|
||
import html
|
||
import os
|
||
import shutil
|
||
import traceback
|
||
|
||
from mab import MatrixBot
|
||
from mab.types import VideoFileProperties
|
||
from mab.filters.text import TextFilter, TextStartsWithFilter, TextEndsWithFilter
|
||
|
||
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"<strong>Скачивание видео по "
|
||
f"<a href=\"{html.escape(text.body)}\">вашей ссылке</a></strong>"
|
||
)
|
||
await _bot.send_text_to_room(room.room_id, response, True)
|
||
# get formats
|
||
d = Downloader(text.body)
|
||
try:
|
||
result = await d.download_video("temp")
|
||
except:
|
||
await _bot.send_text_to_room(room.room_id, traceback.format_exc(), False)
|
||
return
|
||
await _bot.send_text_to_room(room.room_id, "<strong>Скачано, загрузка в Matrix...</strong>", 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
|
||
#
|
||
async def command_help(room: MatrixRoom, text: RoomMessageText) -> None:
|
||
"""!help"""
|
||
if text.sender == _bot.get_client().user_id:
|
||
return
|
||
resp = "<strong>2026-matrix-down</strong><br>"
|
||
resp += "Этот бот используется для скачивания видео/аудио с хостингов, "
|
||
resp += "поддерживаемых <code>yt-dlp</code>. Просто отправьте ссылку в эту "
|
||
resp += "комнату, чтобы бот предложил варианты загрузки."
|
||
await _bot.send_text_to_room(room.room_id, resp, True)
|
||
|
||
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
|
||
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)"""
|
||
if text.sender == _bot.get_client().user_id:
|
||
return
|
||
# glob patterns to use
|
||
glob_img = ["*.png", "*.jpg", ".jpeg", "*.gif", "*.webp"]
|
||
glob_vid = ["*.mp4", "*.avi", "*.mkv"]
|
||
# TODO
|
||
glob_aud = ["*.mp3", "*.m4a", ".wav", "*.ogg", "*.flac"]
|
||
# help message
|
||
resp = "<strong>Сообщение заканчивается на <code> TesT</code>.</strong><br>"
|
||
resp += "Сейчас будут отправлены все файлы из директории <code>test_media"
|
||
resp += "</code>, у которых следующие расширения:<br>"
|
||
resp += f"<ul><li>Изображения: {", ".join([f"<code>{g}</code>" for g in glob_img])}</li></ul>"
|
||
resp += f"<ul><li>Видео: {", ".join([f"<code>{g}</code>" for g in glob_vid])}</li></ul>"
|
||
await _bot.send_text_to_room(room.room_id, resp, is_html=True)
|
||
await asyncio.sleep(1)
|
||
# send images
|
||
for g in [f"test_media/{g}" for g in glob_img]:
|
||
imgs = glob.glob(g)
|
||
for file in imgs:
|
||
try:
|
||
await _bot.send_image_to_room(
|
||
room.room_id,
|
||
file,
|
||
caption=f"Картинка <code>{html.escape(file)}</code>",
|
||
is_html=True
|
||
)
|
||
except:
|
||
await _bot.send_text_to_room(room.room_id, traceback.format_exc(), False)
|
||
await asyncio.sleep(0.5)
|
||
# send videos
|
||
for g in [f"test_media/{g}" for g in glob_vid]:
|
||
imgs = glob.glob(g)
|
||
for file in imgs:
|
||
try:
|
||
await _bot.send_video_to_room(
|
||
room.room_id,
|
||
file,
|
||
caption=f"Видео <code>{html.escape(file)}</code>",
|
||
is_html=True
|
||
)
|
||
except:
|
||
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))
|
||
bot.add_event_callback(command_help, TextFilter())
|
||
|
||
async def deinitizalize(bot: MatrixBot) -> None:
|
||
"""This function deinitializes the bot. It does not stop it."""
|
||
shutil.rmtree("temp", ignore_errors=True) |