From 8812d0c308dced89a11d6ae67813e639496273af Mon Sep 17 00:00:00 2001
From: Dennis Chen <41879777+chennisden@users.noreply.github.com>
Date: Wed, 24 Jul 2024 13:48:43 -0700
Subject: [PATCH] Hash login tokens (#172)

* Delete unused JWT functions

* Hash api tokens for login

The general approach was to create general _get_hashed_item (and similar
hashed functions for the rest of the item functions) in base.py, and
then substitute the regular functions for the hashed functions in
the api_key functions in users.py.
---
 Makefile                |  2 +-
 store/app/crud/base.py  | 22 ++++++++++++++++++++++
 store/app/crud/users.py | 10 ++++++----
 store/app/model.py      | 16 ----------------
 4 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/Makefile b/Makefile
index 18ab2876..ecd72a44 100644
--- a/Makefile
+++ b/Makefile
@@ -43,7 +43,7 @@ start-docker-dynamodb:
 
 format-backend:
 	@black store tests
-	@ruff format store tests
+	@ruff check --fix store tests
 .PHONY: format
 
 format-frontend:
diff --git a/store/app/crud/base.py b/store/app/crud/base.py
index 9b3c3ff4..b612af01 100644
--- a/store/app/crud/base.py
+++ b/store/app/crud/base.py
@@ -10,6 +10,7 @@
 from types_aiobotocore_dynamodb.service_resource import DynamoDBServiceResource
 from types_aiobotocore_s3.service_resource import S3ServiceResource
 
+from store.app.crypto import hash_token
 from store.app.model import RobolistBaseModel
 
 TABLE_NAME = "Robolist"
@@ -76,6 +77,15 @@ async def _add_item(self, item: RobolistBaseModel) -> None:
         item_data["type"] = item.__class__.__name__
         await table.put_item(Item=item_data)
 
+    async def _add_hashed_item(self, item: RobolistBaseModel) -> None:
+        table = await self.db.Table(TABLE_NAME)
+        item_data = item.model_dump()
+        if "type" in item_data:
+            raise ValueError("Cannot add item with 'type' attribute")
+        item_data["type"] = item.__class__.__name__
+        item_data["id"] = hash_token(item_data["id"])
+        await table.put_item(Item=item_data)
+
     async def _delete_item(self, item: RobolistBaseModel | str) -> None:
         table = await self.db.Table(TABLE_NAME)
         if isinstance(item, str):
@@ -83,6 +93,13 @@ async def _delete_item(self, item: RobolistBaseModel | str) -> None:
         else:
             await table.delete_item(Key={"id": item.id})
 
+    async def _delete_hashed_item(self, item: RobolistBaseModel | str) -> None:
+        table = await self.db.Table(TABLE_NAME)
+        if isinstance(item, str):
+            await table.delete_item(Key={"id": hash_token(item)})
+        else:
+            await table.delete_item(Key={"id": hash_token(item.id)})
+
     async def _list_items(
         self,
         item_class: type[T],
@@ -187,6 +204,11 @@ async def _item_exists(self, item_id: str) -> bool:
         item_dict = await table.get_item(Key={"id": item_id})
         return "Item" in item_dict
 
+    async def _hashed_item_exists(self, item_id: str) -> bool:
+        table = await self.db.Table(TABLE_NAME)
+        item_dict = await table.get_item(Key={"id": hash_token(item_id)})
+        return "Item" in item_dict
+
     async def _get_item_batch(
         self,
         item_ids: list[str],
diff --git a/store/app/crud/users.py b/store/app/crud/users.py
index 8379ae78..a5af2d4d 100644
--- a/store/app/crud/users.py
+++ b/store/app/crud/users.py
@@ -5,6 +5,7 @@
 from datetime import datetime
 
 from store.app.crud.base import BaseCrud, GlobalSecondaryIndex
+from store.app.crypto import hash_token
 from store.app.model import APIKey, OAuthKey, User
 from store.settings import settings
 from store.utils import LRUCache
@@ -92,15 +93,16 @@ async def get_user_count(self) -> int:
         return await self._count_items(User)
 
     async def get_api_key(self, id: str) -> APIKey:
-        return await self._get_item(id, APIKey, throw_if_missing=True)
+        hashed_id = hash_token(id)
+        return await self._get_item(hashed_id, APIKey, throw_if_missing=True)
 
     async def add_api_key(self, id: str) -> APIKey:
         token = APIKey.create(id=id)
-        await self._add_item(token)
+        await self._add_hashed_item(token)
         return token
 
     async def delete_api_key(self, token: APIKey | str) -> None:
-        await self._delete_item(token)
+        await self._delete_hashed_item(token)
 
     async def api_key_is_valid(self, token: str) -> bool:
         """Validates a token against the database, with caching.
@@ -119,7 +121,7 @@ async def api_key_is_valid(self, token: str) -> bool:
             last_time, is_valid = LAST_API_KEY_VALIDATION[token]
             if (cur_time - last_time).seconds < settings.crypto.cache_token_db_result_seconds:
                 return is_valid
-        is_valid = await self._item_exists(token)
+        is_valid = await self._hashed_item_exists(token)
         LAST_API_KEY_VALIDATION[token] = (cur_time, is_valid)
         return is_valid
 
diff --git a/store/app/model.py b/store/app/model.py
index 5f352251..0492c850 100644
--- a/store/app/model.py
+++ b/store/app/model.py
@@ -7,11 +7,9 @@
 
 from typing import Self
 
-import jwt
 from pydantic import BaseModel
 
 from store.app.crypto import new_uuid
-from store.settings import settings
 
 
 class RobolistBaseModel(BaseModel):
@@ -72,20 +70,6 @@ def create(cls, id: str) -> Self:
             user_id=id,
         )
 
-    def to_jwt(self) -> str:
-        return jwt.encode(
-            payload={"token": self.id, "user_id": self.user_id},
-            key=settings.crypto.jwt_secret,
-        )
-
-    @classmethod
-    def from_jwt(cls, jwt_token: str) -> Self:
-        data = jwt.decode(
-            jwt=jwt_token,
-            key=settings.crypto.jwt_secret,
-        )
-        return cls(id=data["token"], user_id=data["user_id"])
-
 
 class Bom(BaseModel):
     part_id: str