116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
Minimal metric server with custom binary protocol (cannibalised from Claude).
|
||
|
|
Demonstrates low-overhead sockets vs HTTP/REST.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import socket
|
||
|
|
import struct
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
import os
|
||
|
|
import random
|
||
|
|
import threading
|
||
|
|
|
||
|
|
# Protocol constants (shared with client)
|
||
|
|
MSG_PING = 0x01
|
||
|
|
MSG_PONG = 0x02
|
||
|
|
MSG_METRIC_REQ = 0x03
|
||
|
|
MSG_METRIC_RESP = 0x04
|
||
|
|
MSG_ERROR = 0xFF
|
||
|
|
|
||
|
|
HEADER_FMT = "!BH" # 1 byte type + 2 bytes length (big-endian)
|
||
|
|
HEADER_SIZE = struct.calcsize(HEADER_FMT)
|
||
|
|
|
||
|
|
|
||
|
|
def get_metric(name: str) -> float:
|
||
|
|
metrics = {
|
||
|
|
"cpu": lambda: random.uniform(5.0, 95.0),
|
||
|
|
"memory": lambda: random.uniform(30.0, 80.0),
|
||
|
|
"disk": lambda: random.uniform(40.0, 90.0),
|
||
|
|
"loadavg": lambda: os.getloadavg()[0] if hasattr(os, "getloadavg") else random.uniform(0.5, 4.0),
|
||
|
|
"uptime": lambda: float(int(time.time()) % 100000),
|
||
|
|
}
|
||
|
|
fn = metrics.get(name)
|
||
|
|
if fn is None:
|
||
|
|
raise KeyError(f"Unknown metric: {name}")
|
||
|
|
return fn()
|
||
|
|
|
||
|
|
|
||
|
|
def send_message(sock: socket.socket, msg_type: int, payload: bytes = b"") -> None:
|
||
|
|
header = struct.pack(HEADER_FMT, msg_type, len(payload))
|
||
|
|
sock.sendall(header + payload)
|
||
|
|
|
||
|
|
|
||
|
|
def recv_exact(sock: socket.socket, n: int) -> bytes:
|
||
|
|
buf = bytearray()
|
||
|
|
while len(buf) < n:
|
||
|
|
chunk = sock.recv(n - len(buf))
|
||
|
|
if not chunk:
|
||
|
|
return b""
|
||
|
|
buf.extend(chunk)
|
||
|
|
return bytes(buf)
|
||
|
|
|
||
|
|
|
||
|
|
def recv_message(sock: socket.socket) -> tuple[int, bytes]:
|
||
|
|
header = recv_exact(sock, HEADER_SIZE)
|
||
|
|
if not header:
|
||
|
|
raise ConnectionError("Client disconnected")
|
||
|
|
msg_type, length = struct.unpack(HEADER_FMT, header)
|
||
|
|
payload = recv_exact(sock, length) if length > 0 else b""
|
||
|
|
return msg_type, payload
|
||
|
|
|
||
|
|
|
||
|
|
def handle_client(conn: socket.socket, addr: tuple) -> None:
|
||
|
|
print(f"[+] Connected: {addr[0]}:{addr[1]}")
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
msg_type, payload = recv_message(conn)
|
||
|
|
|
||
|
|
if msg_type == MSG_PING:
|
||
|
|
print(f" <- PING from {addr[0]}")
|
||
|
|
send_message(conn, MSG_PONG)
|
||
|
|
|
||
|
|
elif msg_type == MSG_METRIC_REQ:
|
||
|
|
name = payload.decode("utf-8")
|
||
|
|
print(f" <- METRIC_REQ: {name}")
|
||
|
|
try:
|
||
|
|
value = get_metric(name)
|
||
|
|
send_message(conn, MSG_METRIC_RESP, struct.pack("!d", value))
|
||
|
|
except KeyError as e:
|
||
|
|
send_message(conn, MSG_ERROR, str(e).encode("utf-8"))
|
||
|
|
|
||
|
|
else:
|
||
|
|
send_message(conn, MSG_ERROR, f"Unknown type 0x{msg_type:02X}".encode("utf-8"))
|
||
|
|
|
||
|
|
except ConnectionError:
|
||
|
|
pass
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
print(f"[-] Disconnected: {addr[0]}:{addr[1]}")
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
host = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
|
||
|
|
port = int(sys.argv[2]) if len(sys.argv) > 2 else 9999
|
||
|
|
|
||
|
|
srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
|
|
srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
|
|
srv.bind((host, port))
|
||
|
|
srv.listen(5)
|
||
|
|
|
||
|
|
print(f"Metric server listening on {host}:{port}")
|
||
|
|
print("Available metrics: cpu, memory, disk, loadavg, uptime\n")
|
||
|
|
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
conn, addr = srv.accept()
|
||
|
|
threading.Thread(target=handle_client, args=(conn, addr), daemon=True).start()
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
print("\nShutting down.")
|
||
|
|
finally:
|
||
|
|
srv.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|