convert from using tokens to using api keys

store/app/api/crud/users.py

@@ -3,21 +3,23 @@
 import asyncio
 import uuid
 import warnings
+from typing import cast
 
 from boto3.dynamodb.conditions import Key as KeyCondition
 
 from store.app.api.crud.base import BaseCrud
-from store.app.api.model import Token, User
+from store.app.api.crypto import hash_api_key
+from store.app.api.model import ApiKey, User
 
 
 class UserCrud(BaseCrud):
     async def add_user(self, user: User) -> None:
         table = await self.db.Table("Users")
         await table.put_item(Item=user.model_dump())
 
-    async def get_user(self, user_id: str) -> User | None:
+    async def get_user(self, user_id: uuid.UUID) -> User | None:
         table = await self.db.Table("Users")
-        user_dict = await table.get_item(Key={"user_id": user_id})
+        user_dict = await table.get_item(Key={"user_id": str(user_id)})
         if "Item" not in user_dict:
             return None
         user = User.model_validate(user_dict["Item"])
@@ -34,6 +36,15 @@ async def get_user_from_email(self, email: str) -> User | None:
         user = User.model_validate(items[0])
         return user
 
+    async def get_user_id_from_api_key(self, api_key: uuid.UUID) -> uuid.UUID | None:
+        table = await self.db.Table("ApiKeys")
+        api_key_hash = hash_api_key(api_key)
+        row = await table.get_item(Key={"api_key_hash": api_key_hash})
+        if "Item" not in row:
+            return None
+        user_id = cast(str, row["Item"]["user_id"])
+        return uuid.UUID(user_id)
+
     async def delete_user(self, user: User) -> None:
         table = await self.db.Table("Users")
         await table.delete_item(Key={"user_id": user.user_id})
@@ -48,23 +59,21 @@ async def get_user_count(self) -> int:
         table = await self.db.Table("Users")
         return await table.item_count
 
-    async def add_token(self, token: Token) -> None:
-        table = await self.db.Table("Tokens")
-        await table.put_item(Item=token.model_dump())
-
-    async def get_token(self, token_id: str) -> Token | None:
-        table = await self.db.Table("Tokens")
-        token_dict = await table.get_item(Key={"token_id": token_id})
-        if "Item" not in token_dict:
-            return None
-        token = Token.model_validate(token_dict["Item"])
-        return token
-
-    async def get_user_tokens(self, user_id: str) -> list[Token]:
-        table = await self.db.Table("Tokens")
-        tokens = table.query(IndexName="userIdIndex", KeyConditionExpression=KeyCondition("user_id").eq(user_id))
-        tokens = [Token.model_validate(token) for token in await tokens]
-        return tokens
+    async def add_api_key(self, api_key: uuid.UUID, user_id: uuid.UUID) -> None:
+        row = ApiKey.from_api_key(api_key, user_id)
+        table = await self.db.Table("ApiKeys")
+        await table.put_item(Item=row.model_dump())
+
+    async def check_api_key(self, api_key: uuid.UUID, user_id: uuid.UUID) -> bool:
+        table = await self.db.Table("ApiKeys")
+        row = await table.get_item(Key={"api_key_hash": hash_api_key(api_key)})
+        if "Item" not in row:
+            return False
+        return row["Item"]["user_id"] == str(user_id)
+
+    async def delete_api_key(self, api_key: uuid.UUID) -> None:
+        table = await self.db.Table("ApiKeys")
+        await table.delete_item(Key={"api_key_hash": hash_api_key(api_key)})
 
 
 async def test_adhoc() -> None:

store/app/api/crypto.py

@@ -0,0 +1,33 @@
+"""Defines crypto functions."""
+
+import datetime
+import hashlib
+import uuid
+from typing import Any
+
+import jwt
+
+from store.settings import settings
+
+
+def hash_api_key(api_key: uuid.UUID) -> str:
+    return hashlib.sha256(api_key.bytes).hexdigest()
+
+
+def get_new_user_id() -> uuid.UUID:
+    return uuid.uuid4()
+
+
+def get_new_api_key(user_id: uuid.UUID) -> uuid.UUID:
+    user_id_hash = hashlib.sha1(user_id.bytes).digest()
+    return uuid.UUID(bytes=user_id_hash[:16], version=5)
+
+
+def encode_jwt(data: dict[str, Any], expire_after: datetime.timedelta | None = None) -> str:  # noqa: ANN401
+    if expire_after is not None:
+        data["exp"] = datetime.datetime.utcnow() + expire_after
+    return jwt.encode(data, settings.crypto.jwt_secret, algorithm=settings.crypto.algorithm)
+
+
+def decode_jwt(token: str) -> dict[str, Any]:  # noqa: ANN401
+    return jwt.decode(token, settings.crypto.jwt_secret, algorithms=[settings.crypto.algorithm])

store/app/api/db.py

@@ -43,9 +43,9 @@ async def create_tables(crud: Crud | None = None) -> None:
             ],
         )
         await crud._create_dynamodb_table(
-            name="Tokens",
+            name="ApiKeys",
             keys=[
-                ("token_id", "S", "HASH"),
+                ("api_key_hash", "S", "HASH"),
             ],
             gsis=[
                 ("userIdIndex", "user_id", "S", "HASH"),

store/app/api/email.py

@@ -11,7 +11,7 @@
 
 import aiosmtplib
 
-from store.app.api.token import create_token, load_token
+from store.app.api.crypto import decode_jwt, encode_jwt
 from store.settings import settings
 
 logger = logging.getLogger(__name__)
@@ -40,11 +40,11 @@ class OneTimePassPayload:
     def encode(self) -> str:
         expire_minutes = settings.crypto.expire_otp_minutes
         expire_after = datetime.timedelta(minutes=expire_minutes)
-        return create_token({"email": self.email}, expire_after=expire_after)
+        return encode_jwt({"email": self.email}, expire_after=expire_after)
 
     @classmethod
     def decode(cls, payload: str) -> "OneTimePassPayload":
-        data = load_token(payload)
+        data = decode_jwt(payload)
         return cls(email=data["email"])
 
 
@@ -53,7 +53,7 @@ async def send_otp_email(payload: OneTimePassPayload, login_url: str) -> None:
 
     body = textwrap.dedent(
         f"""
-            <h1><code>don't panic</code><br/><code>stay human</code></h1>
+            <h1><code>K-Scale Labs</code></h1>
             <h2><code><a href="{url}">log in</a></code></h2>
             <p>Or copy-paste this link: {url}</p>
         """
@@ -65,7 +65,7 @@ async def send_otp_email(payload: OneTimePassPayload, login_url: str) -> None:
 async def send_delete_email(email: str) -> None:
     body = textwrap.dedent(
         """
-            <h1><code>don't panic</code><br/><code>stay human</code></h1>
+            <h1><code>K-Scale Labs</code></h1>
             <h2><code>your account has been deleted</code></h2>
         """
     )
@@ -76,7 +76,7 @@ async def send_delete_email(email: str) -> None:
 async def send_waitlist_email(email: str) -> None:
     body = textwrap.dedent(
         """
-            <h1><code>don't panic</code><br/><code>stay human</code></h1>
+            <h1><code>K-Scale Labs</code></h1>
             <h2><code>you're on the waitlist!</code></h2>
             <p>Thanks for signing up! We'll let you know when you can log in.</p>
         """

store/app/api/model.py

@@ -1,22 +1,42 @@
-"""Defines the table models for the API."""
+"""Defines the table models for the API.
+
+These correspond directly with the rows in our database, and provide helper
+methods for converting from our input data into the format the database
+expects (for example, converting a UUID into a string).
+"""
 
 import datetime
+import uuid
 from dataclasses import field
 from decimal import Decimal
 
 from pydantic import BaseModel
 
+from store.app.api.crypto import hash_api_key
+
 
 class User(BaseModel):
     user_id: str  # Primary key
     email: str
 
+    @classmethod
+    def from_uuid(cls, user_id: uuid.UUID, email: str) -> "User":
+        return cls(user_id=str(user_id), email=email)
+
+    def to_uuid(self) -> uuid.UUID:
+        return uuid.UUID(self.user_id)
 
-class Token(BaseModel):
-    token_id: str  # Primary key
+
+class ApiKey(BaseModel):
+    api_key_hash: str  # Primary key
     user_id: str
     issued: Decimal = field(default_factory=lambda: Decimal(datetime.datetime.now().timestamp()))
 
+    @classmethod
+    def from_api_key(cls, api_key: uuid.UUID, user_id: uuid.UUID) -> "ApiKey":
+        api_key_hash = hash_api_key(api_key)
+        return cls(api_key_hash=api_key_hash, user_id=str(user_id))
+
 
 class PurchaseLink(BaseModel):
     name: str