diff --git a/main.py b/main.py index e10ca3e..7ede2c5 100644 --- a/main.py +++ b/main.py @@ -1,20 +1,9 @@ -from __future__ import unicode_literals import simplematrixbotlib as botlib from re import search -import urllib.parse -import youtube_dl -import os +from urllib.parse import quote as sanitize_url +import yt_dlp -try: - passfile = open("pass.txt", "r") - passw = passfile.read().translate({ord(c): None for c in '\n'}) -except: - print("Put the password in pass.txt in this directory") - -creds = botlib.Creds("https://chat.trygve.me", "video_bot", passw) -bot = botlib.Bot(creds) - -prefix = 'https://' +from config import config class MyLogger(object): def debug(self, msg): @@ -27,40 +16,46 @@ class MyLogger(object): print(msg) -def my_hook(d): - if d['status'] == 'finished': +creds = botlib.Creds(config['url'], config['username'], config['password']) +bot = botlib.Bot(creds) + +def progress_hook(d): + if d['status'] == 'started': + print('Started downloading') + elif d['status'] == 'finished': print('Done downloading') - ydl_opts = { - 'format': 'bestvideo[ext=webm]+bestaudio[ext=webm]/best[ext=mp4]/best', - 'outtmpl': '/var/www/html/video_bot/%(title)s.%(ext)s', - 'restrictfilenames': False, + 'format': config['yt_dlp_format'], + 'outtmpl': config['download_dir'] + '/%(title)s.%(ext)s', + 'restrictfilenames': True, 'logger': MyLogger(), - 'progress_hooks': [my_hook], + 'progress_hooks': [progress_hook], } -def filterfunc(variable): - if search('www.youtube.com/watch', variable): - return True +def get_yt_vid_id(word): + if search('www.youtube.com/watch', word): + return word + elif search('youtu.be', word): + return word.replace('https://youtu.be/', 'https://www.youtube.com/watch?v=') else: return False +@bot.listener.on_message_event async def youtube(room, message): match = botlib.MessageMatch(room, message, bot) pieces = str(message).split() - links = filter(filterfunc, pieces) try: - for i in links: - with youtube_dl.YoutubeDL(ydl_opts) as ydl: - info_dict = ydl.extract_info(i, download=False) - fname = urllib.parse.quote(info_dict.get('title', None).replace('|', '_')) - url = 'https://trygve.me/video_bot/' + fname + '.webm' - ydl.download([i]) - await bot.api.send_text_message(room.room_id, url) + for i in pieces: + yt_url = get_yt_vid_id(i) + if yt_url: + with yt_dlp.YoutubeDL(ydl_opts) as ydl: + info_dict = ydl.extract_info(yt_url, download=False) + fname = sanitize_url(yt_dlp.utils.sanitize_filename(info_dict.get('title', None), restricted=True)) + url = 'https://trygve.me/video_bot/' + fname + '.webm' + ydl.download([yt_url]) + await bot.api.send_text_message(room.room_id, url) except Exception as e: - await bot.api.send_text_message(room.room_id, str(e)) - -bot.add_message_listener(youtube) + await bot.api.send_text_message(room.room_id, str(e)) bot.run()