diff --git a/tests/device_tests/test_msg_change_wipe_code_t1.py b/tests/device_tests/test_msg_change_wipe_code_t1.py index 8468324b99..aeff96b910 100644 --- a/tests/device_tests/test_msg_change_wipe_code_t1.py +++ b/tests/device_tests/test_msg_change_wipe_code_t1.py @@ -16,7 +16,9 @@ import pytest -from trezorlib import messages +from trezorlib import device, exceptions, messages + +PinType = messages.PinMatrixRequestType PIN4 = "1234" WIPE_CODE4 = "4321" @@ -27,77 +29,46 @@ pytestmark = pytest.mark.skip_t2 def _set_wipe_code(client, wipe_code): # Set/change wipe code. - ret = client.call_raw(messages.ChangeWipeCode()) - assert isinstance(ret, messages.ButtonRequest) + with client: + if client.features.pin_protection: + pin, _ = client.debug.read_pin() + pins = [pin, wipe_code, wipe_code] + pin_matrices = [ + messages.PinMatrixRequest(type=PinType.Current), + messages.PinMatrixRequest(type=PinType.WipeCodeFirst), + messages.PinMatrixRequest(type=PinType.WipeCodeSecond), + ] + else: + pins = [wipe_code, wipe_code] + pin_matrices = [ + messages.PinMatrixRequest(type=PinType.WipeCodeFirst), + messages.PinMatrixRequest(type=PinType.WipeCodeSecond), + ] - # Confirm intent to set/change wipe code. - client.debug.press_yes() - ret = client.call_raw(messages.ButtonAck()) - - if client.features.pin_protection: - # Send current PIN. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.read_pin_encoded() - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # Send the new wipe code for the first time. - assert isinstance(ret, messages.PinMatrixRequest) - wipe_code_encoded = client.debug.encode_pin(wipe_code) - ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded)) - - # Send the new wipe code for the second time. - assert isinstance(ret, messages.PinMatrixRequest) - wipe_code_encoded = client.debug.encode_pin(wipe_code) - ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded)) - - # Now we're done. - assert isinstance(ret, messages.Success) + client.set_pin(pins) + client.set_expected_responses( + [messages.ButtonRequest()] + + pin_matrices + + [messages.Success(), messages.Features()] + ) + device.change_wipe_code(client) -def _remove_wipe_code(client): - # Remove wipe code - ret = client.call_raw(messages.ChangeWipeCode(remove=True)) - assert isinstance(ret, messages.ButtonRequest) - - # Confirm intent to remove wipe code. - client.debug.press_yes() - ret = client.call_raw(messages.ButtonAck()) - - # Send current PIN. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.read_pin_encoded() - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # Now we're done. - assert isinstance(ret, messages.Success) +def _change_pin(client, old_pin, new_pin): + assert client.features.pin_protection is True + with client: + client.set_pin([old_pin, new_pin, new_pin]) + try: + return device.change_pin(client) + except exceptions.TrezorFailure as f: + return f.failure def _check_wipe_code(client, wipe_code): - # Try to change the PIN to the current wipe code value. The operation should fail. - ret = client.call_raw(messages.ChangePin()) - assert isinstance(ret, messages.ButtonRequest) - - # Confirm intent to change PIN. - client.debug.press_yes() - ret = client.call_raw(messages.ButtonAck()) - - # Send current PIN. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.read_pin_encoded() - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # Send the new wipe code for the first time. - assert isinstance(ret, messages.PinMatrixRequest) - wipe_code_encoded = client.debug.encode_pin(wipe_code) - ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded)) - - # Send the new wipe code for the second time. - assert isinstance(ret, messages.PinMatrixRequest) - wipe_code_encoded = client.debug.encode_pin(wipe_code) - ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded)) - - # Expect failure. - assert isinstance(ret, messages.Failure) + """Check that wipe code is set by changing the PIN to it.""" + old_pin, _ = client.debug.read_pin() + f = _change_pin(client, old_pin, wipe_code) + assert isinstance(f, messages.Failure) @pytest.mark.setup_client(pin=PIN4) @@ -122,11 +93,11 @@ def test_set_remove_wipe_code(client): client.init_device() assert client.features.wipe_code_protection is True - # Check that the PIN is correct. + # Check that the wipe code is correct. _check_wipe_code(client, WIPE_CODE6) # Test remove wipe code. - _remove_wipe_code(client) + device.change_wipe_code(client, remove=True) # Check that there's no wipe code protection now. client.init_device() @@ -138,26 +109,18 @@ def test_set_wipe_code_mismatch(client): assert client.features.wipe_code_protection is False # Let's set a new wipe code. - ret = client.call_raw(messages.ChangeWipeCode()) - assert isinstance(ret, messages.ButtonRequest) - - # Confirm intent to set wipe code. - client.debug.press_yes() - ret = client.call_raw(messages.ButtonAck()) - - # Send the new wipe code for the first time. - assert isinstance(ret, messages.PinMatrixRequest) - wipe_code_encoded = client.debug.encode_pin(WIPE_CODE4) - ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded)) - - # Send the new wipe code for the second time, but different. - assert isinstance(ret, messages.PinMatrixRequest) - wipe_code_encoded = client.debug.encode_pin(WIPE_CODE6) - ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded)) - - # The operation should fail, because the wipe codes are different. - assert isinstance(ret, messages.Failure) - assert ret.code == messages.FailureType.WipeCodeMismatch + with client: + client.set_pin([WIPE_CODE4, WIPE_CODE6]) + client.set_expected_responses( + [ + messages.ButtonRequest(), + messages.PinMatrixRequest(type=PinType.WipeCodeFirst), + messages.PinMatrixRequest(type=PinType.WipeCodeSecond), + messages.Failure(code=messages.FailureType.WipeCodeMismatch), + ] + ) + with pytest.raises(exceptions.TrezorFailure): + device.change_wipe_code(client) # Check that there is no wipe code protection. client.init_device() @@ -170,26 +133,18 @@ def test_set_wipe_code_to_pin(client): assert client.features.wipe_code_protection is None # Let's try setting the wipe code to the curent PIN value. - ret = client.call_raw(messages.ChangeWipeCode()) - assert isinstance(ret, messages.ButtonRequest) - - # Confirm intent to set wipe code. - client.debug.press_yes() - ret = client.call_raw(messages.ButtonAck()) - - # Send current PIN. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.read_pin_encoded() - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # Send the new wipe code. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.read_pin_encoded() - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # The operation should fail, because the wipe code must be different from the PIN. - assert isinstance(ret, messages.Failure) - assert ret.code == messages.FailureType.ProcessError + with client: + client.set_pin([PIN4, PIN4]) + client.set_expected_responses( + [ + messages.ButtonRequest(), + messages.PinMatrixRequest(type=PinType.Current), + messages.PinMatrixRequest(type=PinType.WipeCodeFirst), + messages.Failure(code=messages.FailureType.ProcessError), + ] + ) + with pytest.raises(exceptions.TrezorFailure): + device.change_wipe_code(client) # Check that there is no wipe code protection. client.init_device() @@ -201,26 +156,18 @@ def test_set_pin_to_wipe_code(client): _set_wipe_code(client, WIPE_CODE4) # Try to set the PIN to the current wipe code value. - ret = client.call_raw(messages.ChangePin()) - assert isinstance(ret, messages.ButtonRequest) - - # Confirm intent to set PIN. - client.debug.press_yes() - ret = client.call_raw(messages.ButtonAck()) - - # Send the new PIN for the first time. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.encode_pin(WIPE_CODE4) - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # Send the new PIN for the second time. - assert isinstance(ret, messages.PinMatrixRequest) - pin_encoded = client.debug.encode_pin(WIPE_CODE4) - ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded)) - - # The operation should fail, because the PIN must be different from the wipe code. - assert isinstance(ret, messages.Failure) - assert ret.code == messages.FailureType.ProcessError + with client: + client.set_pin([WIPE_CODE4, WIPE_CODE4]) + client.set_expected_responses( + [ + messages.ButtonRequest(), + messages.PinMatrixRequest(type=PinType.NewFirst), + messages.PinMatrixRequest(type=PinType.NewSecond), + messages.Failure(code=messages.FailureType.ProcessError), + ] + ) + with pytest.raises(exceptions.TrezorFailure): + device.change_pin(client) # Check that there is no PIN protection. client.init_device()