mirror of
https://github.com/trezor/trezor-firmware.git
synced 2024-11-17 21:22:10 +00:00
111 lines
3.3 KiB
Python
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
|