1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2025-03-02 16:26:19 +00:00
trezor-firmware/python/src/trezorlib/transport/udp.py
matejcik 67b879ac07 python: restructure package
This deserves some explanation.

* tests were moved to separate python/tests subdir
* trezorlib was moved to python/src, so that it does not exist on
PYTHONPATH by default
(see https://blog.ionelmc.ro/2014/05/25/python-packaging/ for details)
* everything was updated to understand the new structure
* trezorctl was changed from a top-level executable script to a module
`trezorlib.cli.trezorctl` and is installed via the entry_points
mechanism.
This should make it work normally on Windows!

The package should be installable as normal through pip and pipenv, no
changes are needed on that side.

The source package from pypi will include unit tests. (Device tests were
completely moved out). Wheel will exclude them, because users don't need
them.
That shrinks the .whl from 520 kB to 270 - nice!

python: reorganize remaining unit tests
2019-08-12 12:57:25 +02:00

123 lines
4.0 KiB
Python

# This file is part of the Trezor project.
#
# Copyright (C) 2012-2019 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>.
import socket
from typing import Iterable, Optional, cast
from . import TransportException
from .protocol import ProtocolBasedTransport, get_protocol
class UdpTransport(ProtocolBasedTransport):
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 21324
PATH_PREFIX = "udp"
ENABLED = True
def __init__(self, device: str = None) -> None:
if not device:
host = UdpTransport.DEFAULT_HOST
port = UdpTransport.DEFAULT_PORT
else:
devparts = device.split(":")
host = devparts[0]
port = int(devparts[1]) if len(devparts) > 1 else UdpTransport.DEFAULT_PORT
self.device = (host, port)
self.socket = None # type: Optional[socket.socket]
protocol = get_protocol(self, want_v2=False)
super().__init__(protocol=protocol)
def get_path(self) -> str:
return "{}:{}:{}".format(self.PATH_PREFIX, *self.device)
def find_debug(self) -> "UdpTransport":
host, port = self.device
return UdpTransport("{}:{}".format(host, port + 1))
@classmethod
def _try_path(cls, path: str) -> "UdpTransport":
d = cls(path)
try:
d.open()
if d._ping():
return d
else:
raise TransportException(
"No Trezor device found at address {}".format(path)
)
finally:
d.close()
@classmethod
def enumerate(cls) -> Iterable["UdpTransport"]:
default_path = "{}:{}".format(cls.DEFAULT_HOST, cls.DEFAULT_PORT)
try:
return [cls._try_path(default_path)]
except TransportException:
return []
@classmethod
def find_by_path(cls, path: str, prefix_search: bool = False) -> "UdpTransport":
if prefix_search:
return cast(UdpTransport, super().find_by_path(path, prefix_search))
# This is *technically* type-able: mark `find_by_path` as returning
# the same type from which `cls` comes from.
# Mypy can't handle that though, so here we are.
else:
path = path.replace("{}:".format(cls.PATH_PREFIX), "")
return cls._try_path(path)
def open(self) -> None:
self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.socket.connect(self.device)
self.socket.settimeout(10)
def close(self) -> None:
if self.socket is not None:
self.socket.close()
self.socket = None
def _ping(self) -> bool:
"""Test if the device is listening."""
assert self.socket is not None
resp = None
try:
self.socket.sendall(b"PINGPING")
resp = self.socket.recv(8)
except Exception:
pass
return resp == b"PONGPONG"
def write_chunk(self, chunk: bytes) -> None:
assert self.socket is not None
if len(chunk) != 64:
raise TransportException("Unexpected data length")
self.socket.sendall(chunk)
def read_chunk(self) -> bytes:
assert self.socket is not None
while True:
try:
chunk = self.socket.recv(64)
break
except socket.timeout:
continue
if len(chunk) != 64:
raise TransportException("Unexpected chunk size: %d" % len(chunk))
return bytearray(chunk)