arno/trezor-firmware
1
0
Fork 0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-01 11:58:28 +00:00
Code Issues Releases Wiki Activity
79af52366d
trezor-firmware/python/src/trezorlib/client.py
M1nd3r d058648de8 fix(python): fix device tests for test_session 
[no changelog]
2024-11-28 13:59:17 +01:00

681 lines
26 KiB
Python
Raw Blame History

# This file is part of the Trezor project.
#
# Copyright (C) 2012-2022 SatoshiLabs and contributors
#
# This library is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# as published by the Free Software Foundation.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the License along with this library.
# If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
from __future__ import annotations
import logging
import os
import typing as t
from . import mapping, messages, models
from .mapping import ProtobufMapping
from .tools import parse_path
from .transport import Transport, get_transport
from .transport.thp.channel_data import ChannelData
from .transport.thp.protocol_and_channel import ProtocolAndChannel
from .transport.thp.protocol_v1 import ProtocolV1
from .transport.thp.protocol_v2 import ProtocolV2
if t.TYPE_CHECKING:
from .transport.session import Session
LOG = logging.getLogger(__name__)
MAX_PASSPHRASE_LENGTH = 50
MAX_PIN_LENGTH = 50
PASSPHRASE_ON_DEVICE = object()
PASSPHRASE_TEST_PATH = parse_path("44h/1h/0h/0/0")
OUTDATED_FIRMWARE_ERROR = """
Your Trezor firmware is out of date. Update it with the following command:
trezorctl firmware-update
Or visit https://suite.trezor.io/
""".strip()
LOG = logging.getLogger(__name__)
class TrezorClient:
button_callback: t.Callable[[Session, t.Any], t.Any] | None = None
pin_callback: t.Callable[[Session, t.Any], t.Any] | None = None
_management_session: Session | None = None
_features: messages.Features | None = None
def __init__(
self,
transport: Transport,
protobuf_mapping: ProtobufMapping | None = None,
protocol: ProtocolAndChannel | None = None,
) -> None:
self.transport = transport
if protobuf_mapping is None:
self.mapping = mapping.DEFAULT_MAPPING
else:
self.mapping = protobuf_mapping
if protocol is None:
try:
self.protocol = self._get_protocol()
except Exception as e:
print(e)
else:
self.protocol = protocol
self.protocol.mapping = self.mapping
@classmethod
def resume(
cls,
transport: Transport,
channel_data: ChannelData,
protobuf_mapping: ProtobufMapping | None = None,
) -> TrezorClient:
if protobuf_mapping is None:
protobuf_mapping = mapping.DEFAULT_MAPPING
protocol_v1 = ProtocolV1(transport, protobuf_mapping)
if channel_data.protocol_version == 2:
try:
protocol_v1.write(messages.Ping(message="Sanity check - to resume"))
except Exception as e:
print(type(e))
response = protocol_v1.read()
if (
isinstance(response, messages.Failure)
and response.code == messages.FailureType.InvalidProtocol
):
protocol = ProtocolV2(transport, protobuf_mapping, channel_data)
protocol.write(0, messages.Ping())
response = protocol.read(0)
if not isinstance(response, messages.Success):
LOG.debug("Failed to resume ProtocolV2")
raise Exception("Failed to resume ProtocolV2")
LOG.debug("Protocol V2 detected - can be resumed")
else:
LOG.debug("Failed to resume ProtocolV2")
raise Exception("Failed to resume ProtocolV2")
else:
protocol = ProtocolV1(transport, protobuf_mapping, channel_data)
return TrezorClient(transport, protobuf_mapping, protocol)
def get_session(
self,
passphrase: str | None = None,
derive_cardano: bool = False,
) -> Session:
"""
Returns initialized session (with derived seed).
Will fail if the device is not initialized
"""
from .transport.session import SessionV1, SessionV2
if isinstance(self.protocol, ProtocolV1):
if passphrase is None:
passphrase = ""
return SessionV1.new(self, passphrase, derive_cardano)
if isinstance(self.protocol, ProtocolV2):
return SessionV2.new(self, passphrase, derive_cardano)
raise NotImplementedError # TODO
def resume_session(self, session: Session):
"""
Note: this function potentially modifies the input session.
"""
from trezorlib.transport.session import SessionV1, SessionV2
from trezorlib.debuglink import SessionDebugWrapper
if isinstance(session, SessionDebugWrapper):
session = session._session
if isinstance(session, SessionV2):
return session
elif isinstance(session, SessionV1):
session.init_session()
return session
else:
raise NotImplementedError
def get_management_session(self, new_session: bool = False) -> Session:
from .transport.session import SessionV1, SessionV2
if not new_session and self._management_session is not None:
return self._management_session
if isinstance(self.protocol, ProtocolV1):
self._management_session = SessionV1.new(self, "", False)
elif isinstance(self.protocol, ProtocolV2):
self._management_session = SessionV2(self, b"\x00")
assert self._management_session is not None
return self._management_session
@property
def features(self) -> messages.Features:
if self._features is None:
self._features = self.protocol.get_features()
assert self._features is not None
return self._features
@property
def model(self) -> models.TrezorModel:
f = self.features
model = models.by_name(f.model or "1")
if model is None:
raise RuntimeError(
"Unsupported Trezor model"
f" (internal_model: {f.internal_model}, model: {f.model})"
)
return model
@property
def version(self) -> tuple[int, int, int]:
f = self.features
ver = (
f.major_version,
f.minor_version,
f.patch_version,
)
return ver
def refresh_features(self) -> None:
self.protocol.update_features()
self._features = self.protocol.get_features()
def ensure_unlocked(self) -> None:
# TODO implement
raise NotImplementedError
def _get_protocol(self) -> ProtocolAndChannel:
self.transport.open()
protocol = ProtocolV1(self.transport, mapping.DEFAULT_MAPPING)
protocol.write(messages.Initialize())
response = protocol.read()
self.transport.close()
if isinstance(response, messages.Failure):
if response.code == messages.FailureType.InvalidProtocol:
LOG.debug("Protocol V2 detected")
protocol = ProtocolV2(self.transport, self.mapping)
return protocol
def get_default_client(
path: t.Optional[str] = None,
**kwargs: t.Any,
) -> "TrezorClient":
"""Get a client for a connected Trezor device.
Returns a TrezorClient instance with minimum fuss.
If path is specified, does a prefix-search for the specified device. Otherwise, uses
the value of TREZOR_PATH env variable, or finds first connected Trezor.
If no UI is supplied, instantiates the default CLI UI.
"""
if path is None:
path = os.getenv("TREZOR_PATH")
transport = get_transport(path, prefix_search=True)
return TrezorClient(transport, **kwargs)
# class TrezorClient(t.Generic[UI]):
# """Trezor client, a connection to a Trezor device.
# This class allows you to manage connection state, send and receive protobuf
# messages, handle user interactions, and perform some generic tasks
# (send a cancel message, initialize or clear a session, ping the device).
# """
# model: models.TrezorModel
# transport: "Transport"
# session_id: t.Optional[bytes]
# ui: UI
# features: messages.Features
# def __init__(
# self,
# transport: "Transport",
# ui: UI,
# session_id: t.Optional[bytes] = None,
# derive_cardano: t.Optional[bool] = None,
# model: t.Optional[models.TrezorModel] = None,
# _init_device: bool = True,
# ) -> None:
# """Create a TrezorClient instance.
# You have to provide a `transport`, i.e., a raw connection to the device. You can
# use `trezorlib.transport.get_transport` to find one.
# You have to provide a UI implementation for the three kinds of interaction:
# - button request (notify the user that their interaction is needed)
# - PIN request (on T1, ask the user to input numbers for a PIN matrix)
# - passphrase request (ask the user to enter a passphrase) See `trezorlib.ui` for
# details.
# You can supply a `session_id` you might have saved in the previous session. If
# you do, the user might not need to enter their passphrase again.
# You can provide Trezor model information. If not provided, it is detected from
# the model name reported at initialization time.
# By default, the instance will open a connection to the Trezor device, send an
# `Initialize` message, set up the `features` field from the response, and connect
# to a session. By specifying `_init_device=False`, this step is skipped. Notably,
# this means that `client.features` is unset. Use `client.init_device()` or
# `client.refresh_features()` to fix that, otherwise A LOT OF THINGS will break.
# Only use this if you are _sure_ that you know what you are doing. This feature
# might be removed at any time.
# """
# LOG.info(f"creating client instance for device: {transport.get_path()}")
# # Here, self.model could be set to None. Unless _init_device is False, it will
# # get correctly reconfigured as part of the init_device flow.
# self.model = model # type: ignre ["None" is incompatible with "TrezorModel"]
# if self.model:
# self.mapping = self.model.default_mapping
# else:
# self.mapping = mapping.DEFAULT_MAPPING
# self.transport = transport
# self.ui = ui
# self.session_counter = 0
# self.session_id = session_id
# if _init_device:
# self.init_device(session_id=session_id, derive_cardano=derive_cardano)
# self.resume_session()
# def open(self) -> None:
# if self.session_counter == 0:
# session_id = self.transport.resume_session(b"")
# if self.session_id != session_id:
# print("Failed to resume session, allocated a new session")
# self.session_id = session_id
# self.transport.deprecated_begin_session()
# self.session_counter += 1
# def resume_session(self) -> None:
# new_id = self.transport.resume_session(self.session_id or b"")
# if self.session_id != new_id:
# print("Failed to resume session, allocated a new session")
# self.session_id = new_id
# def close(self) -> None:
# self.session_counter = max(self.session_counter - 1, 0)
# if self.session_counter == 0:
# # TODO call EndSession here?
# self.transport.deprecated_end_session()
# def cancel(self) -> None:
# self._raw_write(messages.Cancel())
# def call_raw(self, msg: "MessageType") -> "MessageType":
# __tracebackhide__ = True # for pytest # pylint: disable=W0612
# self._raw_write(msg)
# x = self._raw_read()
# return x
# def _raw_write(self, msg: "MessageType") -> None:
# __tracebackhide__ = True # for pytest # pylint: disable=W0612
# LOG.debug(
# f"sending message: {msg.__class__.__name__}",
# extra={"protobuf": msg},
# )
# msg_type, msg_bytes = self.mapping.encode(msg)
# LOG.log(
# DUMP_BYTES,
# f"encoded as type {msg_type} ({len(msg_bytes)} bytes): {msg_bytes.hex()}",
# )
# self.transport.write(msg_type, msg_bytes)
# def _raw_read(self) -> "MessageType":
# __tracebackhide__ = True # for pytest # pylint: disable=W0612
# msg_type, msg_bytes = self.transport.read()
# print("type/data", msg_type, msg_bytes)
# LOG.log(
# DUMP_BYTES,
# f"received type {msg_type} ({len(msg_bytes)} bytes): {msg_bytes.hex()}",
# )
# msg = self.mapping.decode(msg_type, msg_bytes)
# LOG.debug(
# f"received message: {msg.__class__.__name__}",
# extra={"protobuf": msg},
# )
# return msg
# def _callback_pin(self, msg: messages.PinMatrixRequest) -> "MessageType":
# try:
# pin = self.ui.get_pin(msg.type)
# except exceptions.Cancelled:
# self.call_raw(messages.Cancel())
# raise
# if any(d not in "123456789" for d in pin) or not (
# 1 <= len(pin) <= MAX_PIN_LENGTH
# ):
# self.call_raw(messages.Cancel())
# raise ValueError("Invalid PIN provided")
# resp = self.call_raw(messages.PinMatrixAck(pin=pin))
# if isinstance(resp, messages.Failure) and resp.code in (
# messages.FailureType.PinInvalid,
# messages.FailureType.PinCancelled,
# messages.FailureType.PinExpected,
# ):
# raise exceptions.PinException(resp.code, resp.message)
# else:
# return resp
# def _callback_passphrase(self, msg: messages.PassphraseRequest) -> "MessageType":
# available_on_device = Capability.PassphraseEntry in self.features.capabilities
# def send_passphrase(
# passphrase: t.Optional[str] = None, on_device: t.Optional[bool] = None
# ) -> "MessageType":
# msg = messages.PassphraseAck(passphrase=passphrase, on_device=on_device)
# resp = self.call_raw(msg)
# if isinstance(resp, messages.Deprecated_PassphraseStateRequest):
# self.session_id = resp.state
# resp = self.call_raw(messages.Deprecated_PassphraseStateAck())
# return resp
# # short-circuit old style entry
# if msg._on_device is True:
# return send_passphrase(None, None)
# try:
# passphrase = self.ui.get_passphrase(available_on_device=available_on_device)
# except exceptions.Cancelled:
# self.call_raw(messages.Cancel())
# raise
# if passphrase is PASSPHRASE_ON_DEVICE:
# if not available_on_device:
# self.call_raw(messages.Cancel())
# raise RuntimeError("Device is not capable of entering passphrase")
# else:
# return send_passphrase(on_device=True)
# # else process host-entered passphrase
# if not isinstance(passphrase, str):
# raise RuntimeError("Passphrase must be a str")
# passphrase = Mnemonic.normalize_string(passphrase)
# if len(passphrase) > MAX_PASSPHRASE_LENGTH:
# self.call_raw(messages.Cancel())
# raise ValueError("Passphrase too long")
# return send_passphrase(passphrase, on_device=False)
# def _callback_button(self, msg: messages.ButtonRequest) -> "MessageType":
# __tracebackhide__ = True # for pytest # pylint: disable=W0612
# # do this raw - send ButtonAck first, notify UI later
# self._raw_write(messages.ButtonAck())
# self.ui.button_request(msg)
# return self._raw_read()
# @session
# def call(self, msg: "MessageType") -> "MessageType":
# self.check_firmware_version()
# resp = self.call_raw(msg)
# while True:
# if isinstance(resp, messages.PinMatrixRequest):
# resp = self._callback_pin(resp)
# elif isinstance(resp, messages.PassphraseRequest):
# resp = self._callback_passphrase(resp)
# elif isinstance(resp, messages.ButtonRequest):
# resp = self._callback_button(resp)
# elif isinstance(resp, messages.Failure):
# print("self.call-failure")
# if resp.code == messages.FailureType.ActionCancelled:
# raise exceptions.Cancelled
# raise exceptions.TrezorFailure(resp)
# else:
# print("self.call-end")
# return resp
# def _refresh_features(self, features: messages.Features) -> None:
# """Update internal fields based on passed-in Features message."""
# if not self.model:
# # Trezor Model One bootloader 1.8.0 or older does not send model name
# model = models.by_internal_name(features.internal_model)
# if model is None:
# model = models.by_name(features.model or "1")
# if model is None:
# raise RuntimeError(
# "Unsupported Trezor model"
# f" (internal_model: {features.internal_model}, model: {features.model})"
# )
# self.model = model
# if features.vendor not in self.model.vendors:
# raise RuntimeError("Unsupported device")
# self.features = features
# self.version = (
# self.features.major_version,
# self.features.minor_version,
# self.features.patch_version,
# )
# self.check_firmware_version(warn_only=True)
# if self.features.session_id is not None:
# self.session_id = self.features.session_id
# self.features.session_id = None
# @session
# def refresh_features(self) -> messages.Features:
# """Reload features from the device.
# Should be called after changing settings or performing operations that affect
# device state.
# """
# resp = self.call_raw(messages.GetFeatures())
# if not isinstance(resp, messages.Features):
# raise exceptions.TrezorException("Unexpected response to GetFeatures")
# self._refresh_features(resp)
# return resp
# def init_device(
# self,
# *,
# session_id: t.Optional[bytes] = None,
# new_session: bool = False,
# derive_cardano: t.Optional[bool] = None,
# ) -> t.Optional[bytes]:
# """Initialize the device and return a session ID.
# You can optionally specify a session ID. If the session still exists on the
# device, the same session ID will be returned and the session is resumed.
# Otherwise a different session ID is returned.
# Specify `new_session=True` to open a fresh session. Since firmware version
# 1.9.0/2.3.0, the previous session will remain cached on the device, and can be
# resumed by calling `init_device` again with the appropriate session ID.
# If neither `new_session` nor `session_id` is specified, the current session ID
# will be reused. If no session ID was cached, a new session ID will be allocated
# and returned.
# # Version notes:
# Trezor One older than 1.9.0 does not have session management. Optional arguments
# have no effect and the function returns None
# Trezor T older than 2.3.0 does not have session cache. Requesting a new session
# will overwrite the old one. In addition, this function will always return None.
# A valid session_id can be obtained from the `session_id` attribute, but only
# after a passphrase-protected call is performed. You can use the following code:
# >>> client.init_device()
# >>> client.ensure_unlocked()
# >>> valid_session_id = client.session_id
# """
# if new_session:
# self.session_id = None
# elif session_id is not None:
# self.session_id = session_id
# print("before init conn")
# resp = self.transport.initialize_connection(
# mapping=self.mapping,
# session_id=session_id,
# derive_cardano=derive_cardano,
# )
# print("here")
# if isinstance(resp, messages.Failure):
# # can happen if `derive_cardano` does not match the current session
# raise exceptions.TrezorFailure(resp)
# if not isinstance(resp, messages.Features):
# raise exceptions.TrezorException("Unexpected response to Initialize")
# if self.session_id is not None and resp.session_id == self.session_id:
# LOG.info("Successfully resumed session")
# elif session_id is not None:
# LOG.info("Failed to resume session")
# # TT < 2.3.0 compatibility:
# # _refresh_features will clear out the session_id field. We want this function
# # to return its value, so that callers can rely on it being either a valid
# # session_id, or None if we can't do that.
# # Older TT FW does not report session_id in Features and self.session_id might
# # be invalid because TT will not allocate a session_id until a passphrase
# # exchange happens.
# reported_session_id = resp.session_id
# self._refresh_features(resp)
# print("there:", reported_session_id)
# return reported_session_id
# def is_outdated(self) -> bool:
# if self.features.bootloader_mode:
# return False
# return self.version < self.model.minimum_version
# def check_firmware_version(self, warn_only: bool = False) -> None:
# if self.is_outdated():
# if warn_only:
# warnings.warn("Firmware is out of date", stacklevel=2)
# else:
# raise exceptions.OutdatedFirmwareError(OUTDATED_FIRMWARE_ERROR)
# @expect(messages.Success, field="message", ret_type=str)
# def ping(
# self,
# msg: str,
# button_protection: bool = False,
# ) -> "MessageType":
# # We would like ping to work on any valid TrezorClient instance, but
# # due to the protection modes, we need to go through self.call, and that will
# # raise an exception if the firmware is too old.
# # So we short-circuit the simplest variant of ping with call_raw.
# if not button_protection:
# # XXX this should be: `with self:`
# try:
# self.open()
# resp = self.call_raw(messages.Ping(message=msg))
# if isinstance(resp, messages.ButtonRequest):
# # device is PIN-locked.
# # respond and hope for the best
# resp = self._callback_button(resp)
# return resp
# finally:
# self.close()
# return self.call(
# messages.Ping(message=msg, button_protection=button_protection)
# )
# def get_device_id(self) -> t.Optional[str]:
# return self.features.device_id
# @session
# def lock(self, *, _refresh_features: bool = True) -> None:
# """Lock the device.
# If the device does not have a PIN configured, this will do nothing.
# Otherwise, a lock screen will be shown and the device will prompt for PIN
# before further actions.
# This call does _not_ invalidate passphrase cache. If passphrase is in use,
# the device will not prompt for it after unlocking.
# To invalidate passphrase cache, use `end_session()`. To lock _and_ invalidate
# passphrase cache, use `clear_session()`.
# """
# # Private argument _refresh_features can be used internally to avoid
# # refreshing in cases where we will refresh soon anyway. This is used
# # in TrezorClient.clear_session()
# self.call(messages.LockDevice())
# if _refresh_features:
# self.refresh_features()
# @session
# def ensure_unlocked(self) -> None:
# """Ensure the device is unlocked and a passphrase is cached.
# If the device is locked, this will prompt for PIN. If passphrase is enabled
# and no passphrase is cached for the current session, the device will also
# prompt for passphrase.
# After calling this method, further actions on the device will not prompt for
# PIN or passphrase until the device is locked or the session becomes invalid.
# """
# from .btc import get_address
# get_address(self, "Testnet", PASSPHRASE_TEST_PATH)
# self.refresh_features()
# def end_session(self) -> None:
# """Close the current session and clear cached passphrase.
# The session will become invalid until `init_device()` is called again.
# If passphrase is enabled, further actions will prompt for it again.
# This is a no-op in bootloader mode, as it does not support session management.
# """
# # since: 2.3.4, 1.9.4
# print("end session")
# try:
# if not self.features.bootloader_mode:
# self.transport.end_session(self.session_id or b"")
# # self.call(messages.EndSession())
# except exceptions.TrezorFailure:
# # A failure most likely means that the FW version does not support
# # the EndSession call. We ignore the failure and clear the local session_id.
# # The client-side end result is identical.
# pass
# except ValueError as e:
# print(e)
# print(e.args)
# self.session_id = None
# @session
# def clear_session(self) -> None:
# """Lock the device and present a fresh session.
# The current session will be invalidated and a new one will be started. If the
# device has PIN enabled, it will become locked.
# Equivalent to calling `lock()`, `end_session()` and `init_device()`.
# """
# self.lock(_refresh_features=False)
# self.end_session()
# self.init_device(new_session=True)
Reference in New Issue View Git Blame Copy Permalink