Refactor

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()