From deb4b43fdd6f69b9235caefffa1704accb5ecb28 Mon Sep 17 00:00:00 2001
From: jaskaranSM <beerabassian32@gmail.com>
Date: Sun, 27 Oct 2019 17:29:02 +0530
Subject: [PATCH] Validate Url and Magnet with Regex

---
 bot/helper/ext_utils/bot_utils.py         | 13 +++++++++----
 bot/helper/mirror_utils/download_tools.py |  9 +++++----
 bot/modules/mirror.py                     |  6 ++++--
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py
index 2d7553b..273804f 100644
--- a/bot/helper/ext_utils/bot_utils.py
+++ b/bot/helper/ext_utils/bot_utils.py
@@ -1,8 +1,12 @@
 from bot import download_dict, download_dict_lock
 import logging
+import re
 
 LOGGER = logging.getLogger(__name__)
 
+MAGNET_REGEX = r"magnet:\?xt=urn:btih:[a-zA-Z0-9]*"
+
+URL_REGEX = r"(?:(?:https?|ftp):\/\/)?[\w/\-?=%.]+\.[\w/\-?=%.]+"
 
 class MirrorStatus:
     STATUS_UPLOADING = "Uploading"
@@ -112,13 +116,14 @@ def get_readable_time(seconds: int) -> str:
 
 
 def is_url(url: str):
-    # TODO: Find the proper way to validate the url
-    if url.startswith('https://') or url.startswith('http://'):
+    url = re.findall(URL_REGEX,url)
+    if url:
         return True
-    return False
+    return False    
 
 
 def is_magnet(url: str):
-    if "magnet" in url:
+    magnet = re.findall(MAGNET_REGEX,url)
+    if magnet:
         return True
     return False
diff --git a/bot/helper/mirror_utils/download_tools.py b/bot/helper/mirror_utils/download_tools.py
index 5800c4b..80d9f19 100644
--- a/bot/helper/mirror_utils/download_tools.py
+++ b/bot/helper/mirror_utils/download_tools.py
@@ -12,13 +12,13 @@ class DownloadHelper:
         self.__is_torrent = False
 
     def add_download(self, link: str):
-        if is_url(link):
+        if is_magnet(link):
+            download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.uid)})
+            self.__is_torrent = True
+        else:
             if link.endswith('.torrent'):
                 self.__is_torrent = True
             download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.uid)})
-        else:
-            download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.uid)})
-            self.__is_torrent = True
         with download_dict_lock:
             download_dict[self.__listener.message.message_id] = DownloadStatus(download.gid,
                                                                                self.__listener.uid)
@@ -60,6 +60,7 @@ class DownloadHelper:
                     sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)        
             new_gid = self.__get_followed_download_gid()
             with download_dict_lock:
+                LOGGER.info(f"{download.name}: Changing GID {download.gid} to {new_gid}")
                 download_dict[self.__listener.message.message_id] = DownloadStatus(new_gid,
                                                                                self.__listener.message.message_id)
 
diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py
index ca3e066..4e4eeec 100644
--- a/bot/modules/mirror.py
+++ b/bot/modules/mirror.py
@@ -92,10 +92,12 @@ class MirrorListener(listeners.MirrorListeners):
             deleteMessage(self.context, self.reply_message)
             with status_reply_dict_lock:
                 del status_reply_dict[self.update.effective_chat.id]
-        except BadRequest:
+        except BadRequest as e:
+        	LOGGER.error(str(e))
             # This means that the message has been deleted because of a /status command
             pass
-        except KeyError:
+        except KeyError as e:
+        	LOGGER.error(str(e))
             pass
         sendMessage(msg, self.context, self.update)
         try: