tests: simplify change_wipe_code test with PIN queuing

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())
+    with client:
-    assert isinstance(ret, messages.ButtonRequest)
+        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.set_pin(pins)
-    client.debug.press_yes()
+        client.set_expected_responses(
-    ret = client.call_raw(messages.ButtonAck())
+            [messages.ButtonRequest()]
-
+            + pin_matrices
-    if client.features.pin_protection:
+            + [messages.Success(), messages.Features()]
-        # Send current PIN.
+        )
-        assert isinstance(ret, messages.PinMatrixRequest)
+        device.change_wipe_code(client)
-        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)
 
 
-def _remove_wipe_code(client):
+def _change_pin(client, old_pin, new_pin):
-    # Remove wipe code
+    assert client.features.pin_protection is True
-    ret = client.call_raw(messages.ChangeWipeCode(remove=True))
+    with client:
-    assert isinstance(ret, messages.ButtonRequest)
+        client.set_pin([old_pin, new_pin, new_pin])
-
+        try:
-    # Confirm intent to remove wipe code.
+            return device.change_pin(client)
-    client.debug.press_yes()
+        except exceptions.TrezorFailure as f:
-    ret = client.call_raw(messages.ButtonAck())
+            return f.failure
-
-    # 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 _check_wipe_code(client, wipe_code):
-    # Try to change the PIN to the current wipe code value. The operation should fail.
+    """Check that wipe code is set by changing the PIN to it."""
-    ret = client.call_raw(messages.ChangePin())
+    old_pin, _ = client.debug.read_pin()
-    assert isinstance(ret, messages.ButtonRequest)
+    f = _change_pin(client, old_pin, wipe_code)
-
+    assert isinstance(f, messages.Failure)
-    # 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)
 
 
 @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())
+    with client:
-    assert isinstance(ret, messages.ButtonRequest)
+        client.set_pin([WIPE_CODE4, WIPE_CODE6])
-
+        client.set_expected_responses(
-    # Confirm intent to set wipe code.
+            [
-    client.debug.press_yes()
+                messages.ButtonRequest(),
-    ret = client.call_raw(messages.ButtonAck())
+                messages.PinMatrixRequest(type=PinType.WipeCodeFirst),
-
+                messages.PinMatrixRequest(type=PinType.WipeCodeSecond),
-    # Send the new wipe code for the first time.
+                messages.Failure(code=messages.FailureType.WipeCodeMismatch),
-    assert isinstance(ret, messages.PinMatrixRequest)
+            ]
-    wipe_code_encoded = client.debug.encode_pin(WIPE_CODE4)
+        )
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
+        with pytest.raises(exceptions.TrezorFailure):
-
+            device.change_wipe_code(client)
-    # 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
 
     # 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())
+    with client:
-    assert isinstance(ret, messages.ButtonRequest)
+        client.set_pin([PIN4, PIN4])
-
+        client.set_expected_responses(
-    # Confirm intent to set wipe code.
+            [
-    client.debug.press_yes()
+                messages.ButtonRequest(),
-    ret = client.call_raw(messages.ButtonAck())
+                messages.PinMatrixRequest(type=PinType.Current),
-
+                messages.PinMatrixRequest(type=PinType.WipeCodeFirst),
-    # Send current PIN.
+                messages.Failure(code=messages.FailureType.ProcessError),
-    assert isinstance(ret, messages.PinMatrixRequest)
+            ]
-    pin_encoded = client.debug.read_pin_encoded()
+        )
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
+        with pytest.raises(exceptions.TrezorFailure):
-
+            device.change_wipe_code(client)
-    # 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
 
     # 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())
+    with client:
-    assert isinstance(ret, messages.ButtonRequest)
+        client.set_pin([WIPE_CODE4, WIPE_CODE4])
-
+        client.set_expected_responses(
-    # Confirm intent to set PIN.
+            [
-    client.debug.press_yes()
+                messages.ButtonRequest(),
-    ret = client.call_raw(messages.ButtonAck())
+                messages.PinMatrixRequest(type=PinType.NewFirst),
-
+                messages.PinMatrixRequest(type=PinType.NewSecond),
-    # Send the new PIN for the first time.
+                messages.Failure(code=messages.FailureType.ProcessError),
-    assert isinstance(ret, messages.PinMatrixRequest)
+            ]
-    pin_encoded = client.debug.encode_pin(WIPE_CODE4)
+        )
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
+        with pytest.raises(exceptions.TrezorFailure):
-
+            device.change_pin(client)
-    # 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
 
     # Check that there is no PIN protection.
     client.init_device()