From ce803a5452c5d12eb7b026fa2c20b521957fb8c8 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Wed, 23 Oct 2024 14:26:59 +0200
Subject: [PATCH] feat(tests): Tests for entropy check workflow in ResetDevice.

---
 .../reset_recovery/test_reset_bip39_t2.py     | 41 ++++++++++++++++++
 .../reset_recovery/test_reset_slip39_basic.py | 42 +++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/tests/device_tests/reset_recovery/test_reset_bip39_t2.py b/tests/device_tests/reset_recovery/test_reset_bip39_t2.py
index 9715d118b3..463f58732f 100644
--- a/tests/device_tests/reset_recovery/test_reset_bip39_t2.py
+++ b/tests/device_tests/reset_recovery/test_reset_bip39_t2.py
@@ -18,6 +18,7 @@ import pytest
 from mnemonic import Mnemonic
 
 from trezorlib import device, messages
+from trezorlib.btc import get_public_node
 from trezorlib.debuglink import LayoutType
 from trezorlib.debuglink import TrezorClientDebugLink as Client
 from trezorlib.exceptions import TrezorFailure
@@ -110,6 +111,46 @@ def test_reset_device_pin(client: Client):
     assert resp.passphrase_protection is True
 
 
+@pytest.mark.setup_client(uninitialized=True)
+def test_reset_entropy_check(client: Client):
+    strength = 128  # 12 words
+
+    with WITH_MOCK_URANDOM, client:
+        IF = InputFlowBip39ResetBackup(client)
+        client.set_input_flow(IF.get())
+
+        # No PIN, no passphrase
+        _, path_xpubs = device.reset_entropy_check(
+            client,
+            strength=strength,
+            passphrase_protection=False,
+            pin_protection=False,
+            label="test",
+            entropy_check_count=2,
+        )
+
+    # Generate the mnemonic locally.
+    internal_entropy = client.debug.state().reset_entropy
+    entropy = generate_entropy(strength, internal_entropy, EXTERNAL_ENTROPY)
+    expected_mnemonic = Mnemonic("english").to_mnemonic(entropy)
+
+    # Check that the device generated the correct mnemonic for the given entropies.
+    assert IF.mnemonic == expected_mnemonic
+
+    # Check that the device is properly initialized.
+    resp = client.call_raw(messages.Initialize())
+    assert resp.initialized is True
+    assert resp.backup_availability == messages.BackupAvailability.NotAvailable
+    assert resp.pin_protection is False
+    assert resp.passphrase_protection is False
+    assert resp.backup_type is messages.BackupType.Bip39
+
+    # Check that the XPUBs are the same as those from the entropy check.
+    for path, xpub in path_xpubs:
+        res = get_public_node(client, path)
+        assert res.xpub == xpub
+
+
 @pytest.mark.setup_client(uninitialized=True)
 def test_reset_failed_check(client: Client):
     strength = 256  # 24 words
diff --git a/tests/device_tests/reset_recovery/test_reset_slip39_basic.py b/tests/device_tests/reset_recovery/test_reset_slip39_basic.py
index 8eb5d7830f..b0d39f9eb1 100644
--- a/tests/device_tests/reset_recovery/test_reset_slip39_basic.py
+++ b/tests/device_tests/reset_recovery/test_reset_slip39_basic.py
@@ -20,6 +20,7 @@ import pytest
 from shamir_mnemonic import MnemonicError, shamir
 
 from trezorlib import device
+from trezorlib.btc import get_public_node
 from trezorlib.debuglink import TrezorClientDebugLink as Client
 from trezorlib.exceptions import TrezorFailure
 from trezorlib.messages import BackupAvailability, BackupType
@@ -76,6 +77,47 @@ def test_reset_device_slip39_basic_256(client: Client):
     reset_device(client, 256)
 
 
+@pytest.mark.setup_client(uninitialized=True)
+def test_reset_entropy_check(client: Client):
+    member_threshold = 3
+
+    strength = 128  # 20 words
+
+    with WITH_MOCK_URANDOM, client:
+        IF = InputFlowSlip39BasicResetRecovery(client)
+        client.set_input_flow(IF.get())
+
+        # No PIN, no passphrase.
+        _, path_xpubs = device.reset_entropy_check(
+            client,
+            strength=strength,
+            passphrase_protection=False,
+            pin_protection=False,
+            label="test",
+            backup_type=BackupType.Slip39_Basic,
+            entropy_check_count=3,
+        )
+
+    # Generate the master secret locally.
+    internal_entropy = client.debug.state().reset_entropy
+    secret = generate_entropy(strength, internal_entropy, EXTERNAL_ENTROPY)
+
+    # Check that all combinations will result in the correct master secret.
+    validate_mnemonics(IF.mnemonics, member_threshold, secret)
+
+    # Check that the device is properly initialized.
+    assert client.features.initialized is True
+    assert client.features.backup_availability == BackupAvailability.NotAvailable
+    assert client.features.pin_protection is False
+    assert client.features.passphrase_protection is False
+    assert client.features.backup_type is BackupType.Slip39_Basic_Extendable
+
+    # Check that the XPUBs are the same as those from the entropy check.
+    for path, xpub in path_xpubs:
+        res = get_public_node(client, path)
+        assert res.xpub == xpub
+
+
 def validate_mnemonics(mnemonics, threshold, expected_ems):
     # We expect these combinations to recreate the secret properly
     for test_group in combinations(mnemonics, threshold):