#!/usr/bin/env python3
"""planck netboot helper: GET template tarball, PUT/GET per-node install artifacts.

PUTs land under PUT_ROOT (tftp/<serial>/...) for bootloader self-update.
"""
import os
import re
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

BASE = "/volume1/plancknetboot"
PUT_ROOT = os.path.join(BASE, "tftp")


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def _safe_target(self, sub):
        # each path component must be [A-Za-z0-9._-], no traversal
        parts = sub.split("/")
        if not parts or any(not re.fullmatch(r"[A-Za-z0-9._-]+", p) for p in parts):
            return None
        return sub

    def do_GET(self):
        if self.path.startswith("/events/"):
            name = self._safe_target(self.path[len("/events/"):].split("?")[0])
            if not name:
                self.send_error(403)
                return
            with open(os.path.join(PUT_ROOT, "events.log"), "a") as log:
                log.write(f"{self.log_date_time_string()} EVENT {name}\n")
            self.send_response(204)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        name = self._safe_target(self.path.lstrip("/"))
        if not name:
            self.send_error(403)
            return
        path = os.path.join(BASE, name)
        if not os.path.isfile(path):
            self.send_error(404)
            return
        size = os.path.getsize(path)
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(size))
        self.end_headers()
        with open(path, "rb") as f:
            while chunk := f.read(1024 * 1024):
                self.wfile.write(chunk)

    def do_PUT(self):
        name = self._safe_target(self.path.lstrip("/"))
        if not name:
            self.send_error(403)
            return
        path = os.path.join(PUT_ROOT, name)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        length = int(self.headers.get("Content-Length", 0))
        if length > 512 * 1024 * 1024:
            self.send_error(413)
            return
        with open(path + ".part", "wb") as f:
            remaining = length
            while remaining > 0:
                chunk = self.rfile.read(min(1024 * 1024, remaining))
                if not chunk:
                    break
                f.write(chunk)
                remaining -= len(chunk)
        os.rename(path + ".part", path)
        self.send_response(201)
        self.send_header("Content-Length", "0")
        self.end_headers()
        with open(os.path.join(PUT_ROOT, "events.log"), "a") as log:
            log.write(f"{self.log_date_time_string()} PUT {name} {length}\n")

    do_POST = do_PUT


if __name__ == "__main__":
    ThreadingHTTPServer(("0.0.0.0", 8000), Handler).serve_forever()