From be5a181ec54c4e95bb9b4835449f8c7e37a17b66 Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Mon, 27 Jan 2025 07:19:55 +0100
Subject: [PATCH] test: add thp device tests [no changelog]

---
 tests/conftest.py                  |  2 +
 tests/device_tests/thp/__init__.py |  0
 tests/device_tests/thp/test_thp.py | 71 ++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 tests/device_tests/thp/__init__.py
 create mode 100644 tests/device_tests/thp/test_thp.py

diff --git a/tests/conftest.py b/tests/conftest.py
index e3bfae1d22..ff30553e8b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -338,6 +338,8 @@ def _client_unlocked(
 
     while True:
         try:
+            if _raw_client.is_invalidated:
+                _raw_client = _get_raw_client(request)
             session = _raw_client.get_management_session()
             wipe_device(session)
             sleep(1.5)  # Makes tests more stable (wait for wipe to finish)
diff --git a/tests/device_tests/thp/__init__.py b/tests/device_tests/thp/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/device_tests/thp/test_thp.py b/tests/device_tests/thp/test_thp.py
new file mode 100644
index 0000000000..e26c3cba6a
--- /dev/null
+++ b/tests/device_tests/thp/test_thp.py
@@ -0,0 +1,71 @@
+import os
+import typing as t
+
+import pytest
+import typing_extensions as tx
+
+from trezorlib.client import ProtocolV2
+from trezorlib.debuglink import TrezorClientDebugLink as Client
+from trezorlib.transport.thp import curve25519
+from trezorlib.transport.thp.protocol_v2 import _hkdf
+
+if t.TYPE_CHECKING:
+    P = tx.ParamSpec("P")
+
+pytestmark = [pytest.mark.protocol("protocol_v2")]
+
+
+def test_allocate_channel(client: Client) -> None:
+    protocol: ProtocolV2 = client.protocol
+    nonce = b"\x1A\x2B\x3B\x4A\x5C\x6D\x7E\x8F"
+
+    # Use valid nonce
+    protocol._send_channel_allocation_request(nonce)
+    protocol._read_channel_allocation_response(nonce)
+
+    # Expect different nonce
+    protocol._send_channel_allocation_request(nonce)
+    with pytest.raises(Exception, match="Invalid channel allocation response."):
+        protocol._read_channel_allocation_response(
+            expected_nonce=b"\xDE\xAD\xBE\xEF\xDE\xAD\xBE\xEF"
+        )
+    client.invalidate()
+
+
+def test_handshake(client: Client) -> None:
+    protocol: ProtocolV2 = client.protocol
+
+    protocol.sync_bit_send = 0
+    protocol.sync_bit_receive = 0
+    host_ephemeral_privkey = curve25519.get_private_key(os.urandom(32))
+    host_ephemeral_pubkey = curve25519.get_public_key(host_ephemeral_privkey)
+
+    protocol._do_channel_allocation()
+    protocol._send_handshake_init_request(host_ephemeral_pubkey)
+    protocol._read_ack()
+    init_response = protocol._read_handshake_init_response()
+
+    trezor_ephemeral_pubkey = init_response[:32]
+    encrypted_trezor_static_pubkey = init_response[32:80]
+    noise_tag = init_response[80:96]
+
+    # TODO check noise_tag is valid
+
+    ck = protocol._send_handshake_completion_request(
+        host_ephemeral_pubkey,
+        host_ephemeral_privkey,
+        trezor_ephemeral_pubkey,
+        encrypted_trezor_static_pubkey,
+    )
+    protocol._read_ack()
+    protocol._read_handshake_completion_response()
+    protocol.key_request, protocol.key_response = _hkdf(ck, b"")
+    protocol.nonce_request = 0
+    protocol.nonce_response = 1
+
+    # TODO - without pairing, the client is damaged and results in fail of the following test
+    # so far no luck in solving it - it should be also tackled in FW, as it causes unexpected FW error
+    protocol._do_pairing()
+
+    # TODO the following is just to make style checker happy
+    assert noise_tag is not None