share/server.py

"""HTTP server and request handler."""

import email.parser
import http.server
import json
import mimetypes
import ssl
import stat
import urllib.parse
from datetime import datetime
from pathlib import Path

from config import CHUNK_SIZE, SNIFF_SIZE, SSL_HANDSHAKE_TIMEOUT


def file_type(path: Path) -> str:
    ct = mimetypes.guess_type(path.name)[0] or ''
    if ct.startswith('image/'):
        return 'image'
    if ct.startswith('text/'):
        return 'text'
    try:
        with path.open('rb') as f:
            return 'text' if b'\x00' not in f.read(SNIFF_SIZE) else 'other'
    except OSError:
        return 'other'


class ShareServer(http.server.ThreadingHTTPServer):
    def __init__(self, addr, handler, *, upload_dir: Path, html: bytes, ssl_ctx: ssl.SSLContext | None = None):
        super().__init__(addr, handler)
        self.upload_dir = upload_dir
        self.html = html
        self.ssl_ctx = ssl_ctx

    def get_request(self):
        client, addr = self.socket.accept()
        if self.ssl_ctx:
            client.settimeout(SSL_HANDSHAKE_TIMEOUT)
            try:
                client = self.ssl_ctx.wrap_socket(client, server_side=True)
            except (ssl.SSLError, OSError):
                client.close()
                raise
            client.settimeout(None)
        return client, addr


class Handler(http.server.BaseHTTPRequestHandler):
    server: ShareServer

    def finish(self):
        try:
            super().finish()
        except (ssl.SSLError, BrokenPipeError, ConnectionResetError, OSError):
            pass

    def log_message(self, fmt, *args):
        print(f"  [{datetime.now():%H:%M:%S}] {args[0]}")

    def _json(self, data):
        self._respond(json.dumps(data).encode(), "application/json")

    def _respond(self, data: bytes, ct: str):
        self.send_response(200)
        self.send_header("Content-Type", ct)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Connection", "close")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(data)

    # ── Routes ──

    ROUTES = {
        "GET": [
            ("/",          "_route_index"),
            ("/index.html","_route_index"),
            ("/files",     "_route_files"),
            ("/dl/",       "_route_download"),
        ],
        "POST":   [("/upload", "_route_upload")],
        "DELETE": [("/del/",   "_route_delete")],
    }

    def _dispatch(self):
        for prefix, method in self.ROUTES.get(self.command, []):
            if self.path == prefix or (len(prefix) > 1 and prefix.endswith("/") and self.path.startswith(prefix)):
                return getattr(self, method)(prefix)
        self.send_error(404)

    do_GET = do_POST = do_DELETE = _dispatch

    def _route_index(self, _prefix):
        self._respond(self.server.html, "text/html; charset=utf-8")

    def _route_files(self, _prefix):
        entries = []
        for p in self.server.upload_dir.iterdir():
            if p.is_file():
                entries.append((p, p.stat()))
        entries.sort(key=lambda e: e[1].st_mtime, reverse=True)
        self._json([{"name": p.name, "size": st.st_size, "type": file_type(p)} for p, st in entries])

    def _safe_path(self, prefix: str) -> Path:
        name = urllib.parse.unquote(self.path[len(prefix):])
        return self.server.upload_dir / Path(name).name

    def _route_download(self, prefix):
        safe = self._safe_path(prefix)
        try:
            st = safe.stat()
        except FileNotFoundError:
            self.send_error(404)
            return
        if not stat.S_ISREG(st.st_mode):
            self.send_error(404)
            return
        ct = mimetypes.guess_type(safe.name)[0] or "application/octet-stream"
        fname = safe.name.replace("\\", "\\\\").replace('"', '\\"')
        self.send_response(200)
        self.send_header("Content-Type", ct)
        self.send_header("Content-Disposition", f'attachment; filename="{fname}"')
        self.send_header("Content-Length", str(st.st_size))
        self.send_header("Connection", "close")
        self.end_headers()
        with safe.open("rb") as f:
            while chunk := f.read(CHUNK_SIZE):
                self.wfile.write(chunk)

    def _route_upload(self, _prefix):
        ct = self.headers.get("Content-Type", "")
        if "multipart/form-data" not in ct:
            self.send_error(400)
            return
        length = int(self.headers.get("Content-Length", 0))
        if not length:
            self.send_error(400)
            return
        body = self.rfile.read(length)
        count = self._save_parts(ct, body)
        self._json({"ok": True, "count": count})

    def _route_delete(self, prefix):
        safe = self._safe_path(prefix)
        if safe.is_file():
            safe.unlink()
        self._json({"ok": True})

    # ── Helpers ──

    def _save_parts(self, content_type: str, body: bytes) -> int:
        header = f"Content-Type: {content_type}\r\n\r\n".encode()
        msg = email.parser.BytesParser().parsebytes(header + body)
        count = 0
        for part in msg.walk():
            fname = part.get_filename()
            if not fname:
                continue
            fname = Path(fname).name or f"paste_{datetime.now():%H%M%S}"
            data = part.get_payload(decode=True)
            if not data:
                continue
            dest = self._unique_path(fname)
            dest.write_bytes(data)
            count += 1
            print(f"  + {dest.name} ({len(data)} bytes)")
        return count

    def _unique_path(self, filename: str) -> Path:
        dest = self.server.upload_dir / filename
        if not dest.exists():
            return dest
        stem, suffix = dest.stem, dest.suffix
        i = 1
        while dest.exists():
            dest = self.server.upload_dir / f"{stem}_{i}{suffix}"
            i += 1
        return dest