(feat) oauth msft + google

litellm/proxy/proxy_server.py

@@ -2885,7 +2885,7 @@ async def user_auth(request: Request):
     return "Email sent!"
 
 
-@app.get("/google-login/key/generate", tags=["experimental"])
+@app.get("/sso/key/generate", tags=["experimental"])
 async def google_login(request: Request):
     """
     Create Proxy API Keys using Google Workspace SSO. Requires setting GOOGLE_REDIRECT_URI in .env
@@ -2894,126 +2894,235 @@ async def google_login(request: Request):
     Example:
 
     """
-    GOOGLE_REDIRECT_URI = os.getenv("PROXY_BASE_URL")
+    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
-    if GOOGLE_REDIRECT_URI is None:
+    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
+    redirect_url = os.getenv("PROXY_BASE_URL", None)
+    if redirect_url is None:
         raise ProxyException(
             message="PROXY_BASE_URL not set. Set it in .env file",
             type="auth_error",
             param="PROXY_BASE_URL",
             code=status.HTTP_500_INTERNAL_SERVER_ERROR,
         )
-    if GOOGLE_REDIRECT_URI.endswith("/"):
-        GOOGLE_REDIRECT_URI += "google-callback"
-    else:
-        GOOGLE_REDIRECT_URI += "/google-callback"
 
-    GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
+    if redirect_url.endswith("/"):
-    if GOOGLE_CLIENT_ID is None:
+        redirect_url += "sso/callback"
-        GOOGLE_CLIENT_ID = (
+    else:
-            "246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com"
+        redirect_url += "/sso/callback"
+    # Google SSO Auth
+    if google_client_id is not None:
+        from fastapi_sso.sso.google import GoogleSSO
+
+        google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
+        if google_client_secret is None:
+            raise ProxyException(
+                message="GOOGLE_CLIENT_SECRET not set. Set it in .env file",
+                type="auth_error",
+                param="GOOGLE_CLIENT_SECRET",
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
+
+        google_sso = GoogleSSO(
+            client_id=google_client_id,
+            client_secret=google_client_secret,
+            redirect_uri=redirect_url,
         )
 
-    verbose_proxy_logger.info(
+        verbose_proxy_logger.info(
-        f"In /google-login/key/generate, \nGOOGLE_REDIRECT_URI: {GOOGLE_REDIRECT_URI}\nGOOGLE_CLIENT_ID: {GOOGLE_CLIENT_ID}"
+            f"In /google-login/key/generate, \nGOOGLE_REDIRECT_URI: {redirect_url}\nGOOGLE_CLIENT_ID: {google_client_id}"
-    )
+        )
-    google_auth_url = f"https://accounts.google.com/o/oauth2/auth?client_id={GOOGLE_CLIENT_ID}&redirect_uri={GOOGLE_REDIRECT_URI}&response_type=code&scope=openid%20profile%20email"
+
-    return RedirectResponse(url=google_auth_url)
+        with google_sso:
+            return await google_sso.get_login_redirect()
+
+    # Microsoft SSO Auth
+    elif microsoft_client_id is not None:
+        from fastapi_sso.sso.microsoft import MicrosoftSSO
+
+        microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
+        microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
+        if microsoft_client_secret is None:
+            raise ProxyException(
+                message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file",
+                type="auth_error",
+                param="MICROSOFT_CLIENT_SECRET",
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
+
+        microsoft_sso = MicrosoftSSO(
+            client_id=microsoft_client_id,
+            client_secret=microsoft_client_secret,
+            tenant=microsoft_tenant,
+            redirect_uri=redirect_url,
+            allow_insecure_http=True,
+        )
+        with microsoft_sso:
+            return await microsoft_sso.get_login_redirect()
 
 
-@app.get("/google-callback", tags=["experimental"], response_model=GenerateKeyResponse)
+@app.get("/sso/callback", tags=["experimental"])
-async def google_callback(code: str, request: Request):
+async def auth_callback(request: Request):
-    import httpx
+    """Verify login"""
+    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
+    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
 
-    GOOGLE_REDIRECT_URI = os.getenv("PROXY_BASE_URL")
+    redirect_url = os.getenv("PROXY_BASE_URL", None)
-    if GOOGLE_REDIRECT_URI is None:
+    if redirect_url is None:
         raise ProxyException(
             message="PROXY_BASE_URL not set. Set it in .env file",
             type="auth_error",
             param="PROXY_BASE_URL",
             code=status.HTTP_500_INTERNAL_SERVER_ERROR,
         )
-    # Add "/google-callback"" to your callback URL
+    if redirect_url.endswith("/"):
-    if GOOGLE_REDIRECT_URI.endswith("/"):
+        redirect_url += "sso/callback"
-        GOOGLE_REDIRECT_URI += "google-callback"
     else:
-        GOOGLE_REDIRECT_URI += "/google-callback"
+        redirect_url += "/sso/callback"
 
-    GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
+    if google_client_id is not None:
-    if GOOGLE_CLIENT_ID is None:
+        from fastapi_sso.sso.google import GoogleSSO
-        GOOGLE_CLIENT_ID = (
+
-            "246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com"
+        google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
+        if google_client_secret is None:
+            raise ProxyException(
+                message="GOOGLE_CLIENT_SECRET not set. Set it in .env file",
+                type="auth_error",
+                param="GOOGLE_CLIENT_SECRET",
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
+        google_sso = GoogleSSO(
+            client_id=google_client_id,
+            redirect_uri=redirect_url,
+            client_secret=google_client_secret,
         )
+        result = await google_sso.verify_and_process(request)
+        return result
 
-    GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
+    elif microsoft_client_id is not None:
-    if GOOGLE_CLIENT_SECRET is None:
+        from fastapi_sso.sso.microsoft import MicrosoftSSO
-        GOOGLE_CLIENT_SECRET = "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"
 
-    verbose_proxy_logger.info(
+        microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
-        f"/google-callback\n GOOGLE_REDIRECT_URI: {GOOGLE_REDIRECT_URI}\n GOOGLE_CLIENT_ID: {GOOGLE_CLIENT_ID}"
+        microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
-    )
+        if microsoft_client_secret is None:
-    # Exchange code for access token
+            raise ProxyException(
-    async with httpx.AsyncClient() as client:
+                message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file",
-        token_url = f"https://oauth2.googleapis.com/token"
+                type="auth_error",
-        data = {
+                param="MICROSOFT_CLIENT_SECRET",
-            "code": code,
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            "client_id": GOOGLE_CLIENT_ID,
+            )
-            "client_secret": GOOGLE_CLIENT_SECRET,
+        if microsoft_tenant is None:
-            "redirect_uri": GOOGLE_REDIRECT_URI,
+            raise ProxyException(
-            "grant_type": "authorization_code",
+                message="MICROSOFT_TENANT not set. Set it in .env file",
-        }
+                type="auth_error",
-        response = await client.post(token_url, data=data)
+                param="MICROSOFT_TENANT",
-
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-    # Process the response, extract user info, etc.
-    if response.status_code == 200:
-        access_token = response.json()["access_token"]
-
-        # Fetch user info using the access token
-        async with httpx.AsyncClient() as client:
-            user_info_url = "https://www.googleapis.com/oauth2/v1/userinfo"
-            headers = {"Authorization": f"Bearer {access_token}"}
-            user_info_response = await client.get(user_info_url, headers=headers)
-
-        # Process user info response
-        if user_info_response.status_code == 200:
-            user_info = user_info_response.json()
-            user_email = user_info.get("email")
-            user_name = user_info.get("name")
-
-            # we can use user_email on litellm proxy now
-
-            # TODO: Handle user info as needed, for example, store it in a database, authenticate the user, etc.
-            response = await generate_key_helper_fn(
-                **{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_email, "team_id": "litellm-dashboard"}  # type: ignore
             )
 
-            key = response["token"]  # type: ignore
+        microsoft_sso = MicrosoftSSO(
-            user_id = response["user_id"]  # type: ignore
+            client_id=microsoft_client_id,
-            litellm_dashboard_ui = "https://litellm-dashboard.vercel.app/"
+            client_secret=microsoft_client_secret,
+            tenant=microsoft_tenant,
+            redirect_uri=redirect_url,
+            allow_insecure_http=True,
+        )
+        result = await microsoft_sso.verify_and_process(request)
+        return result
 
-            # if user set LITELLM_UI_LINK in .env, use that
-            litellm_ui_link_in_env = os.getenv("LITELLM_UI_LINK", None)
-            if litellm_ui_link_in_env is not None:
-                litellm_dashboard_ui = litellm_ui_link_in_env
 
-            litellm_dashboard_ui += (
+# @app.get("/google-callback", tags=["experimental"], response_model=GenerateKeyResponse)
-                "?userID="
+# async def google_callback(code: str, request: Request):
-                + user_id
+#     import httpx
-                + "&accessToken="
-                + key
-                + "&proxyBaseUrl="
-                + os.getenv("PROXY_BASE_URL")
-            )
-            return RedirectResponse(url=litellm_dashboard_ui)
 
-        else:
+#     GOOGLE_REDIRECT_URI = os.getenv("PROXY_BASE_URL")
-            # Handle user info retrieval error
+#     if GOOGLE_REDIRECT_URI is None:
-            raise HTTPException(
+#         raise ProxyException(
-                status_code=user_info_response.status_code,
+#             message="PROXY_BASE_URL not set. Set it in .env file",
-                detail=user_info_response.text,
+#             type="auth_error",
-            )
+#             param="PROXY_BASE_URL",
-    else:
+#             code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-        # Handle the error from the token exchange
+#         )
-        raise HTTPException(status_code=response.status_code, detail=response.text)
+#     # Add "/google-callback"" to your callback URL
+#     if GOOGLE_REDIRECT_URI.endswith("/"):
+#         GOOGLE_REDIRECT_URI += "google-callback"
+#     else:
+#         GOOGLE_REDIRECT_URI += "/google-callback"
+
+#     GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
+#     if GOOGLE_CLIENT_ID is None:
+#         GOOGLE_CLIENT_ID = (
+#             "246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com"
+#         )
+
+#     GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
+#     if GOOGLE_CLIENT_SECRET is None:
+#         GOOGLE_CLIENT_SECRET = "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"
+
+#     verbose_proxy_logger.info(
+#         f"/google-callback\n GOOGLE_REDIRECT_URI: {GOOGLE_REDIRECT_URI}\n GOOGLE_CLIENT_ID: {GOOGLE_CLIENT_ID}"
+#     )
+#     # Exchange code for access token
+#     async with httpx.AsyncClient() as client:
+#         token_url = f"https://oauth2.googleapis.com/token"
+#         data = {
+#             "code": code,
+#             "client_id": GOOGLE_CLIENT_ID,
+#             "client_secret": GOOGLE_CLIENT_SECRET,
+#             "redirect_uri": GOOGLE_REDIRECT_URI,
+#             "grant_type": "authorization_code",
+#         }
+#         response = await client.post(token_url, data=data)
+
+#     # Process the response, extract user info, etc.
+#     if response.status_code == 200:
+#         access_token = response.json()["access_token"]
+
+#         # Fetch user info using the access token
+#         async with httpx.AsyncClient() as client:
+#             user_info_url = "https://www.googleapis.com/oauth2/v1/userinfo"
+#             headers = {"Authorization": f"Bearer {access_token}"}
+#             user_info_response = await client.get(user_info_url, headers=headers)
+
+#         # Process user info response
+#         if user_info_response.status_code == 200:
+#             user_info = user_info_response.json()
+#             user_email = user_info.get("email")
+#             user_name = user_info.get("name")
+
+#             # we can use user_email on litellm proxy now
+
+#             # TODO: Handle user info as needed, for example, store it in a database, authenticate the user, etc.
+#             response = await generate_key_helper_fn(
+#                 **{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_email, "team_id": "litellm-dashboard"}  # type: ignore
+#             )
+
+#             key = response["token"]  # type: ignore
+#             user_id = response["user_id"]  # type: ignore
+#             litellm_dashboard_ui = "https://litellm-dashboard.vercel.app/"
+
+#             # if user set LITELLM_UI_LINK in .env, use that
+#             litellm_ui_link_in_env = os.getenv("LITELLM_UI_LINK", None)
+#             if litellm_ui_link_in_env is not None:
+#                 litellm_dashboard_ui = litellm_ui_link_in_env
+
+#             litellm_dashboard_ui += (
+#                 "?userID="
+#                 + user_id
+#                 + "&accessToken="
+#                 + key
+#                 + "&proxyBaseUrl="
+#                 + os.getenv("PROXY_BASE_URL")
+#             )
+#             return RedirectResponse(url=litellm_dashboard_ui)
+
+#         else:
+#             # Handle user info retrieval error
+#             raise HTTPException(
+#                 status_code=user_info_response.status_code,
+#                 detail=user_info_response.text,
+#             )
+#     else:
+#         # Handle the error from the token exchange
+#         raise HTTPException(status_code=response.status_code, detail=response.text)
 
 
 @router.get(