tests: simplify and extend session_id test

tests/device_tests/test_session_id_and_passphrase.py

@@ -24,16 +24,21 @@ XPUB_PASSPHRASE_A = "xpub6CekxGcnqnJ6osfY4Rrq7W5ogFtR54KUvz4H16XzaQuukMFZCGebEpV
 XPUB_PASSPHRASE_NONE = "xpub6BiVtCpG9fQPxnPmHXG8PhtzQdWC2Su4qWu6XW9tpWFYhxydCLJGrWBJZ5H6qTAHdPQ7pQhtpjiYZVZARo14qHiay2fvrX996oEP42u8wZy"
 XPUB_CARDANO_PASSPHRASE_B = "d80e770f6dfc3edb58eaab68aa091b2c27b08a47583471e93437ac5f8baa61880c7af4938a941c084c19731e6e57a5710e6ad1196263291aea297ce0eec0f177"
 
+ADDRESS_N = parse_path("44h/0h/0h")
+XPUB_REQUEST = messages.GetPublicKey(address_n=ADDRESS_N, coin_name="Bitcoin")
 
-def _get_xpub(client, passphrase):
+
-    """
+def _init_session(client, session_id=None):
-    Are calls are intentionally "raw", i.e. we are not using trezorlib
+    """Call Initialize, check and return the session ID."""
-    to avoid its implicit session management
+    response = client.call(messages.Initialize(session_id=session_id))
-    (clearing session after settings, hidden handling of session ids, etc.).
+    assert isinstance(response, messages.Features)
-    """
+    assert len(response.session_id) == 32
-    response = client.call_raw(
+    return response.session_id
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
+
-    )
+
+def _get_xpub(client, passphrase=None):
+    """Get XPUB and check that the appropriate passphrase flow has happened."""
+    response = client.call_raw(XPUB_REQUEST)
     if passphrase is not None:
         assert isinstance(response, messages.PassphraseRequest)
         response = client.call_raw(messages.PassphraseAck(passphrase=passphrase))
@@ -41,105 +46,142 @@ def _get_xpub(client, passphrase):
     return response.xpub
 
 
-def _enable_passphrase(client):
-    response = client.call_raw(messages.ApplySettings(use_passphrase=True))
-    assert isinstance(response, messages.ButtonRequest)  # confirm dialog
-    client.debug.press_yes()
-    response = client.call_raw(messages.ButtonAck())
-    assert isinstance(response, messages.Success)
-
-
 @pytest.mark.skip_ui
+@pytest.mark.setup_client(passphrase=True)
 def test_session_with_passphrase(client):
-    # Turn on passphrase.
-    _enable_passphrase(client)
-
     # Let's start the communication by calling Initialize.
-    response = client.call_raw(messages.Initialize())
+    session_id = _init_session(client)
-    assert isinstance(response, messages.Features)
-    session_id = response.session_id
-    assert len(session_id) == 32
 
     # GetPublicKey requires passphrase and since it is not cached,
     # Trezor will prompt for it.
-    xpub = _get_xpub(client, passphrase="A")
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
 
     # Call Initialize again, this time with the received session id and then call
     # GetPublicKey. The passphrase should be cached now so Trezor must
     # not ask for it again, whilst returning the same xpub.
-    response = client.call_raw(messages.Initialize(session_id=session_id))
+    new_session_id = _init_session(client, session_id=session_id)
-    assert isinstance(response, messages.Features)
+    assert new_session_id == session_id
-    xpub = _get_xpub(client, passphrase=None)
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
 
     # If we set session id in Initialize to None, the cache will be cleared
     # and Trezor will ask for the passphrase again.
-    response = client.call_raw(messages.Initialize(session_id=None))
+    new_session_id = _init_session(client)
-    assert isinstance(response, messages.Features)
+    assert new_session_id != session_id
-    xpub = _get_xpub(client, passphrase="A")
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
 
     # Unknown session id is the same as setting it to None.
-    response = client.call_raw(messages.Initialize(session_id=b"X" * 32))
+    _init_session(client, session_id=b"X" * 32)
-    assert isinstance(response, messages.Features)
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
-    xpub = _get_xpub(client, passphrase="A")
-    assert xpub == XPUB_PASSPHRASE_A
 
 
 @pytest.mark.skip_ui
 def test_session_enable_passphrase(client):
     # Let's start the communication by calling Initialize.
-    response = client.call_raw(messages.Initialize())
+    session_id = _init_session(client)
-    assert isinstance(response, messages.Features)
-    session_id = response.session_id
-    assert len(session_id) == 32
 
     # Trezor will not prompt for passphrase because it is turned off.
-    xpub = _get_xpub(client, passphrase=None)
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_NONE
-    assert xpub == XPUB_PASSPHRASE_NONE
 
     # Turn on passphrase.
-    _enable_passphrase(client)
+    # Emit the call explicitly to avoid ClearSession done by the library function
+    response = client.call(messages.ApplySettings(use_passphrase=True))
+    assert isinstance(response, messages.Success)
 
     # The session id is unchanged, therefore we do not prompt for the passphrase.
-    response = client.call_raw(messages.Initialize(session_id=session_id))
+    new_session_id = _init_session(client, session_id=session_id)
-    xpub = _get_xpub(client, passphrase=None)
+    assert session_id == new_session_id
-    assert isinstance(response, messages.Features)
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_NONE
-    assert session_id == response.session_id
-    assert xpub == XPUB_PASSPHRASE_NONE
 
     # We clear the session id now, so the passphrase should be asked.
-    response = client.call_raw(messages.Initialize())
+    new_session_id = _init_session(client)
-    xpub = _get_xpub(client, passphrase="A")
+    assert session_id != new_session_id
-    assert isinstance(response, messages.Features)
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
-    assert session_id != response.session_id
+
-    assert xpub == XPUB_PASSPHRASE_A
+
+@pytest.mark.skip_ui
+def test_clear_session(client):
+    session_id = _init_session(client)
+
+    # session id remains the same
+    new_session_id = _init_session(client, session_id=session_id)
+    assert session_id == new_session_id
+
+    # by clearing the session, the id is lost
+    response = client.call(messages.ClearSession())
+    assert isinstance(response, messages.Success)
+
+    # cannot resume the old session now
+    new_session_id = _init_session(client, session_id=session_id)
+    assert session_id != new_session_id
+
+
+@pytest.mark.skip_ui
+@pytest.mark.setup_client(passphrase=True)
+def test_clear_session_passphrase(client):
+    # at first attempt, we are prompted for passphrase
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
+
+    # now the passphrase is cached
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_A
+
+    # ClearSession will erase the cached passphrase
+    response = client.call(messages.ClearSession())
+    assert isinstance(response, messages.Success)
+
+    # we have to enter passphrase again
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
 
 
 @pytest.mark.skip_ui
 @pytest.mark.skip_t1
+@pytest.mark.setup_client(passphrase=True)
+def test_passphrase_on_device(client):
+    _init_session(client)
+
+    # try to get xpub with passphrase on host:
+    response = client.call_raw(XPUB_REQUEST)
+    assert isinstance(response, messages.PassphraseRequest)
+    response = client.call_raw(messages.PassphraseAck(passphrase="A", on_device=False))
+    assert isinstance(response, messages.PublicKey)
+    assert response.xpub == XPUB_PASSPHRASE_A
+
+    # try to get xpub again, passphrase should be cached
+    response = client.call_raw(XPUB_REQUEST)
+    assert isinstance(response, messages.PublicKey)
+    assert response.xpub == XPUB_PASSPHRASE_A
+
+    # make a new session
+    _init_session(client)
+
+    # try to get xpub with passphrase on device:
+    response = client.call_raw(XPUB_REQUEST)
+    assert isinstance(response, messages.PassphraseRequest)
+    response = client.call_raw(messages.PassphraseAck(on_device=True))
+    assert isinstance(response, messages.ButtonRequest)
+    client.debug.input("A")
+    response = client.call_raw(messages.ButtonAck())
+    assert isinstance(response, messages.PublicKey)
+    assert response.xpub == XPUB_PASSPHRASE_A
+
+    # try to get xpub again, passphrase should be cached
+    response = client.call_raw(XPUB_REQUEST)
+    assert isinstance(response, messages.PublicKey)
+    assert response.xpub == XPUB_PASSPHRASE_A
+
+
+@pytest.mark.skip_ui
+@pytest.mark.skip_t1
+@pytest.mark.setup_client(passphrase=True)
 def test_passphrase_always_on_device(client):
     # Let's start the communication by calling Initialize.
-    response = client.call_raw(messages.Initialize())
+    session_id = _init_session(client)
-    assert isinstance(response, messages.Features)
-    session_id = response.session_id
-
-    # Turn on passphrase.
-    _enable_passphrase(client)
 
     # Force passphrase entry on Trezor.
-    response = client.call_raw(messages.ApplySettings(passphrase_always_on_device=True))
+    response = client.call(messages.ApplySettings(passphrase_always_on_device=True))
-
-    assert isinstance(response, messages.ButtonRequest)  # confirm dialog
-    client.debug.press_yes()
-    response = client.call_raw(messages.ButtonAck())
     assert isinstance(response, messages.Success)
 
     # Since we enabled the always_on_device setting, Trezor will send ButtonRequests and ask for it on the device.
-    response = client.call_raw(
+    response = client.call_raw(XPUB_REQUEST)
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
-    )
     assert isinstance(response, messages.ButtonRequest)
     client.debug.input("")  # Input empty passphrase.
     response = client.call_raw(messages.ButtonAck())
@@ -147,20 +189,14 @@ def test_passphrase_always_on_device(client):
     assert response.xpub == XPUB_PASSPHRASE_NONE
 
     # Passphrase will not be prompted. The session id stays the same and the passphrase is cached.
-    response = client.call_raw(messages.Initialize(session_id=session_id))
+    _init_session(client, session_id=session_id)
-    assert isinstance(response, messages.Features)
+    response = client.call_raw(XPUB_REQUEST)
-    response = client.call_raw(
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
-    )
     assert isinstance(response, messages.PublicKey)
     assert response.xpub == XPUB_PASSPHRASE_NONE
 
     # In case we want to add a new passphrase we need to send session_id = None.
-    response = client.call_raw(messages.Initialize(session_id=None))
+    _init_session(client)
-    assert isinstance(response, messages.Features)
+    response = client.call_raw(XPUB_REQUEST)
-    response = client.call_raw(
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
-    )
     assert isinstance(response, messages.ButtonRequest)
     client.debug.input("A")  # Input empty passphrase.
     response = client.call_raw(messages.ButtonAck())
@@ -170,23 +206,15 @@ def test_passphrase_always_on_device(client):
 
 @pytest.mark.skip_ui
 @pytest.mark.skip_t2
+@pytest.mark.setup_client(passphrase=True)
 def test_passphrase_on_device_not_possible_on_t1(client):
-    # Let's start the communication by calling Initialize.
-    response = client.call_raw(messages.Initialize())
-    assert isinstance(response, messages.Features)
-
-    # Turn on passphrase.
-    _enable_passphrase(client)
-
     # This setting makes no sense on T1.
     response = client.call_raw(messages.ApplySettings(passphrase_always_on_device=True))
-
     assert isinstance(response, messages.Failure)
     assert response.code == FailureType.DataError
 
-    response = client.call_raw(
+    # T1 should not accept on_device request
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
+    response = client.call_raw(XPUB_REQUEST)
-    )
     assert isinstance(response, messages.PassphraseRequest)
     response = client.call_raw(messages.PassphraseAck(on_device=True))
     assert isinstance(response, messages.Failure)
@@ -196,9 +224,7 @@ def test_passphrase_on_device_not_possible_on_t1(client):
 @pytest.mark.skip_ui
 @pytest.mark.setup_client(passphrase=True)
 def test_passphrase_ack_mismatch(client):
-    response = client.call_raw(
+    response = client.call_raw(XPUB_REQUEST)
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
-    )
     assert isinstance(response, messages.PassphraseRequest)
     response = client.call_raw(messages.PassphraseAck(passphrase="A", on_device=True))
     assert isinstance(response, messages.Failure)
@@ -208,23 +234,29 @@ def test_passphrase_ack_mismatch(client):
 @pytest.mark.skip_ui
 @pytest.mark.setup_client(passphrase=True)
 def test_passphrase_missing(client):
-    response = client.call_raw(
+    response = client.call_raw(XPUB_REQUEST)
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
-    )
     assert isinstance(response, messages.PassphraseRequest)
     response = client.call_raw(messages.PassphraseAck(passphrase=None))
     assert isinstance(response, messages.Failure)
     assert response.code == FailureType.DataError
 
-    response = client.call_raw(
+    response = client.call_raw(XPUB_REQUEST)
-        messages.GetPublicKey(address_n=parse_path("44'/0'/0'"), coin_name="Bitcoin")
-    )
     assert isinstance(response, messages.PassphraseRequest)
     response = client.call_raw(messages.PassphraseAck(passphrase=None, on_device=False))
     assert isinstance(response, messages.Failure)
     assert response.code == FailureType.DataError
 
 
+def _get_xpub_cardano(client, passphrase):
+    msg = messages.CardanoGetPublicKey(address_n=parse_path("44'/1815'/0'/0/0"))
+    response = client.call_raw(msg)
+    if passphrase is not None:
+        assert isinstance(response, messages.PassphraseRequest)
+        response = client.call_raw(messages.PassphraseAck(passphrase=passphrase))
+    assert isinstance(response, messages.CardanoPublicKey)
+    return response.xpub
+
+
 @pytest.mark.skip_ui
 @pytest.mark.skip_t1
 @pytest.mark.altcoin
@@ -232,43 +264,34 @@ def test_passphrase_missing(client):
 def test_cardano_passphrase(client):
     # Cardano uses a variation of BIP-39 so we need to ask for the passphrase again.
 
-    response = client.call_raw(messages.Initialize())
+    session_id = _init_session(client)
-    assert isinstance(response, messages.Features)
-    session_id = response.session_id
-    assert len(session_id) == 32
 
     # GetPublicKey requires passphrase and since it is not cached,
     # Trezor will prompt for it.
-    xpub = _get_xpub(client, passphrase="A")
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
 
     # The passphrase is now cached for non-Cardano coins.
-    xpub = _get_xpub(client, passphrase=None)
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
 
     # Cardano will prompt for it again.
-    response = client.call_raw(
+    assert _get_xpub_cardano(client, passphrase="B") == XPUB_CARDANO_PASSPHRASE_B
-        messages.CardanoGetPublicKey(address_n=parse_path("44'/1815'/0'/0/0"))
-    )
-    assert isinstance(response, messages.PassphraseRequest)
-    response = client.call_raw(messages.PassphraseAck(passphrase="B"))
-    assert response.xpub == XPUB_CARDANO_PASSPHRASE_B
 
     # But now also Cardano has it cached.
-    response = client.call_raw(
+    assert _get_xpub_cardano(client, passphrase=None) == XPUB_CARDANO_PASSPHRASE_B
-        messages.CardanoGetPublicKey(address_n=parse_path("44'/1815'/0'/0/0"))
-    )
-    assert response.xpub == XPUB_CARDANO_PASSPHRASE_B
 
     # And others behaviour did not change.
-    xpub = _get_xpub(client, passphrase=None)
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
 
     # Initialize with the session id does not destroy the state
-    client.call_raw(messages.Initialize(session_id=session_id))
+    _init_session(client, session_id=session_id)
-    xpub = _get_xpub(client, passphrase=None)
+    assert _get_xpub(client, passphrase=None) == XPUB_PASSPHRASE_A
-    assert xpub == XPUB_PASSPHRASE_A
+    assert _get_xpub_cardano(client, passphrase=None) == XPUB_CARDANO_PASSPHRASE_B
-    response = client.call_raw(
+
-        messages.CardanoGetPublicKey(address_n=parse_path("44'/1815'/0'/0/0"))
+    # New session will destroy the state
-    )
+    _init_session(client)
-    assert response.xpub == XPUB_CARDANO_PASSPHRASE_B
+
+    # GetPublicKey must ask for passphrase again
+    assert _get_xpub(client, passphrase="A") == XPUB_PASSPHRASE_A
+
+    # Cardano must also ask for passphrase again
+    assert _get_xpub_cardano(client, passphrase="B") == XPUB_CARDANO_PASSPHRASE_B