From 8b71471fc34b72e3d611ca9bbfc7446f2b94234b Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Mon, 2 Oct 2023 15:00:49 -0600
Subject: [PATCH] fix media cdn

---
 check_media_cdn.py | 30 +++++++++++++++++++++++++++---
 check_vllm.py      |  2 +-
 icinga2kuma.py     |  3 +++
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/check_media_cdn.py b/check_media_cdn.py
index 9256d56..1ac614c 100755
--- a/check_media_cdn.py
+++ b/check_media_cdn.py
@@ -5,6 +5,7 @@ import json
 import os
 import sys
 import tempfile
+import time
 import traceback
 import urllib
 
@@ -27,6 +28,7 @@ parser.add_argument('--check-domain', required=True, help='The domain that shoul
 parser.add_argument('--media-cdn-redirect', default='true', help='If set, the server must respond with a redirect to the media CDN domain.')
 parser.add_argument('--required-headers', nargs='*', help="If these headers aren't set to the correct value, critical. Use the format 'key=value")
 parser.add_argument('--auth-file', help="File to cache the bot's login details to.")
+parser.add_argument('--retries', type=int, default=5, help="It may take a few seconds for Synapse to send the uploaded file to S3. Retry this many times with 1 second interval between.")
 parser.add_argument('--timeout', type=float, default=90, help='Request timeout limit.')
 parser.add_argument('--warn', type=float, default=2.0, help='Manually set warn level.')
 parser.add_argument('--crit', type=float, default=2.5, help='Manually set critical level.')
@@ -141,7 +143,27 @@ async def main() -> None:
 
     # Check the file. Ignore the non-async thing here, it doesn't matter in this situation.
     # Remember: Cloudflare does not cache non-GET requests.
-    r = requests.head(target_file_url, allow_redirects=False)
+
+    start = time.time()
+    retried = 0
+    for i in range(args.retries):
+        if i % 5 == 0 and i != 0:
+            r = requests.get(target_file_url, allow_redirects=False)
+        else:
+            r = requests.head(target_file_url, allow_redirects=False)
+        headers = dict(r.headers)
+        if len(args.required_headers) == 1:
+            args.required_headers = args.required_headers[0].split(' ')
+        success = []
+        for item in args.required_headers:
+            key, value = item.split('=')
+            _, code = verify_media_header(key, headers, good_value=value)
+            success.append(code)
+        if sum(success) == 0:
+            break
+        time.sleep(1)
+        retried += 1
+    end = time.time()
 
     prints = []
 
@@ -152,8 +174,6 @@ async def main() -> None:
     else:
         prints.append(f'OK: status code is "{r.status_code}"')
 
-    headers = dict(r.headers)
-
     # Check domain
     if args.media_cdn_redirect:
         if 'location' in headers:
@@ -211,9 +231,13 @@ async def main() -> None:
     for msg in prints:
         print(msg)
 
+    print(f'Took {int(end - start)} seconds to validate the file with {retried} retries.')
+
     if clean_msg:
         print(clean_msg)
 
+    print(f'|retries={retried}')
+
     sys.exit(exit_code)
 
 
diff --git a/check_vllm.py b/check_vllm.py
index d248319..d2897ee 100644
--- a/check_vllm.py
+++ b/check_vllm.py
@@ -20,7 +20,7 @@ def main(critical, warning, timeout, target_url, does_not_contain, verify):
         "prompt": "Test prompt",
         "stream": False,
         "temperature": 0,
-        "max_tokens": 50,
+        "max_tokens": 3,
     }
 
     start_time = time.time()
diff --git a/icinga2kuma.py b/icinga2kuma.py
index b26c518..2a6945e 100755
--- a/icinga2kuma.py
+++ b/icinga2kuma.py
@@ -9,6 +9,9 @@ from icinga2api.client import Client
 
 from checker import nagios
 
+# TODO: add query arg to not return an error when the host check is down
+
+
 endpoint = 'https://localhost:8080'  # Icinga2 URL for the API. Defaults to "https://localhost:8080"
 icinga2_user = 'icingaweb2'  # API username. Defaults to "icingaweb2"
 icinga2_pw = ''  # API password or set ICINGA2KUMA_ICINGA2_PW