fixed filedownload tests

TheTechromancer · TheTechromancer · commit e5963f367e9f · 2023-09-21T15:19:25.000-04:00
diff --git a/bbot/modules/filedownload.py b/bbot/modules/filedownload.py
@@ -23,27 +23,24 @@ class filedownload(BaseModule):
             "bashrc",  #  Bash Script or Configuration
             "conf",  #  Configuration File
             "cfg",  #  Configuration File
-            "cr2",  #  Canon RAW Image
             "crt",  #  Certificate File
-            "crw",  #  Canon RAW Image (Older Format)
             "csv",  #  Comma Separated Values File
             "db",  #  SQLite Database File
             "sqlite",  #  SQLite Database File
             "doc",  #  Microsoft Word Document (Old Format)
             "docx",  #  Microsoft Word Document
+            "exe",  #  Windows PE executable
             "ica",  #  Citrix Independent Computing Architecture File
             "indd",  #  Adobe InDesign Document
             "ini",  #  Initialization File
             "jar",  #  Java Archive
-            "jpg",  #  JPEG Image
-            "jpeg",  #  JPEG Image
-            "js",  #  JavaScript File
             "json",  #  JavaScript Object Notation File
             "key",  #  Private Key File
             "pub",  #  Public Key File
             "log",  #  Log File
-            "md",  #  Markdown File
             "markdown",  #  Markdown File
+            "md",  #  Markdown File
+            "msi",  # Windows setup file
             "odg",  #  OpenDocument Graphics (LibreOffice, OpenOffice)
             "odp",  #  OpenDocument Presentation (LibreOffice, OpenOffice)
             "ods",  #  OpenDocument Spreadsheet (LibreOffice, OpenOffice)
@@ -60,8 +57,6 @@ class filedownload(BaseModule):
             "rdp",  #  Remote Desktop Protocol File
             "sh",  #  Shell Script
             "sql",  #  SQL Database Dump
-            "svg",  #  Scalable Vector Graphics
-            "svgz",  #  Compressed SVG
             "swp",  #  Swap File (temporary file, often Vim)
             "sxw",  #  OpenOffice.org Writer document
             "tar",  #  Tar Archive
@@ -154,9 +149,10 @@ def make_filename(self, url, content_type=None):
         if not extension:
             if content_type and content_type in self.mime_db:
                 extension = self.mime_db[content_type]
-            else:
-                self.debug(f'Extension "{extension}" at url "{url}" not in list of watched extensions.')
-                return None, None, None
+
+        if (not extension) or (extension not in self.extensions):
+            self.debug(f'Extension "{extension}" at url "{url}" not in list of watched extensions.')
+            return None, None, None
 
         orig_filename = Path(url_path).stem
         if extension:
diff --git a/bbot/test/test_step_2/module_tests/test_module_filedownload.py b/bbot/test/test_step_2/module_tests/test_module_filedownload.py
@@ -2,8 +2,9 @@
 
 
 class TestFileDownload(ModuleTestBase):
-    targets = ["http://127.0.0.1:8888/Test_PDF"]
+    targets = ["http://127.0.0.1:8888"]
     modules_overrides = ["filedownload", "httpx", "excavate", "speculate"]
+    config_overrides = {"web_spider_distance": 2, "web_spider_depth": 2}
 
     pdf_data = """%PDF-1.
 1 0 obj<</Pages 2 0 R>>endobj
@@ -21,7 +22,10 @@ async def setup_before_prep(self, module_test):
 
     async def setup_after_prep(self, module_test):
         module_test.set_expect_requests(
-            dict(uri="/"), dict(response_data='<a href="/Test_File.txt"/><a href="/Test_PDF"/>')
+            dict(uri="/"),
+            dict(
+                response_data='<a href="/Test_File.txt"/><a href="/Test_PDF"/><a href="/test.html"/><a href="/test2"/>'
+            ),
         )
         module_test.set_expect_requests(
             dict(uri="/Test_File.txt"),
@@ -33,6 +37,14 @@ async def setup_after_prep(self, module_test):
             dict(uri="/Test_PDF"),
             dict(response_data=self.pdf_data, headers={"Content-Type": "application/pdf"}),
         )
+        module_test.set_expect_requests(
+            dict(uri="/test.html"),
+            dict(response_data="<!DOCTYPE html>", headers={"Content-Type": "text/html"}),
+        )
+        module_test.set_expect_requests(
+            dict(uri="/test2"),
+            dict(response_data="<!DOCTYPE html>", headers={"Content-Type": "text/html"}),
+        )
 
     def check(self, module_test, events):
         download_dir = module_test.scan.home / "filedownload"
@@ -50,3 +62,7 @@ def check(self, module_test, events):
         file = pdf_files[0]
         assert file.is_file(), f"File not found at {file}"
         assert open(file).read() == self.pdf_data, f"File at {file} does not contain the correct content"
+
+        # we don't want html files
+        html_files = list(download_dir.glob("*.html"))
+        assert len(html_files) == 0, "HTML files were erroneously downloaded"
