From 503b378c15682b0a842fa439fcd5fb6e44f14aef Mon Sep 17 00:00:00 2001
From: Sarabveer Singh <me@sarabveer.me>
Date: Sat, 13 Apr 2024 22:24:46 -0400
Subject: [PATCH] fix: add timeout to requests to prevent hangs

---
 README.md                 |  1 +
 config.sample.ini         |  1 +
 src/__main__.py           |  1 +
 src/arris_stats_s33.py    |  9 ++++++---
 src/arris_stats_sb8200.py | 21 ++++++++++++++++++---
 src/comcast_xb8_stats.py  |  9 +++++++--
 6 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index c5c5cc1..c5ec260 100644
--- a/README.md
+++ b/README.md
@@ -88,6 +88,7 @@ Config settings can be provided by the config.ini file, or set as environment va
 | `exit_on_html_error` | `True` | Any error retrieving tdata will cause an exit, mostly redundant with exit_on_auth_error |
 | `clear_auth_token_on_html_error` | `True` | This is useful if you don't want to exit, but do want to get a new session if/when getting the stats fails |
 | `sleep_before_exit` | `True` | If you want to sleep before exiting on errors, useful for Docker container when you have `restart = always` |
+| `request_timeout` | `30` | Seconds to wait before request to fetch modem webpage/data times out |
 
 ### InfluxDB Config
 
diff --git a/config.sample.ini b/config.sample.ini
index abbba6b..539a97e 100644
--- a/config.sample.ini
+++ b/config.sample.ini
@@ -10,6 +10,7 @@ exit_on_auth_error = True
 exit_on_html_error = True
 clear_auth_token_on_html_error = True
 sleep_before_exit = True
+request_timeout = 30
 
 # SB8200 Only
 modem_ssl = False
diff --git a/src/__main__.py b/src/__main__.py
index 01cab84..630ad63 100644
--- a/src/__main__.py
+++ b/src/__main__.py
@@ -127,6 +127,7 @@ def get_config(config_path=None):
     'exit_on_html_error': True,
     'clear_auth_token_on_html_error': True,
     'sleep_before_exit': True,
+    'request_timeout': 30,
 
     # SB8200 Only
     'modem_ssl': False,
diff --git a/src/arris_stats_s33.py b/src/arris_stats_s33.py
index 53c5f44..24ed2a3 100644
--- a/src/arris_stats_s33.py
+++ b/src/arris_stats_s33.py
@@ -46,7 +46,8 @@ def get_credential(config):
       url=url,
       json=payload,
       headers=headers,
-      verify=verify_ssl
+      verify=verify_ssl,
+      timeout=config['request_timeout']
     )
 
     if resp.status_code != 200:
@@ -84,7 +85,8 @@ def get_credential(config):
       url=url,
       json=payload,
       headers=headers,
-      verify=verify_ssl
+      verify=verify_ssl,
+      timeout=config['request_timeout']
     )
 
     if resp.status_code != 200:
@@ -144,7 +146,8 @@ def get_json(config, credential):
       url=url,
       json=payload,
       headers=headers,
-      verify=verify_ssl
+      verify=verify_ssl,
+      timeout=config['request_timeout']
     )
     if resp.status_code != 200:
       logging.error('Error retreiving json from %s', url)
diff --git a/src/arris_stats_sb8200.py b/src/arris_stats_sb8200.py
index feab5f1..8c06e0c 100644
--- a/src/arris_stats_sb8200.py
+++ b/src/arris_stats_sb8200.py
@@ -42,11 +42,21 @@ def get_credential(config):
   # have to send as a cookie with subsequent requests
   try:
     if config['modem_new_auth']:
-      resp = requests.get(auth_url, headers={'Authorization': 'Basic ' + auth_hash}, verify=verify_ssl)
+      resp = requests.get(
+        auth_url,
+        headers={'Authorization': 'Basic ' + auth_hash},
+        verify=verify_ssl,
+        timeout=config['request_timeout']
+      )
       cookie = resp.cookies['sessionId']
       logging.debug('cookie: %s', cookie)
     else:
-      resp = requests.get(auth_url, auth=(username, password), verify=verify_ssl)
+      resp = requests.get(
+        auth_url,
+        auth=(username, password),
+        verify=verify_ssl,
+        timeout=config['request_timeout']
+      )
       cookie = None
 
     if resp.status_code != 200:
@@ -99,7 +109,12 @@ def get_html(config, credential):
   logging.info('Retreiving stats from %s', init_url)
 
   try:
-    resp = requests.get(url, cookies=cookies, verify=verify_ssl)
+    resp = requests.get(
+      url,
+      cookies=cookies,
+      verify=verify_ssl,
+      timeout=config['request_timeout']
+    )
     if resp.status_code != 200:
       logging.error('Error retreiving html from %s', url)
       logging.error('Status code: %s', resp.status_code)
diff --git a/src/comcast_xb8_stats.py b/src/comcast_xb8_stats.py
index d299192..a535f25 100644
--- a/src/comcast_xb8_stats.py
+++ b/src/comcast_xb8_stats.py
@@ -24,7 +24,12 @@ def get_credential(config):
   }
 
   try:
-    resp = requests.post(url, data=data, allow_redirects=False)
+    resp = requests.post(
+      url,
+      data=data,
+      allow_redirects=False,
+      timeout=config['request_timeout']
+    )
     cookies = resp.cookies
 
     if resp.status_code != 302:
@@ -52,7 +57,7 @@ def get_html(config, cookies):
   logging.info('Retreiving stats from %s', url)
 
   try:
-    resp = requests.get(url, cookies=cookies)
+    resp = requests.get(url, cookies=cookies, timeout=config['request_timeout'])
     if resp.status_code != 200:
       logging.error('Error retreiving html from %s', url)
       logging.error('Status code: %s', resp.status_code)