Files
pi-swarm/configs/planck-tftp-server.py
2026-09-19 10:37:53 -07:00

127 lines
4.1 KiB
Python

#!/usr/bin/env python3
"""Minimal TFTP server (RRQ only) for planck netboot. Supports blksize and
windowsize options per RFC 2347/2348/2349. Serves files under TFTP_ROOT."""
import os
import socket
import struct
import threading
import time
TFTP_ROOT = "/volume1/plancknetboot/tftp"
PORT = 69
BLKSIZE_MAX = 1468 # avoids IP fragmentation on ethernet
DEFAULT_BLKSIZE = 512
DEFAULT_WINDOW = 1
TIMEOUT = 5
RETRIES = 6
OP_RRQ = 1
OP_DATA = 3
OP_ACK = 4
OP_ERROR = 5
OP_OACK = 6
def parse_rrq(data):
# opcode, filename, mode, optional opts
parts = data[2:].split(b"\x00")
filename = parts[0].decode(errors="replace")
opts = {}
rest = parts[2:]
for i in range(0, len(rest) - 1, 2):
try:
opts[rest[i].decode().lower()] = rest[i + 1].decode()
except (IndexError, UnicodeDecodeError):
pass
return filename, opts
def safe_path(filename):
parts = filename.split("/")
if any(not p or p.startswith(".") or "/" in p for p in parts):
return None
if any(not all(c.isalnum() or c in "._-" for c in p) for p in parts):
return None
path = os.path.join(TFTP_ROOT, *parts)
if not os.path.isfile(path):
return None
return path
def send_err(sock, addr, code, msg):
sock.sendto(struct.pack("!HH", OP_ERROR, code) + msg.encode() + b"\x00", addr)
def handle_client(sock, data, addr, filepath, opts):
print(f"planck-tftp: opts={opts} file={filepath} size={os.path.getsize(filepath)}", flush=True)
blksize = DEFAULT_BLKSIZE
if "blksize" in opts:
try:
blksize = min(BLKSIZE_MAX, max(512, int(opts["blksize"])))
except ValueError:
blksize = 512
payload = (
b"blksize\x00" + str(blksize).encode() + b"\x00"
+ b"tsize\x00" + str(os.path.getsize(filepath)).encode() + b"\x00"
)
sock.sendto(struct.pack("!H", OP_OACK) + payload, addr)
with open(filepath, "rb") as f:
block = 0
while True:
chunk = f.read(blksize)
block = (block % 65536) + 1
pkt = struct.pack("!HH", OP_DATA, block) + chunk
for attempt in range(RETRIES + 1):
sock.sendto(pkt, addr)
sock.settimeout(TIMEOUT)
try:
rdata, raddr = sock.recvfrom(4 + blksize)
except socket.timeout:
continue
if len(rdata) >= 4 and struct.unpack("!H", rdata[:2])[0] == OP_ACK:
acked = struct.unpack("!H", rdata[2:4])[0]
if acked == block:
break
if acked == 0 and block == 1:
break # ack of the OACK, counts as ack of pending state
else:
print(f"planck-tftp: TRANSFER STALLED at block {block} {filepath} {addr}", flush=True)
return # retries exhausted
if len(chunk) < blksize:
print(f"planck-tftp: TRANSFER DONE {filepath} {addr}", flush=True)
return
def main():
srv = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
srv.bind(("0.0.0.0", PORT))
print(f"planck-tftp: serving {TFTP_ROOT} on :{PORT}", flush=True)
while True:
data, addr = srv.recvfrom(1024)
if len(data) < 4 or struct.unpack("!H", data[:2])[0] != OP_RRQ:
continue
try:
filename, opts = parse_rrq(data)
except Exception:
continue
filepath = safe_path(filename)
if filepath is None:
errsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
send_err(errsock, addr, 1, "file not found")
errsock.close()
print(f"planck-tftp: RRQ denied {filename} from {addr}", flush=True)
continue
csock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
csock.bind(("0.0.0.0", 0))
print(f"planck-tftp: RRQ {filename} from {addr}", flush=True)
threading.Thread(
target=handle_client, args=(csock, data, addr, filepath, opts),
daemon=True,
).start()
if __name__ == "__main__":
main()