Merge pull request #652 from guardrails-ai/feat/cli-jwt-auth

new jwt strat

docs-graveyard/cli.md

@@ -15,7 +15,7 @@ In addition to providing a command line interface for validation, the guardrails
 ### Configuration
 In order to access any Validators from the Hub that require authentication, you will need to set up your environment through the `guardrails configure` command.  Before running `guardrails configure`, go to the [Validator Hub]() to generate your access tokens.  [Add more detail on this process].
 
-Once you have your tokens, run `guardrails configure` and enter the client id and client secret you retrieved above.
+Once you have your tokens, run `guardrails configure` and enter the auth token you retrieved above.
 
 Also, if you want to opt out of anonymous metrics collection, you can do this now by specifying `--no-metrics=true`.
 

guardrails/classes/credentials.py

@@ -10,8 +10,7 @@
 @dataclass
 class Credentials(Serializeable):
     id: Optional[str] = None
-    client_id: Optional[str] = None
-    client_secret: Optional[str] = None
+    token: Optional[str] = None
     no_metrics: Optional[bool] = False
 
     @staticmethod

guardrails/cli/configure.py

@@ -11,16 +11,13 @@
 from guardrails.cli.server.hub_client import AuthenticationError, get_auth
 
 
-def save_configuration_file(
-    client_id: str, client_secret: str, no_metrics: bool
-) -> None:
+def save_configuration_file(token: str, no_metrics: bool) -> None:
     home = expanduser("~")
     guardrails_rc = os.path.join(home, ".guardrailsrc")
     with open(guardrails_rc, "w") as rc_file:
         lines = [
             f"id={str(uuid.uuid4())}{os.linesep}",
-            f"client_id={client_id}{os.linesep}",
-            f"client_secret={client_secret}{os.linesep}",
+            f"token={token}{os.linesep}",
             f"no_metrics={str(no_metrics).lower()}",
         ]
         rc_file.writelines(lines)
@@ -29,11 +26,8 @@ def save_configuration_file(
 
 @guardrails.command()
 def configure(
-    client_id: Optional[str] = typer.Option(
-        help="Your Guardrails Hub client ID.", default=""
-    ),
-    client_secret: Optional[str] = typer.Option(
-        help="Your Guardrails Hub client secret.", hide_input=True, default=""
+    token: Optional[str] = typer.Option(
+        help="Your Guardrails Hub auth token.", hide_input=True, default=""
     ),
     no_metrics: Optional[str] = typer.Option(
         help="Opt out of anonymous metrics collection.", default=False
@@ -43,15 +37,13 @@ def configure(
     try:
         notice_message = """
 
-    You can find your tokens at https://hub.guardrailsai.com/tokens
+    You can find your token at https://hub.guardrailsai.com/tokens
     """
         logger.log(level=LEVELS.get("NOTICE"), msg=notice_message)  # type: ignore
-        if not client_id:
-            client_id = typer.prompt("Client ID")
-        if not client_secret:
-            client_secret = typer.prompt("Client secret", hide_input=True)
+        if not token:
+            token = typer.prompt("Token", hide_input=True)
         logger.info("Configuring...")
-        save_configuration_file(client_id, client_secret, no_metrics)  # type: ignore
+        save_configuration_file(token, no_metrics)  # type: ignore
 
         logger.info("Validating credentials...")
         get_auth()
@@ -72,7 +64,7 @@ def configure(
         logger.error(auth_error)
         logger.error(
             """
-            Check that your Client ID and Client secret are correct and try again.
+            Check that your token is correct and try again.
 
             If you don't have your token credentials you can find them here:
 

guardrails/cli/server/auth.py

@@ -1,25 +1,26 @@
-import http.client
-import json
+# import http.client
+# import json
 
-from guardrails.classes.credentials import Credentials
+# from guardrails.classes.credentials import Credentials
 
 
-def get_auth_token(creds: Credentials) -> str:
-    if creds.client_id and creds.client_secret:
-        audience = "https://validator-hub-service.guardrailsai.com"
-        conn = http.client.HTTPSConnection("guardrailsai.us.auth0.com")
-        payload = json.dumps(
-            {
-                "client_id": creds.client_id,
-                "client_secret": creds.client_secret,
-                "audience": audience,
-                "grant_type": "client_credentials",
-            }
-        )
-        headers = {"content-type": "application/json"}
-        conn.request("POST", "/oauth/token", payload, headers)
+# unused - for now
+# def get_auth_token(creds: Credentials) -> str:
+#     if creds.client_id and creds.client_secret:
+#         audience = "https://validator-hub-service.guardrailsai.com"
+#         conn = http.client.HTTPSConnection("guardrailsai.us.auth0.com")
+#         payload = json.dumps(
+#             {
+#                 "client_id": creds.client_id,
+#                 "client_secret": creds.client_secret,
+#                 "audience": audience,
+#                 "grant_type": "client_credentials",
+#             }
+#         )
+#         headers = {"content-type": "application/json"}
+#         conn.request("POST", "/oauth/token", payload, headers)
 
-        res = conn.getresponse()
-        data = json.loads(res.read().decode("utf-8"))
-        return data.get("access_token", "")
-    return ""
+#         res = conn.getresponse()
+#         data = json.loads(res.read().decode("utf-8"))
+#         return data.get("access_token", "")
+#     return ""

guardrails/cli/server/hub_client.py

@@ -3,12 +3,20 @@
 from typing import Any, Dict, Optional
 
 import requests
+from jwt import JWT
+from jwt.exceptions import JWTDecodeError
 
 from guardrails.classes.credentials import Credentials
 from guardrails.cli.logger import logger
-from guardrails.cli.server.auth import get_auth_token
 from guardrails.cli.server.module_manifest import ModuleManifest
 
+TOKEN_EXPIRED_MESSAGE = (
+    "Your token has expired. Please run `guardrails configure` to update your token."
+)
+TOKEN_INVALID_MESSAGE = (
+    "Your token is invalid. Please run `guardrails configure` to update your token."
+)
+
 validator_hub_service = "https://so4sg4q4pb.execute-api.us-east-1.amazonaws.com"
 validator_manifest_endpoint = Template(
     "validator-manifests/${namespace}/${validator_name}"
@@ -52,7 +60,7 @@ def fetch(url: str, token: Optional[str], anonymousUserId: Optional[str]):
 
 
 def fetch_module_manifest(
-    module_name: str, token: str, anonymousUserId: Optional[str] = None
+    module_name: str, token: Optional[str], anonymousUserId: Optional[str] = None
 ) -> Dict[str, Any]:
     namespace, validator_name = module_name.split("/", 1)
     manifest_path = validator_manifest_endpoint.safe_substitute(
@@ -62,9 +70,25 @@ def fetch_module_manifest(
     return fetch(manifest_url, token, anonymousUserId)
 
 
+def get_jwt_token(creds: Credentials) -> Optional[str]:
+    token = creds.token
+
+    # check for jwt expiration
+    if token:
+        try:
+            JWT().decode(token, do_verify=False)
+        except JWTDecodeError as e:
+            # if the error message includes "Expired", then the token is expired
+            if "Expired" in str(e):
+                raise Exception(TOKEN_EXPIRED_MESSAGE)
+            else:
+                raise Exception(TOKEN_INVALID_MESSAGE)
+    return token
+
+
 def fetch_module(module_name: str) -> ModuleManifest:
     creds = Credentials.from_rc_file(logger)
-    token = get_auth_token(creds)
+    token = get_jwt_token(creds)
 
     module_manifest_json = fetch_module_manifest(module_name, token, creds.id)
     return ModuleManifest.from_dict(module_manifest_json)
@@ -90,7 +114,7 @@ def get_validator_manifest(module_name: str):
 def get_auth():
     try:
         creds = Credentials.from_rc_file(logger)
-        token = get_auth_token(creds)
+        token = get_jwt_token(creds)
         auth_url = f"{validator_hub_service}/auth"
         response = fetch(auth_url, token, creds.id)
         if not response:
@@ -106,7 +130,7 @@ def get_auth():
 def post_validator_submit(package_name: str, content: str):
     try:
         creds = Credentials.from_rc_file(logger)
-        token = get_auth_token(creds)
+        token = get_jwt_token(creds)
         submission_url = f"{validator_hub_service}/validator/submit"
 
         headers = {