chore(api): Add backoff

censys-workshop · Oct 23, 2023 · 4b43ba2 · 4b43ba2
1 parent dd8b255
commit 4b43ba2
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,6 +20,7 @@ sqlalchemy = "^2.0.22"
 pydantic-settings = "^2.0.3"
 psycopg2 = "^2.9.9"
 inquirerpy = "^0.3.4"
+backoff = "^2.2.1"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.10.0"

diff --git a/threatfox_censys/threatfox.py b/threatfox_censys/threatfox.py
@@ -1,11 +1,16 @@
 import logging
 from typing import Any
 
+import backoff
 import requests
 
 from .fingerprint import Fingerprint
 
 
+def fatal_code(e: requests.exceptions.RequestException) -> bool:
+    return 400 <= e.response.status_code < 500
+
+
 class ThreatFoxClient:
     """
     Client for the ThreatFox API.
@@ -35,6 +40,12 @@ def __init__(
         self.timeout = timeout
         self.headers = {"API-KEY": self.api_key, "Accept": "application/json"}
 
+    @backoff.on_exception(
+        backoff.expo,
+        requests.exceptions.RequestException,
+        max_time=60,
+        giveup=fatal_code,
+    )
     def _send_request(
         self, endpoint: str, method: str = "GET", data: Any | None = None
     ) -> dict: