Update app/services/key_service.py

app/services/key_service.py

@@ -1,92 +1,100 @@
-import secrets
-from hashlib import sha256
-from datetime import datetime, timedelta, timezone
-
-from app.services.repos import key_state_repo
-
-
-def _hash_key(raw: str) -> str:
-    return sha256(raw.encode("utf-8")).hexdigest()
-
-
-class KeyService:
-    def __init__(self, *, ttl_seconds: int = 60, lock_minutes: int = 5) -> None:
-        self.ttl_seconds = ttl_seconds
-        self.lock_minutes = lock_minutes
-
-    async def generate(self, device_id: str) -> tuple[str, datetime]:
-        key = secrets.token_urlsafe(8)
-        expires_at = datetime.now(timezone.utc) + timedelta(seconds=self.ttl_seconds)
-        await key_state_repo.upsert(
-            device_id,
-            {
-                "key_hash": _hash_key(key),
-                "expires_at": expires_at,
-                "wrong_attempts": 0,
-                "lock_until": None,
-            },
-        )
-        return key, expires_at
-
-    async def validate(self, device_id: str, candidate: str) -> tuple[bool, datetime | None, int]:
-        state = await key_state_repo.get(device_id)
-        now = datetime.now(timezone.utc)
-
-        if state is None:
-            lock_until = now + timedelta(minutes=self.lock_minutes)
-            await key_state_repo.upsert(
-                device_id,
-                {
-                    "wrong_attempts": 1,
-                    "lock_until": lock_until,
-                    "expires_at": now,
-                    "key_hash": "",
-                },
-            )
-            return False, lock_until, 1
-
-        lock_until_value = state.get("lock_until")
-        if isinstance(lock_until_value, datetime) and now < lock_until_value:
-            return False, lock_until_value, int(state.get("wrong_attempts", 0))
-
-        expires_at = state.get("expires_at")
-        if not isinstance(expires_at, datetime) or now > expires_at:
-            lock_until = now + timedelta(minutes=self.lock_minutes)
-            attempts = int(state.get("wrong_attempts", 0)) + 1
-            await key_state_repo.upsert(
-                device_id,
-                {
-                    "wrong_attempts": attempts,
-                    "lock_until": lock_until,
-                    "expires_at": now,
-                },
-            )
-            return False, lock_until, attempts
-
-        candidate_hash = _hash_key(candidate)
-        expected_hash = str(state.get("key_hash", ""))
-        if not secrets.compare_digest(candidate_hash, expected_hash):
-            attempts = int(state.get("wrong_attempts", 0)) + 1
-            lock_until = now + timedelta(minutes=self.lock_minutes)
-            await key_state_repo.upsert(
-                device_id,
-                {
-                    "wrong_attempts": attempts,
-                    "lock_until": lock_until,
-                },
-            )
-            return False, lock_until, attempts
-
-        await key_state_repo.upsert(
-            device_id,
-            {
-                "wrong_attempts": 0,
-                "lock_until": None,
-                "expires_at": now,
-                "key_hash": "",
-            },
-        )
-        return True, None, 0
-
-
-key_service = KeyService()
+import secrets
+from hashlib import sha256
+from datetime import datetime, timedelta, timezone
+
+from app.services.repos import key_state_repo
+
+
+def _hash_key(raw: str) -> str:
+    return sha256(raw.encode("utf-8")).hexdigest()
+
+
+def _to_utc_aware(value: datetime | None) -> datetime | None:
+    if not isinstance(value, datetime):
+        return None
+    if value.tzinfo is None:
+        return value.replace(tzinfo=timezone.utc)
+    return value.astimezone(timezone.utc)
+
+
+class KeyService:
+    def __init__(self, *, ttl_seconds: int = 60, lock_minutes: int = 5) -> None:
+        self.ttl_seconds = ttl_seconds
+        self.lock_minutes = lock_minutes
+
+    async def generate(self, device_id: str) -> tuple[str, datetime]:
+        key = secrets.token_urlsafe(8)
+        expires_at = datetime.now(timezone.utc) + timedelta(seconds=self.ttl_seconds)
+        await key_state_repo.upsert(
+            device_id,
+            {
+                "key_hash": _hash_key(key),
+                "expires_at": expires_at,
+                "wrong_attempts": 0,
+                "lock_until": None,
+            },
+        )
+        return key, expires_at
+
+    async def validate(self, device_id: str, candidate: str) -> tuple[bool, datetime | None, int]:
+        state = await key_state_repo.get(device_id)
+        now = datetime.now(timezone.utc)
+
+        if state is None:
+            lock_until = now + timedelta(minutes=self.lock_minutes)
+            await key_state_repo.upsert(
+                device_id,
+                {
+                    "wrong_attempts": 1,
+                    "lock_until": lock_until,
+                    "expires_at": now,
+                    "key_hash": "",
+                },
+            )
+            return False, lock_until, 1
+
+        lock_until_value = _to_utc_aware(state.get("lock_until"))
+        if lock_until_value is not None and now < lock_until_value:
+            return False, lock_until_value, int(state.get("wrong_attempts", 0))
+
+        expires_at = _to_utc_aware(state.get("expires_at"))
+        if expires_at is None or now > expires_at:
+            lock_until = now + timedelta(minutes=self.lock_minutes)
+            attempts = int(state.get("wrong_attempts", 0)) + 1
+            await key_state_repo.upsert(
+                device_id,
+                {
+                    "wrong_attempts": attempts,
+                    "lock_until": lock_until,
+                    "expires_at": now,
+                },
+            )
+            return False, lock_until, attempts
+
+        candidate_hash = _hash_key(candidate)
+        expected_hash = str(state.get("key_hash", ""))
+        if not secrets.compare_digest(candidate_hash, expected_hash):
+            attempts = int(state.get("wrong_attempts", 0)) + 1
+            lock_until = now + timedelta(minutes=self.lock_minutes)
+            await key_state_repo.upsert(
+                device_id,
+                {
+                    "wrong_attempts": attempts,
+                    "lock_until": lock_until,
+                },
+            )
+            return False, lock_until, attempts
+
+        await key_state_repo.upsert(
+            device_id,
+            {
+                "wrong_attempts": 0,
+                "lock_until": None,
+                "expires_at": now,
+                "key_hash": "",
+            },
+        )
+        return True, None, 0
+
+
+key_service = KeyService()