You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
python-trezor/trezorlib/transport_socket.py

111 lines
3.3 KiB
Python

'''SocketTransport implements TCP socket interface for Transport.'''
import socket
from select import select
from .transport import Transport
class SocketTransportClient(Transport):
def __init__(self, device, *args, **kwargs):
device = device.split(':')
if len(device) < 2:
device = ('0.0.0.0', int(device[0]))
else:
device = (device[0], int(device[1]))
self.socket = None
super(SocketTransportClient, self).__init__(device, *args, **kwargs)
def _open(self):
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.socket.connect(self.device)
self.filelike = self.socket.makefile()
def _close(self):
self.socket.close()
self.socket = None
self.filelike = None
def ready_to_read(self):
rlist, _, _ = select([self.socket], [], [], 0)
return len(rlist) > 0
def _write(self, msg, protobuf_msg):
self.socket.sendall(msg)
def _read(self):
try:
(msg_type, datalen) = self._read_headers(self.filelike)
return (msg_type, self.filelike.read(datalen))
except socket.error:
print("Failed to read from device")
return None
class SocketTransport(Transport):
def __init__(self, device, *args, **kwargs):
device = device.split(':')
if len(device) < 2:
device = ('0.0.0.0', int(device[0]))
else:
device = (device[0], int(device[1]))
self.socket = None
self.client = None
self.filelike = None
super(SocketTransport, self).__init__(device, *args, **kwargs)
def _open(self):
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#self.socket.setblocking(0)
self.socket.bind(self.device)
self.socket.listen(5)
def _disconnect_client(self):
print("Disconnecting client")
if self.client != None:
self.client.close()
self.client = None
self.filelike = None
def _close(self):
self._disconnect_client()
self.socket.close()
self.socket = None
def ready_to_read(self):
if self.filelike:
# Connected
rlist, _, _ = select([self.client], [], [], 0)
return len(rlist) > 0
else:
# Waiting for connection
rlist, _, _ = select([self.socket], [], [], 0)
if len(rlist) > 0:
(self.client, ipaddr) = self.socket.accept()
print("Connected", ipaddr[0])
self.filelike = self.client.makefile()
return self.ready_to_read()
return False
def _write(self, msg, protobuf_msg):
if self.filelike:
# None on disconnected client
try:
self.filelike.write(msg)
self.filelike.flush()
except socket.error:
print("Socket error")
self._disconnect_client()
def _read(self):
try:
(msg_type, datalen) = self._read_headers(self.filelike)
return (msg_type, self.filelike.read(datalen))
except Exception:
print("Failed to read from device")
self._disconnect_client()
return None