95 lines
4.0 KiB
Python
95 lines
4.0 KiB
Python
from dataclasses import dataclass
|
|
from random import randint
|
|
from socket import socket
|
|
import struct
|
|
|
|
MAX_OUTGOING_PAYLOAD_LENGTH = 1446
|
|
MAX_INCOMING_PAYLOAD_LENGTH = 4096
|
|
INVALID_PACKET_REQUEST_ID = 32767
|
|
|
|
# can be overridden to make send_command echo all commands before sending them; useful for troubleshooting
|
|
ECHO_COMMANDS = False
|
|
|
|
# packet types
|
|
PKT_LOGIN = 3
|
|
PKT_LOGIN_RESPONSE = 2
|
|
PKT_COMMAND = 2
|
|
PKT_COMMAND_RESPONSE = 0
|
|
|
|
@dataclass
|
|
class Packet:
|
|
"""A single RCON packet, either sent or received."""
|
|
packet_size: int
|
|
request_id: int
|
|
packet_type: int
|
|
payload: bytes
|
|
|
|
@staticmethod
|
|
def recv(sock: socket) -> Packet:
|
|
"""Read and parse the next full packet from sock, blocking until complete."""
|
|
# first four bytes (little-endian) contain the size of the remainder of the packet
|
|
[packet_size] = struct.unpack('<i', sock.recv(4))
|
|
# then, 4 bytes each for ID & type, arbitrary-length string payload, and two null bytes
|
|
request_id, packet_type, payload, null = struct.unpack(f'<ii{packet_size-4-4-2}sh',
|
|
Packet._recvall(sock, packet_size))
|
|
|
|
return Packet(packet_size, request_id, packet_type, payload)
|
|
|
|
@staticmethod
|
|
def _recvall(sock: socket, n: int) -> bytes:
|
|
"""Block until exactly n bytes have been read from sock."""
|
|
buf = b''
|
|
while len(buf) < n:
|
|
chunk = sock.recv(n - len(buf))
|
|
if not chunk:
|
|
raise ConnectionError("socket closed before full packet received")
|
|
buf += chunk
|
|
return buf
|
|
|
|
@staticmethod
|
|
def send(sock: socket, packet_type: int, payload: str, request_id: int | None = None):
|
|
"""Build a packet from payload and transmit it on sock.
|
|
|
|
Raises ValueError if the encoded payload exceeds MAX_OUTGOING_PAYLOAD_LENGTH.
|
|
"""
|
|
payload_bytes = bytes(payload, encoding='utf-8')
|
|
|
|
if len(payload_bytes) > MAX_OUTGOING_PAYLOAD_LENGTH:
|
|
raise ValueError(f"attempted to send payload of length {len(payload_bytes)} bytes, which exceeds the maximum of {MAX_OUTGOING_PAYLOAD_LENGTH} bytes")
|
|
|
|
# 4 bytes each for ID & type, arbitrary-length string payload, and two null bytes
|
|
packet_size = 4 + 4 + len(payload_bytes) + 2
|
|
|
|
if request_id == None:
|
|
request_id = randint(1, 32766)
|
|
|
|
data = struct.pack(f'<iii{len(payload_bytes)}sh', packet_size, request_id, packet_type, payload_bytes, 0)
|
|
sock.sendall(data)
|
|
|
|
|
|
def login(sock: socket, password: str) -> bool:
|
|
"""Authenticate on an already-connected socket. Returns a bool representing whether login succeeded."""
|
|
Packet.send(sock, PKT_LOGIN, password)
|
|
response = Packet.recv(sock)
|
|
|
|
if response.request_id == -1:
|
|
return False # auth failure
|
|
return True
|
|
|
|
def send_command(sock: socket, command: str) -> str:
|
|
"""Run command on an already-authenticated socket and return the decoded response."""
|
|
# Some commands can have responses greater than the max response length of 4096, so the server will send the response fragmented across several packets. However, the Notchian server doesn't properly implement that fragmentation logic, meaning there can be ambiguity about when the actual last packet is.
|
|
# To handle this, we send two packets with distinct IDs. The first packet has the command we want to run; the second has an invalid packet type. Since the server responds to packets in the order they were received, we read from the socket until we get a response that matches the second packet; everything before was part of the first packet's response.
|
|
if ECHO_COMMANDS: print(command)
|
|
Packet.send(sock, PKT_COMMAND, command)
|
|
Packet.send(sock, 200, "", request_id=INVALID_PACKET_REQUEST_ID)
|
|
|
|
response = b''
|
|
while True:
|
|
received = Packet.recv(sock)
|
|
if received.request_id == INVALID_PACKET_REQUEST_ID: break # we got a response from our second packet
|
|
response += received.payload
|
|
|
|
return response.decode('utf-8', errors='replace')
|
|
|