151 lines
4.8 KiB
Python
151 lines
4.8 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):
|
|
blksize = DEFAULT_BLKSIZE
|
|
window = DEFAULT_WINDOW
|
|
requested = {}
|
|
if "blksize" in opts:
|
|
try:
|
|
blksize = min(BLKSIZE_MAX, max(512, int(opts["blksize"])))
|
|
requested["blksize"] = str(blksize)
|
|
except ValueError:
|
|
pass
|
|
if "windowsize" in opts:
|
|
try:
|
|
window = min(64, max(1, int(opts["windowsize"])))
|
|
requested["windowsize"] = str(window)
|
|
except ValueError:
|
|
pass
|
|
|
|
size = os.path.getsize(filepath)
|
|
if requested:
|
|
payload = b"\x00".join(
|
|
k.encode() + b"\x00" + v.encode() for k, v in requested.items()
|
|
)
|
|
sock.sendto(struct.pack("!H", OP_OACK) + payload + b"\x00", addr)
|
|
else:
|
|
blksize = 512
|
|
window = 1
|
|
|
|
with open(filepath, "rb") as f:
|
|
block = 0
|
|
eof = False
|
|
while not eof:
|
|
# send one window
|
|
for _ in range(window):
|
|
block = (block % 65536) + 1
|
|
chunk = f.read(blksize)
|
|
pkt = struct.pack("!HH", OP_DATA, block) + chunk
|
|
sock.sendto(pkt, addr)
|
|
if len(chunk) < blksize:
|
|
eof = True
|
|
# wait for the ACK of the last block of this window
|
|
expect = block
|
|
retries = 0
|
|
while True:
|
|
sock.settimeout(TIMEOUT)
|
|
try:
|
|
rdata, raddr = sock.recvfrom(4 + blksize)
|
|
except socket.timeout:
|
|
retries += 1
|
|
if retries > RETRIES:
|
|
return
|
|
# retransmit window
|
|
f.seek((expect - window) * blksize if expect >= window else 0)
|
|
block = expect - window if expect >= window else 0
|
|
eof = os.path.getsize(filepath) <= f.tell() + window * blksize
|
|
break
|
|
if len(rdata) >= 4 and struct.unpack("!H", rdata[:2])[0] == OP_ACK:
|
|
acked = struct.unpack("!H", rdata[2:4])[0]
|
|
if acked == expect:
|
|
break
|
|
elif acked < expect:
|
|
continue # stale ack
|
|
else:
|
|
continue # shouldn't happen
|
|
|
|
|
|
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()
|