From 0dae784ce784cdfbb128e54ea9fdd709873b4358 Mon Sep 17 00:00:00 2001
From: NotoriousRebel <matt12299@yahoo.com>
Date: Thu, 7 May 2020 19:59:14 -0400
Subject: [PATCH] Fixed pentesttools module to properly work.

---
 theHarvester/__main__.py               |  4 +-
 theHarvester/discovery/pentesttools.py | 69 +++++++++++++++++++++++---
 theHarvester/lib/core.py               |  1 -
 3 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/theHarvester/__main__.py b/theHarvester/__main__.py
index 1e376916..06f0dfb2 100644
--- a/theHarvester/__main__.py
+++ b/theHarvester/__main__.py
@@ -37,7 +37,7 @@ async def start():
     parser.add_argument('-b', '--source', help='''baidu, bing, bingapi, bufferoverun, certspotter, crtsh, dnsdumpster,
                         dogpile, duckduckgo, exalead, github-code, google,
                         hackertarget, hunter, intelx,
-                        linkedin, linkedin_links, netcraft, otx, securityTrails, spyse, threatcrowd,
+                        linkedin, linkedin_links, netcraft, otx, pentesttools, securityTrails, spyse, threatcrowd,
                         trello, twitter, vhost, virustotal, yahoo, all''')
 
     args = parser.parse_args()
@@ -99,7 +99,7 @@ async def start():
             print(f'\033[94m[*] Searching {source[0].upper() + source[1:]}. \033[0m')
         if store_host:
             host_names = filter(await search_engine.get_hostnames())
-            if source != 'hackertarget':
+            if source != 'hackertarget' and source != 'pentesttools':
                 full_hosts_checker = hostchecker.Checker(host_names)
                 temp_hosts, temp_ips = await full_hosts_checker.check()
                 ips.extend(temp_ips)
diff --git a/theHarvester/discovery/pentesttools.py b/theHarvester/discovery/pentesttools.py
index 2e0beeb9..315c98bc 100644
--- a/theHarvester/discovery/pentesttools.py
+++ b/theHarvester/discovery/pentesttools.py
@@ -1,31 +1,86 @@
 from theHarvester.discovery.constants import *
 from theHarvester.lib.core import *
 import json
-
+import time
+from timeit import default_timer as timer
 
 class SearchPentestTools:
 
     def __init__(self, word):
+        # Script is largely based off https://pentest-tools.com/public/api_client.py.txt
         self.word = word
         self.key = Core.pentest_tools_key()
         if self.key is None:
             raise MissingKey(True)
-        self.total_results = ""
+        self.total_results = []
         self.api = f'https://pentest-tools.com/api?key={self.key}'
         self.proxy = False
 
+    async def poll(self, scan_id):
+        start = timer()
+        while True:
+            current_time = timer()
+            if current_time - start >= 300:
+                status = input('Pentesttools scan has reached 5 minute mark would you like to stop scan: y/n')
+                if status == 'y':
+                    return self.total_results
+            time.sleep(3)
+            # Get the status of our scan
+            scan_status_data = {
+                "op": "get_scan_status",
+                "scan_id": scan_id
+            }
+            responses = await AsyncFetcher.post_fetch(url=self.api,
+                                                      data=json.dumps(scan_status_data),
+                                                      proxy=self.proxy)
+            res_json = json.loads(responses.strip())
+            if res_json["op_status"] == "success":
+                if res_json["scan_status"] != "waiting" and res_json["scan_status"] != "running":
+                    getoutput_data = {
+                        "op": "get_output",
+                        "scan_id": scan_id,
+                        "output_format": 'json'
+                    }
+                    responses = await AsyncFetcher.post_fetch(url=self.api,
+                                                              data=json.dumps(getoutput_data),
+                                                              proxy=self.proxy)
+                    res_json = json.loads(responses.strip())
+                    self.total_results = await self.parse_json(res_json)
+                    break
+            else:
+                print(f"Operation get_scan_status failed because: {res_json['error']}. {res_json['details']}")
+                break
+
+    @staticmethod
+    async def parse_json(json_results):
+        status = json_results['op_status']
+        if status == 'success':
+            scan_tests = json_results['scan_output']['scan_tests']
+            output_data = scan_tests[0]['output_data']
+            host_to_ip = [f'{subdomain[0]}:{subdomain[1]}' for subdomain in output_data if len(subdomain) > 0]
+            return host_to_ip
+        return []
+
+    async def get_hostnames(self) -> list:
+        return self.total_results
+
     async def do_search(self):
         subdomain_payload = {
             'op': 'start_scan',
             'tool_id': 20,
             'tool_params': {
-                'target': f'{self.word}'}
+                'target': f'{self.word}'},
+            "scan_type": "quick"
         }
-
-        responses = await AsyncFetcher.post_fetch(url=self.api, data=json.dumps(subdomain_payload))
-        dct = responses
-        print(dct)
+        responses = await AsyncFetcher.post_fetch(url=self.api,
+                                                  data=json.dumps(subdomain_payload),
+                                                  proxy=self.proxy)
+        res_json = json.loads(responses.strip())
+        if res_json["op_status"] == "success":
+            scan_id = res_json['scan_id']
+            await self.poll(scan_id)
 
     async def process(self, proxy=False):
+        print('Pentesttools may take over 10 minutes to run!')
         self.proxy = proxy
         await self.do_search()  # Only need to do it once.
diff --git a/theHarvester/lib/core.py b/theHarvester/lib/core.py
index 671e8aba..0f29769b 100644
--- a/theHarvester/lib/core.py
+++ b/theHarvester/lib/core.py
@@ -402,7 +402,6 @@ class Core:
 
 class AsyncFetcher:
     proxy_list = Core.proxy_list()
-    print(f'proxy_list: {proxy_list}')
 
     @classmethod
     async def post_fetch(cls, url, headers='', data='', params='', json=False, proxy=False):