From a9558ec1fdd1d0dd85a6588fd56b76852176403d Mon Sep 17 00:00:00 2001
From: Kent Ickler <kickler@blackhillsinfosec.com>
Date: Thu, 13 Jun 2024 18:11:30 -0700
Subject: [PATCH] handle timeouts from urllib (headers/source) (#668)

* timeout handling on urllib (headers & source)

* report: include screenshot from selenium when it exists even if urllib timed-out
---
 Python/modules/objects.py         | 10 ++++++++--
 Python/modules/selenium_module.py |  9 +++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/Python/modules/objects.py b/Python/modules/objects.py
index 3f8fca0..13df909 100644
--- a/Python/modules/objects.py
+++ b/Python/modules/objects.py
@@ -262,8 +262,14 @@ def create_table_html(self):
             </tr>
             """)
         elif self.error_state == 'Timeout':
-            html += ("""</td><td>Hit timeout limit while attempting to
-            screenshot</td></tr>""")
+            html += ("</td><td>Hit timeout limit")
+            if os.path.isfile(self.screenshot_path):
+                html += ("""<br>
+                <div id=\"screenshot\"><a href=\"{1}\"
+                target=\"_blank\"><img style=\"max-height:400px;height: expression(this.height > 400 ? 400: true);\"
+                src=\"{1}\"></a></div></td></tr>""").format(src_path, scr_path)
+            else:
+                html += ("</td></tr>")
         elif self.error_state == 'BadStatus':
             html += ("""</td><td>Unknown error while attempting to
             screenshot</td></tr>""")
diff --git a/Python/modules/selenium_module.py b/Python/modules/selenium_module.py
index 4327325..c9f3f2f 100644
--- a/Python/modules/selenium_module.py
+++ b/Python/modules/selenium_module.py
@@ -219,9 +219,9 @@ def capture_host(cli_parsed, http_object, driver, ua=None):
             req.set_proxy(str(cli_parsed.proxy_ip) + ':' + str(cli_parsed.proxy_port), 'http')
             req.set_proxy(str(cli_parsed.proxy_ip) + ':' + str(cli_parsed.proxy_port), 'https')
         if context is None:
-            opened = urllib.request.urlopen(req)
+            opened = urllib.request.urlopen(req, timeout=cli_parsed.timeout)
         else:
-            opened = urllib.request.urlopen(req, context=context)
+            opened = urllib.request.urlopen(req,timeout=cli_parsed.timeout, context=context)
         headers = dict(opened.info())
         headers['Response Code'] = str(opened.getcode())
     except urllib.error.HTTPError as e:
@@ -266,6 +266,11 @@ def capture_host(cli_parsed, http_object, driver, ua=None):
             headers = {'Error': 'Connection Reset'}
             http_object.error_state = 'ConnReset'
             return http_object, driver
+        elif 'timed out' in str(e):
+            headers = {'Error': 'Timed Out'}
+            http_object.error_state = 'Timeout'
+            print('[*] Socket Timeout when connecting to {0}'.format(http_object.remote_system))
+            return http_object, driver
         else:
             http_object.error_state = 'BadStatus'
             return http_object, driver