From 5d4ed7343542b0d61392b82200ed92b0480ac4f8 Mon Sep 17 00:00:00 2001
From: Houkime <>
Date: Mon, 26 Dec 2022 14:01:36 +0000
Subject: [PATCH] fix(tokens-repo): properly expire new device keys

---
 .../tokens/abstract_tokens_repository.py      |  3 ++
 .../test_repository/test_tokens_repository.py | 37 ++++++++++++++++---
 tests/test_models.py                          |  8 ++++
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py b/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py
index bee76b7..7601ef4 100644
--- a/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py
+++ b/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py
@@ -140,6 +140,9 @@ class AbstractTokensRepository(ABC):
         if not new_device_key:
             raise NewDeviceKeyNotFound
 
+        if not new_device_key.is_valid():
+            raise NewDeviceKeyNotFound
+
         if not self._assert_mnemonic(new_device_key.key, mnemonic_phrase):
             raise NewDeviceKeyNotFound("Phrase is not token!")
 
diff --git a/tests/test_graphql/test_repository/test_tokens_repository.py b/tests/test_graphql/test_repository/test_tokens_repository.py
index 95ba7c0..a95878d 100644
--- a/tests/test_graphql/test_repository/test_tokens_repository.py
+++ b/tests/test_graphql/test_repository/test_tokens_repository.py
@@ -513,15 +513,20 @@ def test_use_not_exists_mnemonic_new_device_key(
         )
 
 
-def test_use_mnemonic_new_device_key(
-    empty_repo, mock_new_device_key_generate_for_mnemonic
-):
+def mnemonic_from_hex(hexkey):
+    return Mnemonic(language="english").to_mnemonic(bytes.fromhex(hexkey))
+
+
+def test_use_mnemonic_new_device_key(empty_repo):
     repo = empty_repo
-    assert repo.get_new_device_key() is not None
+    key = repo.get_new_device_key()
+    assert key is not None
+
+    mnemonic_phrase = mnemonic_from_hex(key.key)
 
     new_token = repo.use_mnemonic_new_device_key(
         device_name="imnew",
-        mnemonic_phrase="captain ribbon toddler settle symbol minute step broccoli bless universe divide bulb",
+        mnemonic_phrase=mnemonic_phrase,
     )
 
     assert new_token.device_name == "imnew"
@@ -532,12 +537,32 @@ def test_use_mnemonic_new_device_key(
         assert (
             repo.use_mnemonic_new_device_key(
                 device_name="imnew",
-                mnemonic_phrase="captain ribbon toddler settle symbol minute step broccoli bless universe divide bulb",
+                mnemonic_phrase=mnemonic_phrase,
             )
             is None
         )
 
 
+def test_use_mnemonic_expired_new_device_key(
+    some_tokens_repo,
+):
+    repo = some_tokens_repo
+    expiration = datetime.now() - timedelta(minutes=5)
+
+    key = repo.get_new_device_key()
+    assert key is not None
+    assert key.expires_at is not None
+    key.expires_at = expiration
+    assert not key.is_valid()
+    repo._store_new_device_key(key)
+
+    with pytest.raises(NewDeviceKeyNotFound):
+        token = repo.use_mnemonic_new_device_key(
+            mnemonic_phrase=mnemonic_from_hex(key.key),
+            device_name="imnew",
+        )
+
+
 def test_use_mnemonic_new_device_key_when_empty(empty_repo):
     repo = empty_repo
 
diff --git a/tests/test_models.py b/tests/test_models.py
index 078dfd5..2263e82 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -2,9 +2,17 @@ import pytest
 from datetime import datetime, timedelta
 
 from selfprivacy_api.models.tokens.recovery_key import RecoveryKey
+from selfprivacy_api.models.tokens.new_device_key import NewDeviceKey
 
 
 def test_recovery_key_expired():
     expiration = datetime.now() - timedelta(minutes=5)
     key = RecoveryKey.generate(expiration=expiration, uses_left=2)
     assert not key.is_valid()
+
+
+def test_new_device_key_expired():
+    expiration = datetime.now() - timedelta(minutes=5)
+    key = NewDeviceKey.generate()
+    key.expires_at = expiration
+    assert not key.is_valid()