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