Merge pull request #73 from NullArray/trivial

Implementation for all

Author: NullArray

api_calls/censys.py

@@ -1,8 +1,6 @@
 import requests
-import threading
 
 import lib.settings
-from lib.output import error
 from lib.errors import AutoSploitAPIConnectionError
 from lib.settings import (
     HOST_FILE,
@@ -17,10 +15,12 @@ class CensysAPIHook(object):
     Censys API hook
     """
 
-    def __init__(self, identity, token, query):
+    def __init__(self, identity=None, token=None, query=None, proxy=None, agent=None, **kwargs):
         self.id = identity
         self.token = token
         self.query = query
+        self.proxy = proxy
+        self.user_agent = agent
         self.host_file = HOST_FILE
 
     def censys(self):
@@ -30,12 +30,15 @@ def censys(self):
         discovered_censys_hosts = set()
         try:
             lib.settings.start_animation("searching Censys with given query '{}'".format(self.query))
-            req = requests.post(API_URLS["censys"], auth=(self.id, self.token), json={"query": self.query})
+            req = requests.post(
+                API_URLS["censys"], auth=(self.id, self.token),
+                json={"query": self.query}, headers=self.user_agent,
+                proxies=self.proxy
+            )
             json_data = req.json()
             for item in json_data["results"]:
                 discovered_censys_hosts.add(str(item["ip"]))
             write_to_file(discovered_censys_hosts, self.host_file)
             return True
         except Exception as e:
-            error(AutoSploitAPIConnectionError(str(e)))
-            return False
+            raise AutoSploitAPIConnectionError(str(e))

api_calls/shodan.py

@@ -3,7 +3,6 @@
 import requests
 
 from lib.settings import start_animation
-from lib.output import error
 from lib.errors import AutoSploitAPIConnectionError
 from lib.settings import (
     API_URLS,
@@ -18,27 +17,30 @@ class ShodanAPIHook(object):
     Shodan API hook, saves us from having to install another dependency
     """
 
-    def __init__(self, token, query, proxy=None):
+    def __init__(self, token=None, query=None, proxy=None, agent=None, **kwargs):
         self.token = token
         self.query = query
         self.proxy = proxy
+        self.user_agent = agent
         self.host_file = HOST_FILE
 
     def shodan(self):
         """
         connect to the API and grab all IP addresses associated with the provided query
         """
-        start_animation("search Shodan with given query '{}'".format(self.query))
+        start_animation("searching Shodan with given query '{}'".format(self.query))
         discovered_shodan_hosts = set()
         try:
-            req = requests.get(API_URLS["shodan"].format(query=self.query, token=self.token))
+            req = requests.get(
+                API_URLS["shodan"].format(query=self.query, token=self.token),
+                proxies=self.proxy, headers=self.user_agent
+            )
             json_data = json.loads(req.content)
             for match in json_data["matches"]:
                 discovered_shodan_hosts.add(match["ip_str"])
             write_to_file(discovered_shodan_hosts, self.host_file)
             return True
         except Exception as e:
-            error(AutoSploitAPIConnectionError(str(e)))
-            return False
+            raise AutoSploitAPIConnectionError(str(e))
 
 

api_calls/zoomeye.py

@@ -6,7 +6,6 @@
 
 from lib.settings import start_animation
 from lib.errors import AutoSploitAPIConnectionError
-from lib.output import error
 from lib.settings import (
     API_URLS,
     HOST_FILE,
@@ -21,9 +20,11 @@ class ZoomEyeAPIHook(object):
     so we're going to use some 'lifted' credentials to login for us
     """
 
-    def __init__(self, query):
+    def __init__(self, query=None, proxy=None, agent=None, **kwargs):
         self.query = query
         self.host_file = HOST_FILE
+        self.proxy = proxy
+        self.user_agent = agent
         self.user_file = "{}/etc/text_files/users.lst".format(os.getcwd())
         self.pass_file = "{}/etc/text_files/passes.lst".format(os.getcwd())
 
@@ -61,9 +62,18 @@ def zoomeye(self):
         discovered_zoomeye_hosts = set()
         try:
             token = self.__get_auth()
-            headers = {"Authorization": "JWT {}".format(str(token["access_token"]))}
+            if self.user_agent is None:
+                headers = {"Authorization": "JWT {}".format(str(token["access_token"]))}
+            else:
+                headers = {
+                    "Authorization": "JWT {}".format(str(token["access_token"])),
+                    "agent": self.user_agent["User-Agent"]
+                }
             params = {"query": self.query, "page": "1", "facet": "ipv4"}
-            req = requests.get(API_URLS["zoomeye"][1].format(query=self.query), params=params, headers=headers)
+            req = requests.get(
+                API_URLS["zoomeye"][1].format(query=self.query),
+                params=params, headers=headers, proxies=self.proxy
+            )
             _json_data = req.json()
             for item in _json_data["matches"]:
                 if len(item["ip"]) > 1:
@@ -74,6 +84,5 @@ def zoomeye(self):
             write_to_file(discovered_zoomeye_hosts, self.host_file)
             return True
         except Exception as e:
-            error(AutoSploitAPIConnectionError(str(e)))
-            return False
+            raise AutoSploitAPIConnectionError(str(e))