allow trusting certificate from custom certificate authorities (#328)

README.md

@@ -234,6 +234,10 @@ Setup KRR for...
 <p align="right">(<a href="#readme-top">back to top</a>)</p>
 
 
+**Trusting custom Certificate Authority (CA) certificate:**
+
+If your llm provider url uses a certificate from a custom CA, in order to trust it, base-64 encode the certificate, and store it in an environment variable named ``CERTIFICATE``
+
 ## Free KRR UI on Robusta SaaS
 
 We highly recommend using the [free Robusta SaaS platform](https://platform.robusta.dev/signup/?utm_source=github&utm_medium=krr-readme). You can:

krr.py

@@ -1,3 +1,15 @@
+import os
+
+from robusta_krr.common.ssl_utils import add_custom_certificate
+
+ADDITIONAL_CERTIFICATE: str = os.environ.get("CERTIFICATE", "")
+
+if add_custom_certificate(ADDITIONAL_CERTIFICATE):
+    print("added custom certificate")
+
+# DO NOT ADD ANY CODE ABOVE THIS
+# ADDING IMPORTS BEFORE ADDING THE CUSTOM CERTS MIGHT INIT HTTP CLIENTS THAT DOESN'T RESPECT THE CUSTOM CERT
+
 from robusta_krr import run
 
 if __name__ == "__main__":

robusta_krr/common/ssl_utils.py

@@ -0,0 +1,40 @@
+import base64
+import os
+
+import certifi
+
+CUSTOM_CERTIFICATE_PATH = "/tmp/custom_ca.pem"
+
+
+def append_custom_certificate(custom_ca: str) -> None:
+    with open(certifi.where(), "ab") as outfile:
+        outfile.write(base64.b64decode(custom_ca))
+
+    os.environ["WEBSOCKET_CLIENT_CA_BUNDLE"] = certifi.where()
+
+
+def create_temporary_certificate(custom_ca: str) -> None:
+    with open(certifi.where(), "rb") as base_cert:
+        base_cert_content = base_cert.read()
+
+    with open(CUSTOM_CERTIFICATE_PATH, "wb") as outfile:
+        outfile.write(base_cert_content)
+        outfile.write(base64.b64decode(custom_ca))
+
+    os.environ["REQUESTS_CA_BUNDLE"] = CUSTOM_CERTIFICATE_PATH
+    os.environ["WEBSOCKET_CLIENT_CA_BUNDLE"] = CUSTOM_CERTIFICATE_PATH
+    certifi.where = lambda: CUSTOM_CERTIFICATE_PATH
+
+
+def add_custom_certificate(custom_ca: str) -> bool:
+    if not custom_ca:
+        return False
+
+    # NOTE: Sometimes (Openshift) the certifi.where() is not writable, so we need to
+    #       use a temporary file in case of PermissionError.
+    try:
+        append_custom_certificate(custom_ca)
+    except PermissionError:
+        create_temporary_certificate(custom_ca)
+
+    return True