From a327e44fefbec7e609003b010d896e2c411d9fc3 Mon Sep 17 00:00:00 2001
From: Jonathan M Dolbee <jonathandolbee@gmail.com>
Date: Thu, 24 Feb 2022 21:30:29 -0600
Subject: [PATCH] Added timeout to session requests

I've been having issues lately with shodan.count with the python client.
It seems to hang at this get request. I went ahead and gave the get request a timeout of 1, while allowing retries up to five times.
This works for me, if there is a better way please feel free to let me know!

Most of the time, it simply needs one retry. I played with the timeout time, and anything above1 takes way to long, 1 seemed to be the best.
---
 shodan/client.py | 38 ++++++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/shodan/client.py b/shodan/client.py
index db50b4c..c0add8f 100644
--- a/shodan/client.py
+++ b/shodan/client.py
@@ -304,20 +304,30 @@ def _request(self, function, params, service='shodan', method='get'):
             while (1.0 / self.api_rate_limit) + self._api_query_time >= time.time():
                 time.sleep(0.1 / self.api_rate_limit)
 
-        # Send the request
-        try:
-            method = method.lower()
-            if method == 'post':
-                data = self._session.post(base_url + function, params)
-            elif method == 'put':
-                data = self._session.put(base_url + function, params=params)
-            elif method == 'delete':
-                data = self._session.delete(base_url + function, params=params)
-            else:
-                data = self._session.get(base_url + function, params=params)
-            self._api_query_time = time.time()
-        except Exception:
-            raise APIError('Unable to connect to Shodan')
+        # Determine options for timeout settings and retrying
+        MAX_RETRIES = 5
+        TIMEOUT_SECONDS = 1
+        current_retries = 0
+        finished_request = False
+        while(not finished_request):
+            # Send the request
+            try:
+                method = method.lower()
+                if method == 'post':
+                    data = self._session.post(base_url + function, params, timeout=TIMEOUT_SECONDS)
+                elif method == 'put':
+                    data = self._session.put(base_url + function, params=params, timeout=TIMEOUT_SECONDS)
+                elif method == 'delete':
+                    data = self._session.delete(base_url + function, params=params, timeout=TIMEOUT_SECONDS)
+                else:
+                    data = self._session.get(base_url + function, params=params, timeout=TIMEOUT_SECONDS)
+                finished_request = True
+            except requests.exceptions.RequestException as e:
+                current_retries+=1
+                if current_retries >= MAX_RETRIES:
+                    raise APIError('Max retries reached!')
+            except Exception:
+                raise APIError('Unable to connect to Shodan')
 
         # Check that the API key wasn't rejected
         if data.status_code == 401: