Add optional parameter to reset API key functions (#923)

Author: tomchop

core/schemas/user.py

@@ -1,3 +1,4 @@
+import re
 import secrets
 from typing import ClassVar
 
@@ -28,8 +29,13 @@ class User(BaseModel, database_arango.ArangoYetiConnector):
     def load(cls, object: dict) -> "User":
         return cls(**object)
 
-    def reset_api_key(self) -> None:
-        self.api_key = secrets.token_hex(32)
+    def reset_api_key(self, api_key=None) -> None:
+        if api_key:
+            if not re.match(r"^[a-f0-9]{64}$", api_key):
+                raise ValueError("Invalid API key: must match ^[a-f0-9]{64}$")
+            self.api_key = api_key
+        else:
+            self.api_key = secrets.token_hex(32)
 
 
 class UserSensitive(User):

tests/schemas/user.py

@@ -27,3 +27,21 @@ def test_set_user_password(self) -> None:
         self.assertEqual(user.username, "tomchop")
         self.assertTrue(user.verify_password("test"))
         self.assertFalse(user.verify_password("password"))
+
+    def test_reset_api_key(self) -> None:
+        old_api_key = self.user1.api_key
+        self.user1.reset_api_key()
+        self.user1.save()
+
+        user = UserSensitive.find(username="tomchop")
+        assert user is not None
+        self.assertNotEqual(old_api_key, self.user1.api_key)
+        self.assertRegex(user.api_key, r"[a-f0-9]{64}")
+
+    def test_reset_api_key_with_param(self) -> None:
+        self.user1.reset_api_key(api_key="1234123412341234123412341234123412341234123412341234123412341234")
+        self.user1.save()
+
+        user = UserSensitive.find(username="tomchop")
+        assert user is not None
+        self.assertEqual(self.user1.api_key, "1234123412341234123412341234123412341234123412341234123412341234")

yetictl/cli.py

@@ -21,14 +21,17 @@ def list_users():
 @click.argument("username")
 @click.argument("password")
 @click.option("--admin", is_flag=True, default=False)
-def create_user(username: str, password: str, admin: bool = False) -> None:
+@click.option("--api_key")
+def create_user(username: str, password: str, admin: bool = False, api_key: str | None = None) -> None:
     """Creates a new user in the system."""
     user = UserSensitive.find(username=username)
     if user:
         raise RuntimeError(f"User with username {username} already exists")
     user = UserSensitive(username=username, admin=admin)
     user.set_password(password)
     user.save()
+    if api_key:
+        user.reset_api_key(api_key=api_key)
     click.echo(
         f"User {username} succesfully created! API key: {username}:{user.api_key}"
     )