refactor: wip
This commit is contained in:
@@ -7,7 +7,7 @@ import socket
|
||||
import threading
|
||||
import time
|
||||
|
||||
from judas_protocol import Message
|
||||
from judas_protocol import Category, ControlAction, Message
|
||||
|
||||
from judas_server.backend.client import Client
|
||||
|
||||
@@ -60,13 +60,14 @@ class BackendServer:
|
||||
)
|
||||
time.sleep(1)
|
||||
|
||||
def _send_ack(self, client: Client) -> None:
|
||||
def _send_ack(self, client: Client, target_id: str) -> None:
|
||||
"""Send an ACK message to a client.
|
||||
|
||||
Args:
|
||||
client (Client): The client to send the ACK to.
|
||||
target_id (str): The id of the ACK'd message.
|
||||
"""
|
||||
ack: bytes = Message.ack().to_bytes()
|
||||
ack: bytes = Message.ack(target_id=target_id).to_bytes()
|
||||
self.logger.debug(f"[>] Sending ACK to {client}")
|
||||
client.outbound += ack
|
||||
|
||||
@@ -78,43 +79,13 @@ class BackendServer:
|
||||
"""
|
||||
conn, addr = sock.accept()
|
||||
self.logger.info(f"[+] Accepted connection from {addr}")
|
||||
conn.setblocking(False)
|
||||
|
||||
# wait for hello message to get mac_id
|
||||
|
||||
conn.settimeout(5)
|
||||
try:
|
||||
message = conn.recv(1024)
|
||||
if not message:
|
||||
self.logger.error(f"[-] No data received from {addr}")
|
||||
conn.close()
|
||||
return
|
||||
except socket.timeout:
|
||||
self.logger.error(f"[-] Timeout waiting for hello from {addr}")
|
||||
conn.close()
|
||||
return
|
||||
conn.settimeout(None)
|
||||
|
||||
message = message.split(b"\n")[0] # get first line only
|
||||
message = Message.from_bytes(message)
|
||||
|
||||
mac_id = message.payload.get("mac", None)
|
||||
if mac_id is None:
|
||||
self.logger.error(
|
||||
f"[-] No mac_id provided by {addr}, closing connection"
|
||||
)
|
||||
conn.close()
|
||||
return
|
||||
|
||||
client = Client(id_=mac_id, addr=addr, socket=conn)
|
||||
self.clients[mac_id] = client
|
||||
|
||||
self._send_ack(client)
|
||||
client = Client(mac_id=None, addr=addr, socket=conn)
|
||||
|
||||
events = selectors.EVENT_READ | selectors.EVENT_WRITE
|
||||
self.selector.register(conn, events, data=client)
|
||||
|
||||
self.logger.info(f"[+] Registered client {client}")
|
||||
self.logger.info(f"[+] Registered client {client}, HELLO pending...")
|
||||
|
||||
def _disconnect(self, client: Client) -> None:
|
||||
"""Disconnect a client and clean up resources.
|
||||
@@ -122,13 +93,54 @@ class BackendServer:
|
||||
Args:
|
||||
sock (socket.socket): The client socket to disconnect.
|
||||
"""
|
||||
self.logger.info(f"[-] Disconnecting {client}")
|
||||
self.logger.info(f"[-] Disconnecting {client}...")
|
||||
self.logger.debug("[*] Sending DNR message...")
|
||||
|
||||
try:
|
||||
self.selector.unregister(client.socket)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error unregistering client {client}: {e}")
|
||||
|
||||
client.disconnect()
|
||||
|
||||
def _send_outbound(
|
||||
self, sock: socket.socket, client: Client, data: bytes
|
||||
) -> None:
|
||||
"""Queue data to be sent to a client.
|
||||
|
||||
Args:
|
||||
client (Client): The client to send data to.
|
||||
data (bytes): The data to send.
|
||||
"""
|
||||
self.logger.debug(f"[>] Sending data to {client}: {client.outbound!r}")
|
||||
sent = sock.send(client.outbound)
|
||||
|
||||
client.outbound = client.outbound[sent:]
|
||||
|
||||
def _receive_inbound(
|
||||
self, sock: socket.socket, client: Client, packet_size: int = 4096
|
||||
) -> None:
|
||||
"""Receive data from a client socket.
|
||||
|
||||
Args:
|
||||
sock (socket.socket): The client socket to receive data from.
|
||||
client (Client): The client object.
|
||||
packet_size (int): The maximum amount of data to be received at once.
|
||||
Returns:
|
||||
bytes: The received data.
|
||||
"""
|
||||
recv_data = sock.recv(1024)
|
||||
if recv_data:
|
||||
self.logger.debug(
|
||||
f"[<] Received data from {client}: {recv_data!r}"
|
||||
)
|
||||
client.inbound += recv_data
|
||||
|
||||
# set last seen
|
||||
client.last_seen = time.time()
|
||||
else:
|
||||
self._disconnect(client)
|
||||
|
||||
def _handle_connection(
|
||||
self, key: selectors.SelectorKey, mask: int
|
||||
) -> None:
|
||||
@@ -138,39 +150,76 @@ class BackendServer:
|
||||
key (selectors.SelectorKey): The selector key for the client.
|
||||
mask (int): The event mask.
|
||||
"""
|
||||
sock: socket.socket = key.fileobj
|
||||
sock: socket.socket = key.fileobj # type: ignore
|
||||
client = key.data
|
||||
|
||||
try:
|
||||
if mask & selectors.EVENT_READ:
|
||||
recv_data = sock.recv(1024)
|
||||
if recv_data:
|
||||
self.logger.debug(
|
||||
f"[<] Received data from {client}: {recv_data!r}"
|
||||
)
|
||||
client.inbound += recv_data
|
||||
self._receive_inbound(sock, client)
|
||||
if client.inbound:
|
||||
if client.mac_id is None:
|
||||
# expect HELLO message
|
||||
try:
|
||||
msg = Message.from_bytes(client.inbound)
|
||||
if (
|
||||
msg.category == Category.CONTROL
|
||||
and msg.action == ControlAction.HELLO
|
||||
and msg.payload.get("mac") is not None
|
||||
):
|
||||
client.mac_id = msg.payload["mac"]
|
||||
if (
|
||||
client.mac_id in self.clients
|
||||
and self.clients[client.mac_id].status
|
||||
== "connected"
|
||||
):
|
||||
old_client: Client = self.clients[
|
||||
client.mac_id
|
||||
]
|
||||
self.logger.warning(
|
||||
f"Client {client.mac_id} is already connected from {old_client.addr}, disconnecting old client..."
|
||||
)
|
||||
self._disconnect(old_client)
|
||||
# TODO: tell client not to reconnect
|
||||
self.clients[client.mac_id] = client
|
||||
self.logger.info(
|
||||
f"[+] Registered new client {client}"
|
||||
)
|
||||
else:
|
||||
self.logger.error(
|
||||
f"Expected HELLO message from {client}, got {msg}"
|
||||
)
|
||||
self._disconnect(client)
|
||||
return
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to parse HELLO message from {client}: {e}"
|
||||
)
|
||||
self._disconnect(client)
|
||||
return
|
||||
while b"\n" in client.inbound:
|
||||
line, client.inbound = client.inbound.split(b"\n", 1)
|
||||
self.logger.info(
|
||||
f"[<] Complete message from {client}: {line!r}"
|
||||
)
|
||||
try:
|
||||
msg = Message.from_bytes(line)
|
||||
self.logger.info(f"[.] Parsed message {msg.id}")
|
||||
if msg.ack_required:
|
||||
self._send_ack(client, target_id=msg.id)
|
||||
|
||||
self._send_ack(client)
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to parse message from {client}: {e}"
|
||||
)
|
||||
self._disconnect(client)
|
||||
return
|
||||
|
||||
# set last seen
|
||||
client.last_seen = time.time()
|
||||
else:
|
||||
self._disconnect(client)
|
||||
|
||||
if mask & selectors.EVENT_WRITE:
|
||||
if client.outbound:
|
||||
self.logger.debug(
|
||||
f"[>] Sending data to {client}: {client.outbound!r}"
|
||||
)
|
||||
sent = sock.send(client.outbound)
|
||||
|
||||
client.outbound = client.outbound[sent:]
|
||||
# TODO: wait for ACK from client
|
||||
self._send_outbound(sock, client, client.outbound)
|
||||
|
||||
except ConnectionResetError as e:
|
||||
self.logger.error(f"Connection reset by {client}, disconnect: {e}")
|
||||
@@ -213,7 +262,7 @@ class BackendServer:
|
||||
self.logger.warning(f"Client {client_id} not found")
|
||||
return None
|
||||
return {
|
||||
"id": client.id,
|
||||
"id": client.mac_id,
|
||||
"addr": client.addr,
|
||||
"last_seen": client.last_seen,
|
||||
"status": client.status,
|
||||
|
||||
@@ -17,12 +17,13 @@ class Client:
|
||||
"""Represents a client."""
|
||||
|
||||
def __init__(
|
||||
self, id_: str, addr: tuple[str, int], socket: socket.socket
|
||||
self, mac_id: str | None, addr: tuple[str, int], socket: socket.socket
|
||||
) -> None:
|
||||
"""Initialize the client.
|
||||
|
||||
Args:
|
||||
id_ (str): The unique identifier for the client.
|
||||
mac_id (str | None): The unique identifier for the client.
|
||||
Can be None if not yet assigned.
|
||||
addr (tuple[str, int]): The (IP, port) address of the client.
|
||||
socket (socket.socket): The socket object for communication.
|
||||
"""
|
||||
@@ -31,7 +32,7 @@ class Client:
|
||||
)
|
||||
self.logger.debug(f"Initializing Client {addr}...")
|
||||
|
||||
self.id: str = id_
|
||||
self.mac_id: str | None = mac_id
|
||||
self.last_seen: float = 0.0 # unix timestanp of last inbound message
|
||||
self.status: ClientStatus = ClientStatus.CONNECTED
|
||||
|
||||
@@ -41,10 +42,10 @@ class Client:
|
||||
self.outbound: bytes = b""
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"Client({self.id} ({self.addr[0]}:{self.addr[1]}))"
|
||||
return f"Client({self.mac_id} ({self.addr[0]}:{self.addr[1]}))"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Client({self.id}, {self.addr})"
|
||||
return f"Client({self.mac_id}, {self.addr})"
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Disconnect the client and close the socket."""
|
||||
|
||||
Reference in New Issue
Block a user