#!/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()