diff --git a/README.md b/README.md index 43baa8e..44798e5 100644 --- a/README.md +++ b/README.md @@ -234,6 +234,10 @@ Setup KRR for...
+**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: diff --git a/krr.py b/krr.py index 3c7ac1e..319ee61 100644 --- a/krr.py +++ b/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__": diff --git a/robusta_krr/common/ssl_utils.py b/robusta_krr/common/ssl_utils.py new file mode 100644 index 0000000..fd77d61 --- /dev/null +++ b/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