81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
'''UDP Socket implementation of Transport.'''
|
|
|
|
import socket
|
|
from select import select
|
|
import time
|
|
from .transport import Transport, ConnectionError
|
|
|
|
class FakeRead(object):
|
|
# Let's pretend we have a file-like interface
|
|
def __init__(self, func):
|
|
self.func = func
|
|
|
|
def read(self, size):
|
|
return self.func(size)
|
|
|
|
class UdpTransport(Transport):
|
|
def __init__(self, device, *args, **kwargs):
|
|
self.buffer = ''
|
|
|
|
device = device.split(':')
|
|
if len(device) < 2:
|
|
if not device[0]:
|
|
# Default port used by trezor v2
|
|
device = ('127.0.0.1', 21324)
|
|
else:
|
|
device = ('127.0.0.1', int(device[0]))
|
|
else:
|
|
device = (device[0], int(device[1]))
|
|
|
|
self.socket = None
|
|
super(UdpTransport, self).__init__(device, *args, **kwargs)
|
|
|
|
def _open(self):
|
|
self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self.socket.connect(self.device)
|
|
|
|
def _close(self):
|
|
self.socket.close()
|
|
self.socket = None
|
|
self.buffer = ''
|
|
|
|
def ready_to_read(self):
|
|
rlist, _, _ = select([self.socket], [], [], 0)
|
|
return len(rlist) > 0
|
|
|
|
def _write_chunk(self, chunk):
|
|
if len(chunk) != 64:
|
|
raise Exception("Unexpected data length")
|
|
|
|
self.socket.sendall(chunk)
|
|
|
|
def _write(self, msg, protobuf_msg):
|
|
raise NotImplemented()
|
|
|
|
def _read(self):
|
|
(session_id, msg_type, datalen) = self._read_headers(FakeRead(self._raw_read))
|
|
return (session_id, msg_type, self._raw_read(datalen))
|
|
|
|
def _raw_read(self, length):
|
|
start = time.time()
|
|
while len(self.buffer) < length:
|
|
data = self.socket.recv(64)
|
|
if not len(data):
|
|
if time.time() - start > 10:
|
|
# Over 10 s of no response, let's check if
|
|
# device is still alive
|
|
if not self.is_connected():
|
|
raise ConnectionError("Connection failed")
|
|
else:
|
|
# Restart timer
|
|
start = time.time()
|
|
|
|
time.sleep(0.001)
|
|
continue
|
|
|
|
self.buffer += data
|
|
|
|
ret = self.buffer[:length]
|
|
self.buffer = self.buffer[length:]
|
|
return ret
|