From 4c99010eeea1d21e1370326fd1dcd16133b4b99a Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 25 Jun 2024 13:55:54 -0700
Subject: [PATCH] fix - verify license without api request

---
 litellm/proxy/auth/litellm_license.py | 65 +++++++++++++++++++++++++++
 litellm/proxy/auth/public_key.pem     |  9 ++++
 requirements.txt                      |  1 +
 3 files changed, 75 insertions(+)
 create mode 100644 litellm/proxy/auth/public_key.pem

diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py
index ffd9f5273..ec51f904c 100644
--- a/litellm/proxy/auth/litellm_license.py
+++ b/litellm/proxy/auth/litellm_license.py
@@ -1,6 +1,14 @@
 # What is this?
 ## If litellm license in env, checks if it's valid
+import base64
+import json
 import os
+from datetime import datetime
+
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+
+from litellm._logging import verbose_proxy_logger
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
 
 
@@ -15,6 +23,20 @@ class LicenseCheck:
     def __init__(self) -> None:
         self.license_str = os.getenv("LITELLM_LICENSE", None)
         self.http_handler = HTTPHandler()
+        self.public_key = None
+        self.read_public_key()
+
+    def read_public_key(self):
+        # current dir
+        current_dir = os.path.dirname(os.path.realpath(__file__))
+
+        # check if public_key.pem exists
+        _path_to_public_key = os.path.join(current_dir, "public_key.pem")
+        if os.path.exists(_path_to_public_key):
+            with open(_path_to_public_key, "rb") as key_file:
+                self.public_key = serialization.load_pem_public_key(key_file.read())
+        else:
+            self.public_key = None
 
     def _verify(self, license_str: str) -> bool:
         url = "{}/verify_license/{}".format(self.base_url, license_str)
@@ -35,11 +57,54 @@ class LicenseCheck:
             return False
 
     def is_premium(self) -> bool:
+        """
+        1. verify_license_without_api_request: checks if license was generate using private / public key pair
+        2. _verify: checks if license is valid calling litellm API. This is the old way we were generating/validating license
+        """
         try:
             if self.license_str is None:
                 return False
+            elif self.verify_license_without_api_request(
+                public_key=self.public_key, license_key=self.license_str
+            ):
+                return True
             elif self._verify(license_str=self.license_str):
                 return True
             return False
         except Exception as e:
             return False
+
+    def verify_license_without_api_request(self, public_key, license_key):
+        try:
+            # Decode the license key
+            decoded = base64.b64decode(license_key)
+            message, signature = decoded.split(b".", 1)
+
+            # Verify the signature
+            public_key.verify(
+                signature,
+                message,
+                padding.PSS(
+                    mgf=padding.MGF1(hashes.SHA256()),
+                    salt_length=padding.PSS.MAX_LENGTH,
+                ),
+                hashes.SHA256(),
+            )
+
+            # Decode and parse the data
+            license_data = json.loads(message.decode())
+
+            # debug information provided in license data
+            verbose_proxy_logger.debug("License data: %s", license_data)
+
+            # Check expiration date
+            expiration_date = datetime.strptime(
+                license_data["expiration_date"], "%Y-%m-%d"
+            )
+            if expiration_date < datetime.now():
+                return False, "License has expired"
+
+            return True
+
+        except Exception as e:
+            return False
diff --git a/litellm/proxy/auth/public_key.pem b/litellm/proxy/auth/public_key.pem
new file mode 100644
index 000000000..12a69dde2
--- /dev/null
+++ b/litellm/proxy/auth/public_key.pem
@@ -0,0 +1,9 @@
+-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAmfBuNiNzDkNWyce23koQ
+w0vq3bSVHkq7fd9Sw/U1q7FwRwL221daLTyGWssd8xAoQSFXAJKoBwzJQ9wd+o44
+lfL54E3a61nfjZuF+D9ntpXZFfEAxLVtIahDeQjUz4b/EpgciWIJyUfjCJrQo6LY
+eyAZPTGSO8V3zHyaU+CFywq5XCuCnfZqCZeCw051St59A2v8W32mXSCJ+A+x0hYP
+yXJyRRFcefSFG5IBuRHr4Y24Vx7NUIAoco5cnxJho9g2z3J/Hb0GKW+oBNvRVumk
+nuA2Ljmjh4yI0OoTIW8ZWxemvCCJHSjdfKlMyb+QI4fmeiIUZzP5Au+F561Styqq
+YQIDAQAB
+-----END PUBLIC KEY-----
diff --git a/requirements.txt b/requirements.txt
index fbf2bfc1d..8c5e4ab3b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -31,6 +31,7 @@ azure-identity==1.16.1 # for azure content safety
 opentelemetry-api==1.25.0
 opentelemetry-sdk==1.25.0
 opentelemetry-exporter-otlp==1.25.0
+cryptography==42.0.7
 
 ### LITELLM PACKAGE DEPENDENCIES
 python-dotenv==1.0.0 # for env