Merge pull request #955 from netenglabs/fix-http-timeout

Node: fix `timeout` type error when using http as transport
netenglabs · Jun 3, 2024 · 9362ae2 · 9362ae2
2 parents 1ce2a5c + efab7d5
commit 9362ae2
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/suzieq/poller/worker/nodes/node.py b/suzieq/poller/worker/nodes/node.py
@@ -1338,15 +1338,15 @@ class EosNode(Node):
     async def _rest_connect(self):
         '''Check that connectivity and authentication works'''
 
-        timeout = self.cmd_timeout
+        timeout = aiohttp.ClientTimeout(self.cmd_timeout)
 
         auth = aiohttp.BasicAuth(self.username,
                                  password=self.password or 'vagrant')
         url = f"https://{self.address}:{self.port}/command-api"
 
         try:
             self._conn = aiohttp.ClientSession(
-                auth=auth, timeout=self.connect_timeout,
+                auth=auth, timeout=timeout,
                 connector=aiohttp.TCPConnector(ssl=False))
             async with self._conn.post(url, timeout=timeout) as response:
                 _ = response.status
@@ -1414,7 +1414,7 @@ async def _rest_gather(self, service_callback, cmd_list, cb_token,
                 await self._post_result(service_callback, result, cb_token)
                 return
 
-        timeout = timeout or self.cmd_timeout
+        timeout = aiohttp.ClientTimeout(timeout or self.cmd_timeout)
 
         now = int(time.time() * 1000)
         data = {
@@ -2154,8 +2154,8 @@ async def get_api_key(self):
         if not self._retry:
             return
         async with self._cmd_pacer.wait(self.per_cmd_auth):
-            async with self._conn.get(url, timeout=self.connect_timeout) \
-                    as response:
+            timeout = aiohttp.ClientTimeout(self.connect_timeout)
+            async with self._conn.get(url, timeout=timeout) as response:
                 status, xml = response.status, await response.text()
                 if status == 200:
                     data = xmltodict.parse(xml)
@@ -2215,9 +2215,10 @@ async def _rest_connect(self):
         # In case of PANOS, getting here means REST is up
         if not self._conn:
             async with self._cmd_pacer.wait(self.per_cmd_auth):
+                timeout = aiohttp.ClientTimeout(self.connect_timeout)
                 try:
                     self._conn = aiohttp.ClientSession(
-                        conn_timeout=self.connect_timeout,
+                        conn_timeout=timeout,
                         connector=aiohttp.TCPConnector(ssl=False),
                     )
                     if self.api_key is None:
@@ -2240,7 +2241,7 @@ async def _rest_gather(self, service_callback, cmd_list, cb_token,
         if not cmd_list:
             return result
 
-        timeout = timeout or self.connect_timeout
+        timeout = aiohttp.ClientTimeout(timeout or self.connect_timeout)
 
         url = f"https://{self.address}:{self.port}/api/"