From f4081ebd3b3b1d759c36d71d67c6b7ad28b02408 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9ophile=20Diot?= <theophile.diot900@gmail.com>
Date: Thu, 16 Feb 2023 13:19:50 +0100
Subject: [PATCH] Handle more errors with Bunkernet job

---
 .../core/bunkernet/jobs/bunkernet-data.py     |  5 +++++
 .../core/bunkernet/jobs/bunkernet-register.py | 19 ++++++++++++++-----
 src/common/core/bunkernet/jobs/bunkernet.py   | 17 +++++++++++++----
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/src/common/core/bunkernet/jobs/bunkernet-data.py b/src/common/core/bunkernet/jobs/bunkernet-data.py
index 65574b59b..80a4aa3f2 100755
--- a/src/common/core/bunkernet/jobs/bunkernet-data.py
+++ b/src/common/core/bunkernet/jobs/bunkernet-data.py
@@ -78,6 +78,11 @@ try:
             "BunkerNet API is rate limiting us, trying again later...",
         )
         _exit(0)
+    elif status == 403:
+        logger.warning(
+            "BunkerNet has banned this instance, retrying a register later...",
+        )
+        _exit(0)
     elif data["result"] != "ok":
         logger.error(
             f"Received error from BunkerNet API while sending db request : {data['data']}, removing instance ID",
diff --git a/src/common/core/bunkernet/jobs/bunkernet-register.py b/src/common/core/bunkernet/jobs/bunkernet-register.py
index 4095a4ccb..741068133 100755
--- a/src/common/core/bunkernet/jobs/bunkernet-register.py
+++ b/src/common/core/bunkernet/jobs/bunkernet-register.py
@@ -66,14 +66,19 @@ try:
                 "BunkerNet API is rate limiting us, trying again later...",
             )
             _exit(0)
+        elif status == 403:
+            logger.warning(
+                "BunkerNet has banned this instance, retrying a register later...",
+            )
+            _exit(0)
         elif status != 200:
             logger.error(
                 f"Error {status} from BunkerNet API : {data['data']}",
             )
             _exit(1)
-        elif data["result"] != "ok":
+        elif data.get("result", "ko") != "ok":
             logger.error(
-                f"Received error from BunkerNet API while sending register request : {data['data']}"
+                f"Received error from BunkerNet API while sending register request : {data.get('data', {})}"
             )
             _exit(1)
         bunkernet_id = data["data"]
@@ -98,15 +103,19 @@ try:
                 "BunkerNet API is rate limiting us, trying again later...",
             )
             retry = True
+        elif status == 403:
+            logger.warning(
+                "BunkerNet has banned this instance, retrying a register later...",
+            )
         elif status == 401:
             logger.warning(
                 "Instance ID is not registered, removing it and retrying a register later...",
             )
             remove("/var/cache/bunkerweb/bunkernet/instance.id")
             _exit(2)
-        elif data["result"] != "ok":
+        elif data.get("result", "ko") != "ok":
             logger.error(
-                f"Received error from BunkerNet API while sending ping request : {data['data']}, removing instance ID",
+                f"Received error from BunkerNet API while sending ping request : {data.get('data', {})}, removing instance ID",
             )
             retry = True
         if not retry:
@@ -115,7 +124,7 @@ try:
         logger.warning("Waiting 1s and trying again ...")
         sleep(1)
 
-    if bunkernet_ping:
+    if bunkernet_ping and status != 403:
         logger.info("Connectivity with BunkerWeb is successful !")
         status = 1
         if not isfile("/var/cache/bunkerweb/bunkernet/instance.id"):
diff --git a/src/common/core/bunkernet/jobs/bunkernet.py b/src/common/core/bunkernet/jobs/bunkernet.py
index eff3b6005..c8dc908d3 100644
--- a/src/common/core/bunkernet/jobs/bunkernet.py
+++ b/src/common/core/bunkernet/jobs/bunkernet.py
@@ -1,9 +1,12 @@
-import requests, traceback
+from typing import Literal, Optional, Tuple, Union
+import requests
 from os import getenv
 from os.path import exists
 
 
-def request(method, url, _id=None):
+def request(
+    method: Union[Literal["POST"], Literal["GET"]], url: str, _id: Optional[str] = None
+) -> Tuple[bool, Optional[int], Union[str, dict]]:
     data = {"integration": get_integration(), "version": get_version()}
     headers = {"User-Agent": f"BunkerWeb/{get_version()}"}
     if _id is not None:
@@ -19,11 +22,17 @@ def request(method, url, _id=None):
         status = resp.status_code
         if status == 429:
             return True, 429, "rate limited"
-        raw_data = resp.json()
+        elif status == 403:
+            return True, 403, "forbidden"
+
+        raw_data: dict = resp.json()
+
         assert "result" in raw_data
         assert "data" in raw_data
+    except requests.ReadTimeout:
+        return False, None, "request timed out"
     except Exception as e:
-        return False, None, traceback.format_exc()
+        return False, None, f"request failed: {e}"
     return True, status, raw_data