229 lines
7.7 KiB
Python
229 lines
7.7 KiB
Python
"""A server for multiple devices that can be connected to the simulator."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import select
|
|
import signal
|
|
import socket
|
|
import sys
|
|
from threading import Event
|
|
from typing import Protocol
|
|
|
|
from sbot_interface.devices.util import get_globals
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
g = get_globals()
|
|
|
|
|
|
class Board(Protocol):
|
|
"""The interface for all board simulators that can be connected to the simulator."""
|
|
|
|
asset_tag: str
|
|
software_version: str
|
|
|
|
def handle_command(self, command: str) -> str | bytes:
|
|
"""
|
|
Process a command string and return the response.
|
|
|
|
Bytes type are treated as tag-length-value (TLV) encoded data.
|
|
"""
|
|
pass
|
|
|
|
|
|
class DeviceServer:
|
|
"""
|
|
A server for a single device that can be connected to the simulator.
|
|
|
|
The process_data method is called when data is received from the socket.
|
|
Line-delimited commands are processed and responses are sent back.
|
|
"""
|
|
|
|
def __init__(self, board: Board) -> None:
|
|
self.board = board
|
|
# create TCP socket server
|
|
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self.server_socket.bind(('127.0.0.1', 0))
|
|
self.server_socket.listen(1) # only allow one connection per device
|
|
self.server_socket.setblocking(True)
|
|
LOGGER.info(
|
|
f'Started server for {self.board_type} ({self.board.asset_tag}) '
|
|
f'on port {self.port}'
|
|
)
|
|
|
|
self.device_socket: socket.socket | None = None
|
|
self.buffer = b''
|
|
|
|
def process_data(self, data: bytes) -> bytes | None:
|
|
"""Process incoming data if a line has been received and return the response."""
|
|
self.buffer += data
|
|
if b'\n' in self.buffer:
|
|
# Sleep to simulate processing time
|
|
g.sleep(g.timestep / 1000)
|
|
data, self.buffer = self.buffer.split(b'\n', 1)
|
|
return self.run_command(data.decode().strip())
|
|
else:
|
|
return None
|
|
|
|
def run_command(self, command: str) -> bytes:
|
|
"""
|
|
Process a command and return the response.
|
|
|
|
Wraps the board's handle_command method and deals with exceptions and data types.
|
|
"""
|
|
LOGGER.debug(f'> {command}')
|
|
try:
|
|
response = self.board.handle_command(command)
|
|
if isinstance(response, bytes):
|
|
LOGGER.debug(f'< {len(response)} bytes')
|
|
return response
|
|
else:
|
|
LOGGER.debug(f'< {response}')
|
|
return response.encode() + b'\n'
|
|
except Exception as e:
|
|
LOGGER.exception(f'Error processing command: {command}')
|
|
return f'NACK:{e}\n'.encode()
|
|
|
|
def flush_buffer(self) -> None:
|
|
"""Clear the internal buffer of received data."""
|
|
self.buffer = b''
|
|
|
|
def socket(self) -> socket.socket:
|
|
"""
|
|
Return the socket to select on.
|
|
|
|
If the device is connected, return the device socket.
|
|
Otherwise, return the server socket.
|
|
"""
|
|
if self.device_socket is not None:
|
|
# ignore the server socket while we are connected
|
|
return self.device_socket
|
|
else:
|
|
return self.server_socket
|
|
|
|
def accept(self) -> None:
|
|
"""Accept a connection from a device and set the device socket to blocking."""
|
|
if self.device_socket is not None:
|
|
self.disconnect_device()
|
|
self.device_socket, _ = self.server_socket.accept()
|
|
self.device_socket.setblocking(True)
|
|
LOGGER.info(f'Connected to {self.asset_tag} from {self.device_socket.getpeername()}')
|
|
|
|
def disconnect_device(self) -> None:
|
|
"""Close the device socket, flushing the buffer first."""
|
|
self.flush_buffer()
|
|
if self.device_socket is not None:
|
|
self.device_socket.close()
|
|
self.device_socket = None
|
|
LOGGER.info(f'Disconnected from {self.asset_tag}')
|
|
|
|
def close(self) -> None:
|
|
"""Close the server and client sockets."""
|
|
self.disconnect_device()
|
|
self.server_socket.close()
|
|
|
|
def __del__(self) -> None:
|
|
self.close()
|
|
|
|
@property
|
|
def port(self) -> int:
|
|
"""Return the port number of the server socket."""
|
|
if self.server_socket is None:
|
|
return -1
|
|
return int(self.server_socket.getsockname()[1])
|
|
|
|
@property
|
|
def asset_tag(self) -> str:
|
|
"""Return the asset tag of the board."""
|
|
return self.board.asset_tag
|
|
|
|
@property
|
|
def board_type(self) -> str:
|
|
"""Return the class name of the board object."""
|
|
return self.board.__class__.__name__
|
|
|
|
|
|
class SocketServer:
|
|
"""
|
|
A server for multiple devices that can be connected to the simulator.
|
|
|
|
The run method blocks until the stop_event is set.
|
|
"""
|
|
|
|
def __init__(self, devices: list[DeviceServer]) -> None:
|
|
self.devices = devices
|
|
self.stop_event = Event()
|
|
g.stop_event = self.stop_event
|
|
# flag to indicate that we are exiting because the usercode has completed
|
|
self.completed = False
|
|
|
|
def run(self) -> None:
|
|
"""
|
|
Run the server, accepting connections and processing data.
|
|
|
|
This method blocks until the stop_event is set.
|
|
"""
|
|
while not self.stop_event.is_set():
|
|
# select on all server sockets and device sockets
|
|
sockets = [device.socket() for device in self.devices]
|
|
|
|
readable, _, _ = select.select(sockets, [], [], 0.5)
|
|
|
|
for device in self.devices:
|
|
try:
|
|
if device.server_socket in readable:
|
|
device.accept()
|
|
|
|
if device.device_socket in readable and device.device_socket is not None:
|
|
try:
|
|
if sys.platform == 'win32':
|
|
data = device.device_socket.recv(4096)
|
|
else:
|
|
data = device.device_socket.recv(4096, socket.MSG_DONTWAIT)
|
|
except ConnectionError:
|
|
device.disconnect_device()
|
|
continue
|
|
|
|
if not data:
|
|
device.disconnect_device()
|
|
else:
|
|
response = device.process_data(data)
|
|
if response is not None:
|
|
try:
|
|
device.device_socket.sendall(response)
|
|
except ConnectionError:
|
|
device.disconnect_device()
|
|
continue
|
|
except Exception as e:
|
|
LOGGER.exception(f"Failure in simulated boards: {e}")
|
|
|
|
LOGGER.info('Stopping server')
|
|
for device in self.devices:
|
|
device.close()
|
|
|
|
if self.stop_event.is_set() and self.completed is False:
|
|
# Stop the usercode
|
|
os.kill(os.getpid(), signal.SIGINT)
|
|
|
|
def links(self) -> dict[str, dict[str, str]]:
|
|
"""Return a mapping of asset tags to ports, grouped by board type."""
|
|
return {
|
|
device.asset_tag: {
|
|
'board_type': device.board_type,
|
|
'port': str(device.port),
|
|
}
|
|
for device in self.devices
|
|
}
|
|
|
|
def links_formatted(self, address: str = '127.0.0.1') -> str:
|
|
"""
|
|
Return a formatted string of all the links to the devices.
|
|
|
|
The format is 'socket://address:port/board_type/asset_tag'.
|
|
Each link is separated by a newline.
|
|
"""
|
|
return '\n'.join(
|
|
f"socket://{address}:{data['port']}/{data['board_type']}/{asset_tag}"
|
|
for asset_tag, data in self.links().items()
|
|
)
|