Issue #83: Use Asyncio (VirusTotal).

Nekmo · Aug 11, 2023 · 51b5822 · 51b5822
1 parent 9da7d93
commit 51b5822
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 18 deletions.
diff --git a/dirhunt/sources/__init__.py b/dirhunt/sources/__init__.py
@@ -18,7 +18,7 @@
 
 SOURCE_CLASSES: List[Type["SourceBase"]] = [
     Robots,
-    # VirusTotal,
+    VirusTotal,
     Google,
     CommonCrawl,
     CrtSh,

diff --git a/dirhunt/sources/virustotal.py b/dirhunt/sources/virustotal.py
@@ -1,8 +1,9 @@
 import string
+from typing import Iterable
+
 from bs4 import BeautifulSoup
-from requests import RequestException
 
-from dirhunt.sessions import Sessions
+from dirhunt.exceptions import SourceError
 from dirhunt.sources.base import SourceBase
 
 
@@ -12,19 +13,16 @@
 
 
 class VirusTotal(SourceBase):
-    def callback(self, domain):
+    async def search_by_domain(self, domain: str) -> Iterable[str]:
+        """Search by domain in VirusTotal."""
         url = VT_URL.format(domain=domain)
-        session = Sessions().get_session()
-        try:
-            with session.get(url) as response:
-                html = response.text
-        except RequestException as e:
-            self.add_error("Error on Crt.sh source: {}".format(e))
-            return
+        async with self.sources.crawler.session.get(url) as response:
+            response.raise_for_status()
+            html = await response.text()
         if ABUSE in html:
-            self.add_error(ABUSE_MESSAGE_ERROR.format(url=url))
-            return
+            raise SourceError(ABUSE_MESSAGE_ERROR.format(url=url))
         soup = BeautifulSoup(html, "html.parser")
-
-        for url in soup.select("#detected-urls .enum a"):
-            self.add_result(url.text.strip(string.whitespace))
+        return [
+            url.text.strip(string.whitespace)
+            for url in soup.select("#detected-urls .enum a")
+        ]
diff --git a/dirhunt/tests/test_sources.py b/dirhunt/tests/test_sources.py
@@ -95,7 +95,7 @@ def test_urls(self, m):
             ),
         )
         with patch.object(VirusTotal, "add_result") as mock_add_result:
-            VirusTotal(lambda x: x, None).callback(domain)
+            VirusTotal(lambda x: x, None).search_by_domain(domain)
             mock_add_result.assert_has_calls(
                 [call(detect_url) for detect_url in detect_urls]
             )
@@ -106,7 +106,7 @@ def test_abuse(self, m):
         url = VT_URL.format(domain=domain)
         m.get(url, text=ABUSE)
         with patch.object(VirusTotal, "add_error") as mock_add_error:
-            VirusTotal(lambda x: x, lambda x: x).callback(domain)
+            VirusTotal(lambda x: x, lambda x: x).search_by_domain(domain)
             mock_add_error.assert_called()