arno/trezor-firmware
1
0
Fork 0
mirror of https://github.com/trezor/trezor-firmware.git synced 2025-02-07 13:12:41 +00:00
Code Issues Releases Wiki Activity

tests(python): improve thp device tests

Browse Source 
[no changelog]
This commit is contained in:
M1nd3r 2025-01-29 21:41:36 +01:00
parent d8506f8f43
commit d636c3fa0d
1 changed files with 97 additions and 148 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

245
tests/device_tests/thp/test_thp.py
View File

@ -5,6 +5,7 @@ import typing as t
import pytest import pytest
import typing_extensions as tx import typing_extensions as tx
from trezorlib import protobuf
from trezorlib.client import ProtocolV2 from trezorlib.client import ProtocolV2
from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.debuglink import TrezorClientDebugLink as Client
from trezorlib.messages import ( from trezorlib.messages import (
@ -12,7 +13,9 @@ from trezorlib.messages import (
ButtonRequest, ButtonRequest,
ThpCodeEntryChallenge, ThpCodeEntryChallenge,
ThpCodeEntryCommitment, ThpCodeEntryCommitment,
ThpCodeEntryCpaceHostTag,
ThpCodeEntryCpaceTrezor, ThpCodeEntryCpaceTrezor,
ThpCodeEntrySecret,
ThpEndRequest, ThpEndRequest,
ThpEndResponse, ThpEndResponse,
ThpPairingMethod, ThpPairingMethod,
@ -29,11 +32,26 @@ from trezorlib.transport.thp.protocol_v2 import MANAGEMENT_SESSION_ID, _hkdf
if t.TYPE_CHECKING: if t.TYPE_CHECKING:
P = tx.ParamSpec("P") P = tx.ParamSpec("P")
MT = t.TypeVar("MT", bound=protobuf.MessageType)
pytestmark = [pytest.mark.protocol("protocol_v2")] pytestmark = [pytest.mark.protocol("protocol_v2")]
protocol: ProtocolV2
def _prepare_protocol(client: Client):
global protocol
protocol = client.protocol
protocol.sync_bit_send = 0
protocol.sync_bit_receive = 0
def test_allocate_channel(client: Client) -> None: def test_allocate_channel(client: Client) -> None:
protocol: ProtocolV2 = client.protocol global protocol
_prepare_protocol(client)
# protocol: ProtocolV2 = client.protocol
nonce = b"\x1A\x2B\x3B\x4A\x5C\x6D\x7E\x8F" nonce = b"\x1A\x2B\x3B\x4A\x5C\x6D\x7E\x8F"
# Use valid nonce # Use valid nonce
@ -50,10 +68,10 @@ def test_allocate_channel(client: Client) -> None:
def test_handshake(client: Client) -> None: def test_handshake(client: Client) -> None:
protocol: ProtocolV2 = client.protocol global protocol
_prepare_protocol(client)
# 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_privkey = curve25519.get_private_key(os.urandom(32))
host_ephemeral_pubkey = curve25519.get_public_key(host_ephemeral_privkey) host_ephemeral_pubkey = curve25519.get_public_key(host_ephemeral_privkey)
@ -88,10 +106,27 @@ def test_handshake(client: Client) -> None:
assert noise_tag is not None assert noise_tag is not None
def _send_message(
message: MT,
session_id: int = MANAGEMENT_SESSION_ID,
):
global protocol
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(session_id, message_type, message_data)
protocol._read_ack()
def _read_message(message_type: type[MT]) -> MT:
global protocol
_, msg_type, msg_data = protocol.read_and_decrypt()
msg = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(msg, message_type)
return msg
def test_pairing_qr_code(client: Client) -> None: def test_pairing_qr_code(client: Client) -> None:
protocol: ProtocolV2 = client.protocol global protocol
protocol.sync_bit_send = 0 _prepare_protocol(client)
protocol.sync_bit_receive = 0
# Generate ephemeral keys # Generate ephemeral keys
host_ephemeral_privkey = curve25519.get_private_key(os.urandom(32)) host_ephemeral_privkey = curve25519.get_private_key(os.urandom(32))
@ -101,96 +136,54 @@ def test_pairing_qr_code(client: Client) -> None:
protocol._do_handshake(host_ephemeral_privkey, host_ephemeral_pubkey) protocol._do_handshake(host_ephemeral_privkey, host_ephemeral_pubkey)
# Send StartPairingReqest message _send_message(ThpPairingRequest())
message = ThpPairingRequest()
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data) _read_message(ButtonRequest)
# Read ACK _send_message(ButtonAck())
protocol._read_ack()
# Read button request
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ButtonRequest)
# Send button ACK
message = ButtonAck()
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
protocol._read_ack()
client.debug.press_yes() client.debug.press_yes()
# Read PairingRequestApproved _read_message(ThpPairingRequestApproved)
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpPairingRequestApproved) _send_message(ThpSelectMethod(selected_pairing_method=ThpPairingMethod.QrCode))
message = ThpSelectMethod(selected_pairing_method=ThpPairingMethod.QrCode) _read_message(ThpPairingPreparationsFinished)
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
# Read ACK
protocol._read_ack()
# Read ThpPairingPreparationsFinished
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpPairingPreparationsFinished)
# QR Code shown # QR Code shown
# Read button request _read_message(ButtonRequest)
_, msg_type, msg_data = protocol.read_and_decrypt() _send_message(ButtonAck())
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ButtonRequest)
# Send button ACK
message = ButtonAck()
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
protocol._read_ack()
# Read code from "Trezor's display" using debuglink
state = client.debug.state(thp_channel_id=protocol.channel_id.to_bytes(2, "big")) state = client.debug.state(thp_channel_id=protocol.channel_id.to_bytes(2, "big"))
code = state.thp_pairing_code_qr_code
# Compute tag for response
sha_ctx = hashlib.sha256(protocol.handshake_hash) sha_ctx = hashlib.sha256(protocol.handshake_hash)
sha_ctx.update(state.thp_pairing_code_qr_code) sha_ctx.update(state.thp_pairing_code_qr_code)
tag = sha_ctx.digest() tag = sha_ctx.digest()
message_type, message_data = protocol.mapping.encode(ThpQrCodeTag(tag=tag)) _send_message(ThpQrCodeTag(tag=tag))
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
protocol._read_ack() secret_msg = _read_message(ThpQrCodeSecret)
# Read ThpQrCodeSecret # Check that the `code` was derived from the revealed secret
_, msg_type, msg_data = protocol.read_and_decrypt() sha_ctx = hashlib.sha256(ThpPairingMethod.QrCode.to_bytes(1, "big"))
maaa = protocol.mapping.decode(msg_type, msg_data) sha_ctx.update(protocol.handshake_hash)
assert isinstance(maaa, ThpQrCodeSecret) sha_ctx.update(secret_msg.secret)
computed_code = sha_ctx.digest()[:16]
assert code == computed_code
message = ThpEndRequest() _send_message(ThpEndRequest())
message_type, message_data = protocol.mapping.encode(message) _read_message(ThpEndResponse)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
# Read ACK
protocol._read_ack()
# Read ThpEndResponse
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpEndResponse)
protocol._has_valid_channel = True protocol._has_valid_channel = True
@pytest.mark.skip("Cpace is not implemented yet") # @pytest.mark.skip("Cpace is not implemented yet")
def test_pairing_code_entry(client: Client) -> None: def test_pairing_code_entry(client: Client) -> None:
protocol: ProtocolV2 = client.protocol global protocol
protocol.sync_bit_send = 0 _prepare_protocol(client)
protocol.sync_bit_receive = 0
# Generate ephemeral keys # Generate ephemeral keys
host_ephemeral_privkey = curve25519.get_private_key(os.urandom(32)) host_ephemeral_privkey = curve25519.get_private_key(os.urandom(32))
@ -200,101 +193,57 @@ def test_pairing_code_entry(client: Client) -> None:
protocol._do_handshake(host_ephemeral_privkey, host_ephemeral_pubkey) protocol._do_handshake(host_ephemeral_privkey, host_ephemeral_pubkey)
# Send StartPairingReqest message _send_message(ThpPairingRequest())
message = ThpPairingRequest()
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data) _read_message(ButtonRequest)
# Read ACK _send_message(ButtonAck())
protocol._read_ack()
# Read button request
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ButtonRequest)
# Send button ACK
message = ButtonAck()
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
protocol._read_ack()
client.debug.press_yes() client.debug.press_yes()
# Read PairingRequestApproved _read_message(ThpPairingRequestApproved)
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpPairingRequestApproved) _send_message(ThpSelectMethod(selected_pairing_method=ThpPairingMethod.CodeEntry))
message = ThpSelectMethod(selected_pairing_method=ThpPairingMethod.CodeEntry) commitment_msg = _read_message(ThpCodeEntryCommitment)
message_type, message_data = protocol.mapping.encode(message) commitment = commitment_msg.commitment
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
# Read ACK
protocol._read_ack()
# Read ThpCodeEntryCommitment
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpCodeEntryCommitment)
challenge = b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xAA\xBB\xCC\xDD\xEE\xFF" challenge = b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xAA\xBB\xCC\xDD\xEE\xFF"
message = ThpCodeEntryChallenge(challenge=challenge) _send_message(ThpCodeEntryChallenge(challenge=challenge))
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data) cpace_trezor = _read_message(ThpCodeEntryCpaceTrezor)
# Read ACK cpace_trezor_public_key = cpace_trezor.cpace_trezor_public_key
protocol._read_ack()
# Read ThpCodeEntryCpaceTrezor
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpCodeEntryCpaceTrezor)
_ = maaa.cpace_trezor_public_key
# Code Entry code shown # Code Entry code shown
# Read button request _read_message(ButtonRequest)
_, msg_type, msg_data = protocol.read_and_decrypt() _send_message(ButtonAck())
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ButtonRequest)
# Send button ACK
message = ButtonAck()
message_type, message_data = protocol.mapping.encode(message)
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
protocol._read_ack()
state = client.debug.state(thp_channel_id=protocol.channel_id.to_bytes(2, "big")) state = client.debug.state(thp_channel_id=protocol.channel_id.to_bytes(2, "big"))
code = state.thp_pairing_code_entry_code
sha_ctx = hashlib.sha256(protocol.handshake_hash) # TODO fix missing CPACE
sha_ctx.update(state.thp_pairing_code_entry_code) cpace_shared_secret = b"\x01"
sha_ctx = hashlib.sha256(cpace_shared_secret)
tag = sha_ctx.digest() tag = sha_ctx.digest()
message_type, message_data = protocol.mapping.encode(ThpQrCodeTag(tag=tag)) cpace_host_public_key = cpace_trezor_public_key
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data)
protocol._read_ack() _send_message(
ThpCodeEntryCpaceHostTag(
cpace_host_public_key=cpace_host_public_key,
tag=tag,
)
)
# Read ThpQrCodeSecret secret_msg = _read_message(ThpCodeEntrySecret)
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpQrCodeSecret)
message = ThpEndRequest() # Check `commitment` and `code`
message_type, message_data = protocol.mapping.encode(message) sha_ctx = hashlib.sha256(secret_msg.secret)
computed_commitment = sha_ctx.digest()
assert commitment == computed_commitment
assert code == b"" # TODO implement
protocol._encrypt_and_write(MANAGEMENT_SESSION_ID, message_type, message_data) _send_message(ThpEndRequest())
# Read ACK _read_message(ThpEndResponse)
protocol._read_ack()
# Read ThpEndResponse
_, msg_type, msg_data = protocol.read_and_decrypt()
maaa = protocol.mapping.decode(msg_type, msg_data)
assert isinstance(maaa, ThpEndResponse)
protocol._has_valid_channel = True protocol._has_valid_channel = True