From 87959028e5a0ffffe545b821660fa53e6f50dfe0 Mon Sep 17 00:00:00 2001
From: Serene-Arc <serenical@gmail.com>
Date: Tue, 25 May 2021 18:59:32 +1000
Subject: [PATCH] Add blacklist for web filetypes

---
 bdfr/site_downloaders/download_factory.py     | 20 ++++++++++++++++++-
 .../site_downloaders/test_download_factory.py | 13 ++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/bdfr/site_downloaders/download_factory.py b/bdfr/site_downloaders/download_factory.py
index 8eff2b8..cbfee2d 100644
--- a/bdfr/site_downloaders/download_factory.py
+++ b/bdfr/site_downloaders/download_factory.py
@@ -24,7 +24,8 @@ class DownloadFactory:
         sanitised_url = DownloadFactory.sanitise_url(url)
         if re.match(r'(i\.)?imgur.*\.gifv$', sanitised_url):
             return Imgur
-        elif re.match(r'.*/.*\.\w{3,4}(\?[\w;&=]*)?$', sanitised_url):
+        elif re.match(r'.*/.*\.\w{3,4}(\?[\w;&=]*)?$', sanitised_url) and \
+                not DownloadFactory.is_web_resource(sanitised_url):
             return Direct
         elif re.match(r'erome\.com.*', sanitised_url):
             return Erome
@@ -55,3 +56,20 @@ class DownloadFactory:
         split_url = split_url.netloc + split_url.path
         split_url = re.sub(beginning_regex, '', split_url)
         return split_url
+
+    @staticmethod
+    def is_web_resource(url: str) -> bool:
+        web_extensions = (
+            'asp',
+            'cfm',
+            'cfml',
+            'css',
+            'html',
+            'js',
+            'php',
+            'xhtml',
+        )
+        if re.match(rf'(?i).*/.*\.({"|".join(web_extensions)})$', url):
+            return True
+        else:
+            return False
diff --git a/tests/site_downloaders/test_download_factory.py b/tests/site_downloaders/test_download_factory.py
index d5e84d8..4b5356c 100644
--- a/tests/site_downloaders/test_download_factory.py
+++ b/tests/site_downloaders/test_download_factory.py
@@ -72,3 +72,16 @@ def test_factory_lever_bad(test_url: str):
 def test_sanitise_url(test_url: str, expected: str):
     result = DownloadFactory.sanitise_url(test_url)
     assert result == expected
+
+
+@pytest.mark.parametrize(('test_url', 'expected'), (
+    ('www.example.com/test.asp', True),
+    ('www.example.com/test.html', True),
+    ('www.example.com/test.js', True),
+    ('www.example.com/test.xhtml', True),
+    ('www.example.com/test.mp4', False),
+    ('www.example.com/test.png', False),
+))
+def test_is_web_resource(test_url: str, expected: bool):
+    result = DownloadFactory.is_web_resource(test_url)
+    assert result == expected