From 8195a4eaccff6b481d6e37d7d2e2d08533f28605 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 9 Jul 2024 16:25:23 -0700
Subject: [PATCH] check if premium user for sso / allowed ip

---
 ui/litellm-dashboard/src/app/page.tsx          |  1 +
 ui/litellm-dashboard/src/components/admins.tsx | 14 ++++++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 1aaf8cbc3d..26bdb7af24 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -221,6 +221,7 @@ const CreateKeyPage = () => {
               searchParams={searchParams}
               accessToken={accessToken}
               showSSOBanner={showSSOBanner}
+              premiumUser={premiumUser}
             />
           ) : page == "api_ref" ? (
             <APIRef 
diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx
index ad25011653..94211868f4 100644
--- a/ui/litellm-dashboard/src/components/admins.tsx
+++ b/ui/litellm-dashboard/src/components/admins.tsx
@@ -40,6 +40,7 @@ interface AdminPanelProps {
   accessToken: string | null;
   setTeams: React.Dispatch<React.SetStateAction<Object[] | null>>;
   showSSOBanner: boolean;
+  premiumUser: boolean;
 }
 
 import {
@@ -59,6 +60,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
   searchParams,
   accessToken,
   showSSOBanner,
+  premiumUser,
 }) => {
   const [form] = Form.useForm();
   const [memberForm] = Form.useForm();
@@ -105,6 +107,12 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
 
   const handleShowAllowedIPs = async () => {
     try {
+      if (premiumUser !== true) {
+        message.error(
+          "This feature is only available for premium users. Please upgrade your account."
+        )
+        return
+      }
       if (accessToken) {
         const data = await getAllowedIPs(accessToken);
         setAllowedIPs(data && data.length > 0 ? data : [all_ip_address_allowed]);
@@ -116,7 +124,9 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
       message.error(`Failed to fetch allowed IPs ${error}`);
       setAllowedIPs([all_ip_address_allowed]);
     } finally {
-      setIsAllowedIPModalVisible(true);
+      if (premiumUser === true) {
+        setIsAllowedIPModalVisible(true);
+      }
     }
   };
   
@@ -605,7 +615,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
         <Title level={4}> ✨ Security Settings</Title>
         <div style={{ display: 'flex', flexDirection: 'column', gap: '1rem', marginTop: '1rem' }}>
     <div>
-      <Button onClick={() => setIsAddSSOModalVisible(true)}>Add SSO</Button>
+      <Button onClick={() => premiumUser === true ? setIsAddSSOModalVisible(true) : message.error("Only premium users can add SSO")}>Add SSO</Button>
     </div>
     <div>
       <Button onClick={handleShowAllowedIPs}>Allowed IPs</Button>