endpoint to check if an api key is valid or not (#343)

* endpoint to check if an api key is valid or not * tests pass
kscalelabs · Aug 30, 2024 · 3810329 · 3810329
1 parent 221ccfc
commit 3810329
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 8 deletions.
diff --git a/frontend/src/components/pages/APIKeys.tsx b/frontend/src/components/pages/APIKeys.tsx
@@ -114,6 +114,7 @@ const APIKeys = () => {
   const auth = useAuthentication();
   const [apiKeys, setApiKeys] = useState<KeysResponse | null>(null);
   const [readonly, setReadonly] = useState(true);
+  const [creatingKey, setCreatingKey] = useState(false);
 
   useEffect(() => {
     const fetchUser = async () => {
@@ -133,6 +134,7 @@ const APIKeys = () => {
   }, [auth]);
 
   const createKey = async () => {
+    setCreatingKey(true);
     const { data, error } = await auth.client.POST("/keys/new", {
       body: {
         readonly,
@@ -144,6 +146,7 @@ const APIKeys = () => {
     } else {
       setApiKeys(apiKeys ? [...apiKeys, data.key] : null);
     }
+    setCreatingKey(false);
   };
 
   return (
@@ -187,7 +190,11 @@ const APIKeys = () => {
                     )}
                   ></div>
                 </label>
-                <Button onClick={createKey} variant="primary">
+                <Button
+                  onClick={createKey}
+                  variant="primary"
+                  disabled={creatingKey}
+                >
                   Create {readonly ? "Read-only " : "Read-write "}Key
                 </Button>
               </div>

diff --git a/frontend/src/gen/api.ts b/frontend/src/gen/api.ts
@@ -588,6 +588,23 @@ export interface paths {
         patch?: never;
         trace?: never;
     };
+    "/users/validate-api-key": {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        /** Validate Api Key Endpoint */
+        get: operations["validate_api_key_endpoint_users_validate_api_key_get"];
+        put?: never;
+        post?: never;
+        delete?: never;
+        options?: never;
+        head?: never;
+        patch?: never;
+        trace?: never;
+    };
     "/users/github/client-id": {
         parameters: {
             query?: never;
@@ -2205,6 +2222,26 @@ export interface operations {
             };
         };
     };
+    validate_api_key_endpoint_users_validate_api_key_get: {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        requestBody?: never;
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": boolean;
+                };
+            };
+        };
+    };
     github_client_id_endpoint_users_github_client_id_get: {
         parameters: {
             query?: never;

diff --git a/store/app/crud/base.py b/store/app/crud/base.py
@@ -205,13 +205,25 @@ async def _list_me(
         sorted_items = sorted(response, key=sort_key, reverse=True)
         return sorted_items[(page - 1) * ITEMS_PER_PAGE : page * ITEMS_PER_PAGE], page * ITEMS_PER_PAGE < len(response)
 
-    async def _count_items(self, item_class: type[T]) -> int:
+    async def _count_items(
+        self,
+        item_class: type[T],
+        expression_attribute_values: dict[str, Any] | None = None,
+    ) -> int:
         table = await self.db.Table(TABLE_NAME)
-        item_dict = await table.scan(
-            IndexName="type_index",
-            Select="COUNT",
-            FilterExpression=Key("type").eq(item_class.__name__),
-        )
+        if expression_attribute_values is None:
+            item_dict = await table.scan(
+                IndexName="type_index",
+                Select="COUNT",
+                FilterExpression=Key("type").eq(item_class.__name__),
+            )
+        else:
+            item_dict = await table.scan(
+                IndexName="type_index",
+                Select="COUNT",
+                FilterExpression=Key("type").eq(item_class.__name__),
+                ExpressionAttributeValues=expression_attribute_values,
+            )
         return item_dict["Count"]
 
     def _validate_item(self, data: dict[str, Any], item_class: type[T]) -> T:

diff --git a/store/app/crud/users.py b/store/app/crud/users.py
@@ -148,6 +148,9 @@ async def add_api_key(
         source: APIKeySource,
         permissions: APIKeyPermissionSet,
     ) -> APIKey:
+        api_key_count = await self._count_items(APIKey, expression_attribute_values={":user_id": user_id})
+        if api_key_count >= 10:
+            raise ValueError("User has reached the maximum number of API keys (10)")
         api_key = APIKey.create(user_id=user_id, source=source, permissions=permissions)
         await self._add_item(api_key)
         return api_key

diff --git a/store/app/model.py b/store/app/model.py
@@ -430,13 +430,15 @@ def get_artifact_name(
 def get_artifact_url(
     *,
     artifact: Artifact | None = None,
+    artifact_id: str | None = None,
     artifact_type: ArtifactType | None = None,
     listing_id: str | None = None,
     name: str | None = None,
     size: ArtifactSize = "large",
 ) -> str:
     artifact_name = get_artifact_name(
         artifact=artifact,
+        artifact_id=artifact_id,
         listing_id=listing_id,
         name=name,
         artifact_type=artifact_type,

diff --git a/store/app/routers/artifacts.py b/store/app/routers/artifacts.py
@@ -59,7 +59,8 @@ async def artifact_url(
     # TODO: Use CloudFront API to return a signed CloudFront URL.
     return RedirectResponse(
         url=get_artifact_url(
-            artifact_type=artifact_type,
+            artifact_type=artifact.artifact_type,
+            artifact_id=artifact.id,
             listing_id=listing_id,
             name=s3_filename,
             size=size,

diff --git a/store/app/routers/users.py b/store/app/routers/users.py
@@ -354,5 +354,17 @@ async def update_profile(
         )
 
 
+@users_router.get("/validate-api-key")
+async def validate_api_key_endpoint(
+    crud: Annotated[Crud, Depends(Crud.get)],
+    api_key_id: Annotated[str, Depends(get_request_api_key_id)],
+) -> bool:
+    try:
+        await crud.get_api_key(api_key_id)
+        return True
+    except ItemNotFoundError:
+        return False
+
+
 users_router.include_router(github_auth_router, prefix="/github")
 users_router.include_router(google_auth_router, prefix="/google")