Merge pull request #496 from blacklanternsecurity/async-optimize

Explicitly cleanup http client

baddns/base.py

@@ -32,6 +32,9 @@ def infomsg(self, msg):
         else:
             log.debug(msg)
 
+    async def cleanup(self):
+        pass
+
 
 def get_all_modules(*args, **kwargs):
     return [m for m in BadDNS_base.__subclasses__()]

baddns/cli.py

@@ -91,6 +91,7 @@ async def execute_module(ModuleClass, target, custom_nameservers, signatures, si
             for finding in findings:
                 print(finding.to_dict())
 
+    await module_instance.cleanup()
     return findings
 
 

baddns/lib/httpmanager.py

@@ -64,3 +64,8 @@ async def dispatchHttp(self):
             else:
                 setattr(self, f"{protocol}_denyredirects_results", result)
             idx += 1
+
+    async def close(self):
+        """Clean up the http_client when done."""
+        await self.http_client.aclose()
+        log.debug("HTTP client closed successfully.")

baddns/modules/cname.py

@@ -194,3 +194,8 @@ def analyze(self):
                 )
 
         return findings
+
+    async def cleanup(self):
+        if self.target_httpmanager:
+            await self.target_httpmanager.close()
+            log.debug("HTTP Manager cleaned up successfully.")

baddns/modules/references.py

@@ -214,3 +214,8 @@ def analyze(self):
         if self.cname_findings_direct:
             findings.extend(self._convert_findings(self.cname_findings_direct))
         return findings
+
+    async def cleanup(self):
+        if self.target_httpmanager:
+            await self.target_httpmanager.close()
+            log.debug("HTTP Manager cleaned up successfully.")

baddns/modules/txt.py

@@ -110,3 +110,8 @@ def analyze(self):
         if self.cname_findings:
             findings.extend(self._convert_findings(self.cname_findings))
         return findings
+
+    async def cleanup(self):
+        if self.target_httpmanager:
+            await self.target_httpmanager.close()
+            log.debug("HTTP Manager cleaned up successfully.")