asyncify sqlite

bbot/modules/gowitness.py

@@ -1,5 +1,5 @@
 import asyncio
-import sqlite3
+import aiosqlite
 import multiprocessing
 from pathlib import Path
 from contextlib import suppress
@@ -34,6 +34,7 @@ class gowitness(BaseModule):
         "idle_timeout": "Skip the current gowitness batch if it stalls for longer than this many seconds",
     }
     deps_common = ["chromium"]
+    deps_pip = ["aiosqlite"]
     deps_ansible = [
         {
             "name": "Download gowitness",
@@ -136,7 +137,8 @@ async def handle_batch(self, *events):
             return
 
         # emit web screenshots
-        for filename, screenshot in self.new_screenshots.items():
+        new_screenshots = await self.get_new_screenshots()
+        for filename, screenshot in new_screenshots.items():
             url = screenshot["url"]
             final_url = screenshot["final_url"]
             filename = self.screenshot_path / screenshot["filename"]
@@ -150,7 +152,8 @@ async def handle_batch(self, *events):
             )
 
         # emit URLs
-        for url, row in self.new_network_logs.items():
+        new_network_logs = await self.get_new_network_logs()
+        for url, row in new_network_logs.items():
             ip = row["ip"]
             status_code = row["status_code"]
             tags = [f"status-{status_code}", f"ip-{ip}", "spider-danger"]
@@ -168,7 +171,8 @@ async def handle_batch(self, *events):
                 )
 
         # emit technologies
-        for _, row in self.new_technologies.items():
+        new_technologies = await self.get_new_technologies()
+        for _, row in new_technologies.items():
             parent_id = row["url_id"]
             parent_url = self.screenshots_taken[parent_id]
             parent_event = event_dict[parent_url]
@@ -207,59 +211,53 @@ def construct_command(self):
         command += ["--timeout", str(self.timeout)]
         return command
 
-    @property
-    def new_screenshots(self):
+    async def get_new_screenshots(self):
         screenshots = {}
         if self.db_path.is_file():
-            with sqlite3.connect(str(self.db_path)) as con:
-                con.row_factory = sqlite3.Row
+            async with aiosqlite.connect(str(self.db_path)) as con:
+                con.row_factory = aiosqlite.Row
                 con.text_factory = self.helpers.smart_decode
-                cur = con.cursor()
-                res = self.cur_execute(cur, "SELECT * FROM urls")
-                for row in res:
-                    row = dict(row)
-                    _id = row["id"]
-                    if _id not in self.screenshots_taken:
-                        self.screenshots_taken[_id] = row["url"]
-                        screenshots[_id] = row
+                async with con.execute("SELECT * FROM urls") as cur:
+                    async for row in cur:
+                        row = dict(row)
+                        _id = row["id"]
+                        if _id not in self.screenshots_taken:
+                            self.screenshots_taken[_id] = row["url"]
+                            screenshots[_id] = row
         return screenshots
 
-    @property
-    def new_network_logs(self):
+    async def get_new_network_logs(self):
         network_logs = dict()
         if self.db_path.is_file():
-            with sqlite3.connect(str(self.db_path)) as con:
-                con.row_factory = sqlite3.Row
-                cur = con.cursor()
-                res = self.cur_execute(cur, "SELECT * FROM network_logs")
-                for row in res:
-                    row = dict(row)
-                    url = row["final_url"]
-                    if url not in self.connections_logged:
-                        self.connections_logged.add(url)
-                        network_logs[url] = row
+            async with aiosqlite.connect(str(self.db_path)) as con:
+                con.row_factory = aiosqlite.Row
+                async with con.execute("SELECT * FROM network_logs") as cur:
+                    async for row in cur:
+                        row = dict(row)
+                        url = row["final_url"]
+                        if url not in self.connections_logged:
+                            self.connections_logged.add(url)
+                            network_logs[url] = row
         return network_logs
 
-    @property
-    def new_technologies(self):
+    async def get_new_technologies(self):
         technologies = dict()
         if self.db_path.is_file():
-            with sqlite3.connect(str(self.db_path)) as con:
-                con.row_factory = sqlite3.Row
-                cur = con.cursor()
-                res = self.cur_execute(cur, "SELECT * FROM technologies")
-                for row in res:
-                    _id = row["id"]
-                    if _id not in self.technologies_found:
-                        self.technologies_found.add(_id)
-                        row = dict(row)
-                        technologies[_id] = row
+            async with aiosqlite.connect(str(self.db_path)) as con:
+                con.row_factory = aiosqlite.Row
+                async with con.execute("SELECT * FROM technologies") as cur:
+                    async for row in cur:
+                        _id = row["id"]
+                        if _id not in self.technologies_found:
+                            self.technologies_found.add(_id)
+                            row = dict(row)
+                            technologies[_id] = row
         return technologies
 
-    def cur_execute(self, cur, query):
+    async def cur_execute(self, cur, query):
         try:
-            return cur.execute(query)
-        except sqlite3.OperationalError as e:
+            return await cur.execute(query)
+        except aiosqlite.OperationalError as e:
             self.warning(f"Error executing query: {query}: {e}")
             return []