share/server.py

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

import email.parser
import email.policy
import hashlib
import hmac
import http.server
import json
import mimetypes
import re
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

# Keeps `?token=...` out of the log / journald.
_TOKEN_RE = re.compile(r"([?&]token=)[^&\s\"]*")


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'


def _disposition_filename(name: str) -> str:
    """Content-Disposition filename params: ASCII fallback + RFC 5987 UTF-8 form.

    HTTP headers are latin-1 only, so a non-ASCII name must be percent-encoded.
    """
    ascii_name = name.encode("ascii", "replace").decode("ascii")
    ascii_name = ascii_name.replace("\\", "_").replace('"', "_").replace("?", "_")
    return f'filename="{ascii_name}"; filename*=UTF-8\'\'{urllib.parse.quote(name, safe="")}'


class ShareServer(http.server.ThreadingHTTPServer):
    def __init__(self, addr, handler, *, upload_dir: Path, html: bytes, ssl_ctx: ssl.SSLContext | None = None, token: str = "", token_is_hash: bool = False):
        super().__init__(addr, handler)
        self.upload_dir = upload_dir
        self.html = html
        self.ssl_ctx = ssl_ctx
        if not token:
            self.token_hash = ""
        elif token_is_hash:
            self.token_hash = token
        else:
            self.token_hash = hashlib.sha256(token.encode()).hexdigest()

    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 end_headers(self):
        # Stops a `?token=` link from leaking the token via Referer on outbound clicks.
        self.send_header("Referrer-Policy", "no-referrer")
        super().end_headers()

    def log_message(self, fmt, *args):
        msg = str(args[0]) if args else fmt
        print(f"  [{datetime.now():%H:%M:%S}] {_TOKEN_RE.sub(r'\1<redacted>', msg)}")

    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)

    # ── Auth ──

    @staticmethod
    def _hash(token: str) -> str:
        return hashlib.sha256(token.encode()).hexdigest()

    def _match(self, value: str) -> bool:
        return hmac.compare_digest(self._hash(value), self.server.token_hash)

    def _query_token(self) -> str:
        qs = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query)
        return qs.get("token", [""])[0]

    def _cookie(self, token: str) -> str:
        return f"share_token={self._hash(token)}; Path=/; HttpOnly; Secure; SameSite=Strict; Max-Age=31536000"

    def _check_auth(self) -> bool:
        # NB: `?token=` is deliberately not accepted here — a token in the URL
        # leaks into logs, browser history and Referer. It is exchanged for a
        # cookie on the index route only (see _dispatch).
        if not self.server.token_hash:
            return True
        auth = self.headers.get("Authorization", "")
        if auth.startswith("Bearer ") and self._match(auth[7:]):
            return True
        cookie = self.headers.get("Cookie", "")
        for part in cookie.split(";"):
            part = part.strip()
            if part.startswith("share_token=") and hmac.compare_digest(part[12:], self.server.token_hash):
                return True
        return False

    _AUTH_PAGE = b'''<!DOCTYPE html>
<html><head>
<meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>Share</title>
<style>
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:system-ui,-apple-system,sans-serif;background:#0a0a0a;color:#e0e0e0;min-height:100vh;display:flex;align-items:center;justify-content:center}
.card{background:#141414;border-radius:12px;padding:40px 32px;width:100%;max-width:360px;text-align:center}
h1{font-size:1.3rem;color:#fff;margin-bottom:24px}
input{width:100%;padding:12px 14px;background:#1a1a1a;border:1px solid #333;border-radius:8px;color:#e0e0e0;font-size:.95rem;outline:none;margin-bottom:16px;transition:border-color .2s}
input:focus{border-color:#4fc3f7}
button{width:100%;padding:12px;background:#4fc3f7;color:#000;border:none;border-radius:8px;font-size:.95rem;font-weight:600;cursor:pointer;transition:opacity .2s}
button:hover{opacity:.85}
.err{color:#ef5350;font-size:.85rem;margin-top:12px;min-height:1.2em}
</style>
</head><body>
<div class="card">
<h1>Share</h1>
<form id="f">
<input type="password" id="tok" placeholder="Token" autofocus>
<button type="submit">Login</button>
</form>
<div class="err" id="err"></div>
</div>
<script>
const f=document.getElementById('f'),tok=document.getElementById('tok'),err=document.getElementById('err');
f.onsubmit=e=>{e.preventDefault();err.textContent='';
fetch('auth',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({token:tok.value})})
.then(r=>{if(r.ok)location.reload();else{err.textContent='Invalid token';tok.select()}})
.catch(()=>{err.textContent='Connection error'})};
</script>
</body></html>'''

    def _token_to_cookie(self, token: str):
        """Swap a valid `?token=` link for a cookie, then redirect to a clean URL."""
        self.send_response(302)
        self.send_header("Set-Cookie", self._cookie(token))
        self.send_header("Location", "/")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", "0")
        self.send_header("Connection", "close")
        self.end_headers()

    def _send_auth_page(self):
        self.send_response(401)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(self._AUTH_PAGE)))
        self.send_header("Connection", "close")
        self.end_headers()
        self.wfile.write(self._AUTH_PAGE)

    # ── Routes ──

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

    def _dispatch(self):
        parsed = urllib.parse.urlparse(self.path)
        path_only = parsed.path

        for prefix, method in self.ROUTES.get(self.command, []):
            if path_only == prefix or (len(prefix) > 1 and prefix.endswith("/") and path_only.startswith(prefix)):
                if method not in ("_route_auth",) and not self._check_auth():
                    if self.command == "GET" and path_only in ("/", "/index.html"):
                        tok = self._query_token()
                        if tok and self._match(tok):
                            return self._token_to_cookie(tok)
                        return self._send_auth_page()
                    self.send_response(401)
                    body = json.dumps({"error": "unauthorized"}).encode()
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(body)))
                    self.send_header("Connection", "close")
                    self.end_headers()
                    self.wfile.write(body)
                    return
                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:
        parsed = urllib.parse.urlparse(self.path)
        name = urllib.parse.unquote(parsed.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"
        self.send_response(200)
        self.send_header("Content-Type", ct)
        self.send_header("Content-Disposition", f"attachment; {_disposition_filename(safe.name)}")
        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_auth(self, _prefix):
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length) if length else b""
        try:
            data = json.loads(body)
        except (json.JSONDecodeError, ValueError):
            data = {}
        if data.get("token") and self._match(data["token"]):
            resp = json.dumps({"ok": True}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(resp)))
            self.send_header("Set-Cookie", self._cookie(data["token"]))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(resp)
        else:
            resp = json.dumps({"ok": False}).encode()
            self.send_response(401)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(resp)))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(resp)

    def _route_logout(self, _prefix):
        self.send_response(302)
        self.send_header("Set-Cookie", "share_token=; Path=/; HttpOnly; Secure; SameSite=Strict; Max-Age=0")
        self.send_header("Location", "/")
        self.send_header("Connection", "close")
        self.end_headers()

    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()
        # policy=HTTP decodes headers as UTF-8; the default compat32 policy mangles
        # non-ASCII filenames into replacement characters.
        msg = email.parser.BytesParser(policy=email.policy.HTTP).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