Merge pull request #16 from spencer-hanson/master

Add requests session support

dnsdumpster/DNSDumpsterAPI.py

@@ -17,14 +17,17 @@ class DNSDumpsterAPI(object):
 
     """DNSDumpsterAPI Main Handler"""
 
-    def __init__(self, verbose=False):
+    def __init__(self, verbose=False, session=None):
         self.verbose = verbose
+        if not session:
+            self.session = requests.Session()
+        else:
+            self.session = session
 
     def display_message(self, s):
         if self.verbose:
             print('[verbose] %s' % s)
 
-
     def retrieve_results(self, table):
         res = []
         trs = table.findAll('tr')
@@ -63,17 +66,16 @@ def retrieve_txt_record(self, table):
 
     def search(self, domain):
         dnsdumpster_url = 'https://dnsdumpster.com/'
-        s = requests.session()
 
-        req = s.get(dnsdumpster_url)
+        req = self.session.get(dnsdumpster_url)
         soup = BeautifulSoup(req.content, 'html.parser')
         csrf_middleware = soup.findAll('input', attrs={'name': 'csrfmiddlewaretoken'})[0]['value']
         self.display_message('Retrieved token: %s' % csrf_middleware)
 
         cookies = {'csrftoken': csrf_middleware}
         headers = {'Referer': dnsdumpster_url}
         data = {'csrfmiddlewaretoken': csrf_middleware, 'targetip': domain}
-        req = s.post(dnsdumpster_url, cookies=cookies, data=data, headers=headers)
+        req = self.session.post(dnsdumpster_url, cookies=cookies, data=data, headers=headers)
 
         if req.status_code != 200:
             print(
@@ -101,7 +103,7 @@ def search(self, domain):
         # Network mapping image
         try:
             tmp_url = 'https://dnsdumpster.com/static/map/{}.png'.format(domain)
-            image_data = base64.b64encode(requests.get(tmp_url).content)
+            image_data = base64.b64encode(self.session.get(tmp_url).content)
         except:
             image_data = None
         finally:
@@ -112,7 +114,7 @@ def search(self, domain):
         try:
             pattern = r'https://dnsdumpster.com/static/xls/' + domain + '-[0-9]{12}\.xlsx'
             xls_url = re.findall(pattern, req.content.decode('utf-8'))[0]
-            xls_data = base64.b64encode(requests.get(xls_url).content)
+            xls_data = base64.b64encode(self.session.get(xls_url).content)
         except Exception as err:
             print(err)
             xls_data = None