mini/mini.py

from __future__ import annotations
import asyncio as aio, base64 as b64, concurrent.futures as cf, datetime as dt, hashlib as hl, html as h, json as j, mimetypes as mt, re as r, sqlite3 as s3, time as t, urllib.parse as up
from pathlib import Path as P


class Time:
    @staticmethod
    def now():
        return dt.datetime.now(dt.timezone.utc)


class Ids:
    @staticmethod
    def uuid7():
        import uuid as u

        ms = int(t.time() * 1000)
        b = u.uuid4().bytes
        return str(u.UUID(bytes=(ms.to_bytes(6, "big") + b"\x70" + b[1:2] + b"\x80" + b[3:])[:16]))


class Esc:
    @staticmethod
    def x(s):
        return h.escape(str(s), quote=True)


class Field:
    def __init__(
        self, required=False, regex=None, type_=str, default=None, min_len=None, max_len=None
    ):
        self.required = required
        self.rx = r.compile(regex) if regex else None
        self.type_ = type_
        self.default = default
        self.min_len = min_len
        self.max_len = max_len

    def clean(self, name, val):
        if val is None or val == "":
            if self.required:
                raise ValueError(f"{name} is_required")
            return self.default
        v = self.type_(val)
        s = str(v)
        if self.min_len is not None and len(s) < self.min_len:
            raise ValueError(f"{name} too_short")
        if self.max_len is not None and len(s) > self.max_len:
            raise ValueError(f"{name} too_long")
        if self.rx and not self.rx.search(s):
            raise ValueError(f"{name} invalid")
        return v


class Validator:
    def __init__(self, schema: dict[str, Field]):
        self.schema = schema

    def openapi(self):
        t = {str: "string", int: "integer", float: "number", bool: "boolean"}
        ps = {}
        rq = []
        for k, f in self.schema.items():
            mp = {str: "string", int: "integer", float: "number", bool: "boolean"}
            o = {"type": mp.get(f.type_, "string")}
            if f.rx is not None:
                o["pattern"] = f.rx.pattern
            if f.min_len is not None:
                o["minLength"] = f.min_len
            if f.max_len is not None:
                o["maxLength"] = f.max_len
            if f.default is not None:
                o["default"] = f.default
            ps[k] = o
            if f.required:
                rq.append(k)
        o = {"type": "object", "properties": ps}
        if rq:
            o["required"] = rq
        return o

    def clean(self, data: dict):
        out = {}
        errs = {}
        for k, f in self.schema.items():
            try:
                out[k] = f.clean(k, data.get(k))
            except Exception as e:
                errs[k] = str(e)
        if errs:
            raise ValueError(errs)
        return out


class Request:
    def __init__(
        self,
        method,
        path,
        query,
        headers,
        cookies,
        path_params,
        query_params,
        form,
        body,
        json_data,
    ):
        self.method = method
        self.path = path
        self.query = query
        self.headers = headers
        self.cookies = cookies
        self.path_params = path_params
        self.query_params = query_params
        self.form = form
        self.body = body
        self.json = json_data

    def param(self, name, default=None):
        if name in self.path_params:
            return self.path_params[name]
        if name in self.query_params:
            return self.query_params[name]
        if name in self.form:
            return self.form[name]
        if isinstance(self.json, dict) and name in self.json:
            return self.json[name]
        return default

    @property
    def all_params(self):
        d = {}
        d.update(self.query_params)
        d.update(self.form)
        if isinstance(self.json, dict):
            d.update(self.json)
        d.update(self.path_params)
        return d


class Response:
    def __init__(self, body=b"", status=200, ctype="text/plain", headers=None):
        self.body = body if isinstance(body, bytes) else str(body).encode()
        self.status = status
        self.ctype = ctype
        self.headers = headers or {}

    @staticmethod
    def text(s, st=200):
        return Response(s, st, "text/plain")

    @staticmethod
    def html(s, st=200):
        return Response(s, st, "text/html")

    @staticmethod
    def json(o, st=200):
        return Response(j.dumps(o), st, "application/json")

    @staticmethod
    def redirect(u):
        return Response(b"", 302, "text/plain", {"Location": u})


class Template:
    def __init__(self, dir="templates"):
        self.dir = P(dir)

    def render(self, name, ctx: dict):
        s = (self.dir / name).read_text()
        for k, v in ctx.items():
            s = s.replace("{{" + k + "}}", Esc.x(v))
        return Response.html(s)


class DB:
    def __init__(self, path="./data/app.db"):
        self.path = P(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._w = aio.Lock()
        self._idx = set()
        self._queue = aio.Queue()
        self.con = s3.connect(self.path, check_same_thread=False, detect_types=s3.PARSE_DECLTYPES)
        self.con.row_factory = s3.Row

    def _cols(self, t):
        return [x[1] for x in self.con.execute(f"PRAGMA table_info({t})")]

    def _ensure(self, t, row: dict):
        try:
            self.con.execute(
                f"CREATE TABLE IF NOT EXISTS {t} (id INTEGER PRIMARY KEY AUTOINCREMENT,uid TEXT,created_at TEXT,updated_at TEXT,deleted_at TEXT)"
            )
        except Exception:
            pass
        cols = self._cols(t)
        for k in list(row.keys()):
            if k not in cols:
                try:
                    self.con.execute(f"ALTER TABLE {t} ADD COLUMN {k} TEXT")
                except Exception:
                    pass

    def _stamp(self, row):
        n = Time.now().isoformat()
        row.setdefault("uid", Ids.uuid7())
        row.setdefault("created_at", n)
        row["updated_at"] = n
        return row

    def _maybe_index(self, t, keys):
        if len(keys) > 1:
            k = (t, tuple(sorted(keys)))
            if k not in self._idx:
                try:
                    self.con.execute(
                        f"CREATE INDEX IF NOT EXISTS idx_{t}_{'_'.join(sorted(keys))} ON {t} ({','.join(sorted(keys))})"
                    )
                    self.con.commit()
                except Exception:
                    pass
                self._idx.add(k)

    async def insert(self, t, **kw):
        async with self._w:
            kw = self._stamp(dict(kw))
            for _ in range(6):
                try:
                    self._ensure(t, kw)
                    c = ",".join(kw.keys())
                    q = ",".join(["?"] * len(kw))
                    cur = self.con.execute(
                        f"INSERT INTO {t} ({c}) VALUES ({q})",
                        tuple(str(v) if v is not None else None for v in kw.values()),
                    )
                    self.con.commit()
                    kw["id"] = cur.lastrowid
                    return kw
                except s3.OperationalError:
                    self._ensure(t, kw)

    async def update(self, t, where, what):
        async with self._w:
            w = {where: [where]} if isinstance(where, list) else where
            what = dict(what)
            what["updated_at"] = Time.now().isoformat()
            for _ in range(6):
                try:
                    self._ensure(t, {**w, **what})
                    self._maybe_index(t, list(w.keys()))
                    self.con.execute(
                        f"UPDATE {t} SET {','.join(f'{k}=?' for k in what)} WHERE {' AND '.join(f'{k}=?' for k in w)} AND deleted_at IS NULL",
                        tuple([str(v) for v in what.values()] + [str(v) for v in w.values()]),
                    )
                    self.con.commit()
                    return True
                except s3.OperationalError:
                    self._ensure(t, {**w, **what})

    async def upsert(self, t, match: dict, **kw):
        async with self._w:
            pass
        rows = await self.find(t, match)
        if rows:
            await self.update(t, match, kw)
            return {**rows[0], **kw}
        return await self.insert(t, **{**match, **kw})

    async def delete(self, t, where):
        return await self.update(t, where, {"deleted_at": Time.now().isoformat()})

    async def find(self, t, where=None):
        async with self._w:
            try:
                if not where:
                    return [
                        dict(x)
                        for x in self.con.execute(f"SELECT * FROM {t} WHERE deleted_at IS NULL")
                    ]
                self._maybe_index(t, list(where.keys()))
                return [
                    dict(x)
                    for x in self.con.execute(
                        f"SELECT * FROM {t} WHERE {' AND '.join(f'{k}=?' for k in where)} AND deleted_at IS NULL",
                        tuple(str(v) for v in where.values()),
                    )
                ]
            except Exception:
                return []

    async def insert_many(self, t, rows: list[dict]):
        return [await self.insert(t, **x) for x in rows]

    async def update_many(self, t, items: list[tuple[dict, dict]]):
        return [await self.update(t, w, x) for w, x in items]

    async def upsert_many(self, t, items: list[tuple[dict, dict]]):
        return [await self.upsert(t, m, **x) for m, x in items]

    def _later_wrap(self, fn):
        async def w(*a, **k):
            opt = dict(k.get("_opt", {})) if isinstance(k.get("_opt"), dict) else {}
            preview = dict(k)
            preview.pop("_opt", None)
            if fn.__name__ in ("insert", "upsert"):
                preview = self._stamp(preview) if fn.__name__ == "insert" else preview
            await self._queue.put((fn, a, k))
            return preview

        return w

    def __getattr__(self, n):
        if n.endswith("_later") and n[:-6] in (
            "insert",
            "update",
            "upsert",
            "delete",
            "insert_many",
            "upsert_many",
            "update_many",
        ):
            return self._later_wrap(getattr(self, f"{n[:-6]}"))
        raise AttributeError(n)

    async def drain(self):
        while True:
            fn, a, k = await self._queue.get()
            try:
                await fn(*a, **{x: y for x, y in k.items() if x != "_opt"})
            except Exception:
                pass
            self._queue.task_done()


class Sessions:
    def __init__(self, db: DB, cookie="sid", cap=50 * 1024 * 1024):
        self.cookie = cookie
        self.db = db
        self.cache = {}
        self.size = 0
        self.cap = cap

    def _sid(self, cookies):
        return cookies.get(self.cookie) or Ids.uuid7()

    async def load(self, sid):
        if sid in self.cache:
            return self.cache[sid]
        r = await self.db.find("sessions", {"uid": sid})
        d = j.loads(r[0].get("data", "{}")) if r else {}
        self.cache[sid] = d
        return d

    async def save(self, sid, data):
        self.cache[sid] = data
        b = j.dumps(data).encode()
        self.size += len(b)
        r = await self.db.find("sessions", {"uid": sid})
        if r:
            await self.db.update("sessions", {"uid": sid}, {"data": j.dumps(data)})
        else:
            (
                await self.db.con.execute(
                    "INSERT INTO sessions (uid,data) VALUES (?,?)", (sid, j.dumps(data))
                )
                if False
                else await self.db.insert("sessions", uid=sid, data=j.dumps(data))
            )
        if self.size > self.cap:
            self.size = 0


class PubSub:
    def __init__(self):
        self.topics: dict[str, set] = {}

    def sub(self, topic, q):
        self.topics.setdefault(topic, set()).add(q)

    def unsub(self, topic, q):
        self.topics.get(topic, set()).discard(q)

    async def pub(self, topic, msg):
        for q in list(self.topics.get(topic, [])):
            try:
                q.put_nowait(msg)
            except Exception:
                pass


class Route:
    def __init__(self, method, pattern, handler, auth=False, schema_in=None, schema_out=None):
        self.method = method
        self.handler = handler
        self.auth = auth
        self.schema_in = schema_in
        self.schema_out = schema_out
        self.schema_in = schema_in
        self.schema_out = schema_out
        rx = r.sub(r"\{(\w+)\}", r"(?P<\1>[^/]+)", pattern)
        self.pattern = pattern
        self.rx = r.compile(f"^{rx}$")


class Mini:
    def __init__(self, data_dir="./data", version="v1"):
        self.root = P(data_dir)
        self.root.mkdir(parents=True, exist_ok=True)
        self.version = version
        self.routes: list[Route] = []
        self.db = DB(str(self.root / "app.db"))
        self.pubsub = PubSub()
        self.tpl = Template(str(self.root / "templates"))
        self.pool = cf.ThreadPoolExecutor(max_workers=32)
        mt.init()
        e = self._env_load()
        import uuid as _uu

        sc = e.get("session_cookie")
        xx = None
        yy = None
        zz = None
        if not sc:
            sc = f"sid_{_uu.uuid4().hex[:8]}"
            e["session_cookie"] = sc
            self._env_save(e)
        self.session_cookie = sc
        self.sessions = Sessions(self.db, cookie=sc)

    def route(self, path, method="GET", auth=False, schema_in=None, schema_out=None):
        def d(fn):
            self.routes.append(Route(method, path, fn, auth, schema_in, schema_out))
            try:
                self.openapi()
            except Exception:
                pass
            return fn

        return d

    def get(self, path, **k):
        return self.route(path, "GET", **k)

    def put(self, path, **k):
        return self.route(path, "PUT", **k)

    def post(self, path, **k):
        return self.route(path, "POST", **k)

    def static(self, url_prefix="/static", dir="static"):
        d = self.root / dir
        d.mkdir(parents=True, exist_ok=True)

        async def h(req: Request):
            rel = up.unquote(req.path_params.get("_p", ""))
            f = (d / rel).resolve()
            if d.resolve() not in f.parents and f != d.resolve():
                return Response.text("forbidden", 403)
            if not f.exists():
                return Response.text("nf", 404)
            c, enc = mt.guess_type(str(f))
            c = c or "application/octet-stream"
            rng = req.headers.get("range", "")
            data = f.read_bytes()
            if rng.startswith("bytes="):
                try:
                    s, e = rng[6:].split("-")
                    s = int(s or 0)
                    e = int(e or len(data) - 1)
                    return Response(
                        data[s : e + 1],
                        206,
                        c,
                        {"Content-Range": f"bytes {s}-{e}/{len(data)}", "Accept-Ranges": "bytes"},
                    )
                except Exception:
                    pass
            return Response(
                data,
                200,
                c,
                {
                    "Accept-Ranges": "bytes",
                    "Cache-Control": (
                        f"public,max-age=31536000"
                        if f"/{self.version}/" in req.path
                        else "no-cache"
                    ),
                },
            )

        self.routes.append(Route("GET", f"{url_prefix}/{{_p}}", h))

    def openapi(self):
        pp = {}
        rr = None
        oo = None
        dd = None
        for rt in self.routes:
            doc = (rt.handler.__doc__ or "").strip()
            summ = doc.splitlines()[0] if doc else f"{rt.method} {rt.pattern}"
            op = {"summary": summ, "description": doc, "responses": {"200": {"description": "ok"}}}
            if rt.schema_in:
                op["requestBody"] = {
                    "content": {"application/json": {"schema": rt.schema_in.openapi()}}
                }
            if rt.schema_out:
                op["responses"]["200"]["content"] = {
                    "application/json": {"schema": rt.schema_out.openapi()}
                }
            pp.setdefault(rt.pattern, {}).setdefault(rt.method.lower(), op)
        self._spec = {
            "openapi": "3.0.0",
            "info": {"title": "mini", "version": self.version},
            "paths": pp,
        }
        try:
            (self.root / "openapi.json").write_text(j.dumps(self._spec, indent=2))
        except Exception:
            pass
        return self._spec

    def _parse(self, raw: bytes):
        head, _, body = raw.partition(b"\r\n\r\n")
        lines = head.decode("iso-8859-1").split("\r\n")
        m, p, _ = lines[0].split(" ", 2)
        path, q = p.split("?", 1) if "?" in p else (p, "")
        hd = {}
        for l in lines[1:]:
            if ":" in l:
                k, v = l.split(":", 1)
                hd[k.strip().lower()] = v.strip()
        ck = {}
        for c in hd.get("cookie", "").split(";"):
            if "=" in c:
                k, v = c.strip().split("=", 1)
                ck[k] = up.unquote(v)
        qp = {
            k: v[0] if len(v) == 1 else v for k, v in up.parse_qs(q, keep_blank_values=True).items()
        }
        fm = {}
        jd = None
        ct = hd.get("content-type", "")
        if m in ("POST", "PUT", "PATCH"):
            if "application/json" in ct:
                try:
                    jd = j.loads(body.decode() or "null")
                except Exception:
                    jd = None
            elif "application/x-www-form-urlencoded" in ct:
                fm = {
                    k: v[0] if len(v) == 1 else v
                    for k, v in up.parse_qs(body.decode(), keep_blank_values=True).items()
                }
        return m, up.unquote(path), q, hd, ck, qp, fm, body, jd

    def _check_auth(self, hd):
        a = hd.get("authorization", "")
        if not a.startswith("Basic "):
            return None
        try:
            u, p = b64.b64decode(a[6:]).decode().split(":", 1)
            return {"user": u, "pass": p}
        except Exception:
            return None

    async def _handle_conn(self, r, w):
        try:
            raw = await aio.wait_for(r.read(65536), timeout=10)
            if not raw:
                w.close()
                return
            m, path, q, hd, ck, qp, fm, body, jd = self._parse(raw)
            for rt in self.routes:
                if rt.method != m:
                    continue
                mo = rt.rx.match(path)
                if not mo:
                    continue
                if rt.auth and not self._check_auth(hd):
                    return await self._send(
                        w,
                        Response(
                            b"auth", 401, "text/plain", {"WWW-Authenticate": 'Basic realm="mini"'}
                        ),
                        None,
                    )
                pp = {k: up.unquote(v) for k, v in mo.groupdict().items()}
                req = Request(m, path, q, hd, ck, pp, qp, fm, body, jd)
                sid = ck.get(self.session_cookie) or Ids.uuid7()
                new = self.session_cookie not in ck
                req.session = await self.sessions.load(sid)
                req.validated = {}
                vv = None
                ee = None
                if rt.schema_in:
                    try:
                        req.validated = rt.schema_in.clean(req.all_params)
                    except Exception as ex:
                        return await self._send(
                            w, Response.json({"ok": False, "errors": str(ex)}, 422), None
                        )
                res = rt.handler(req)
                if aio.iscoroutine(res):
                    res = await res
                else:
                    res = (
                        await aio.get_event_loop().run_in_executor(self.pool, rt.handler, req)
                        if callable(rt.handler) and res is None
                        else res
                    )
                if isinstance(res, (dict, list)):
                    res = Response.json(res)
                elif isinstance(res, str):
                    res = Response.html(res)
                await self.sessions.save(sid, req.session)
                extra = (
                    {}
                    if not new
                    else {"Set-Cookie": f"{self.session_cookie}={sid}; Path=/; HttpOnly"}
                )
                return await self._send(w, res, extra)
            await self._send(w, Response.text("nf", 404), None)
        except Exception as e:
            try:
                await self._send(w, Response.text("err", 500), None)
            except Exception:
                pass
        finally:
            try:
                w.close()
            except Exception:
                pass

    async def _send(self, w, res: Response, extra):
        hds = {
            "Content-Type": res.ctype,
            "Content-Length": str(len(res.body)),
            "Connection": "close",
            **res.headers,
            **(extra or {}),
        }
        w.write(
            f"HTTP/1.1 {res.status} OK\r\n{''.join(f'{k}: {v}\r\n' for k,v in hds.items())}\r\n".encode()
            + res.body
        )
        await w.drain()

    async def _run(self, port):
        aio.create_task(self.db.drain())
        s = await aio.start_server(self._handle_conn, "0.0.0.0", port)
        print(f"mini on {port}")
        async with s:
            await s.serve_forever()

    def _env_path(self):
        return self.root / ".env.json"

    def _env_load(self):
        try:
            import json as _j

            return _j.loads((self.root / ".env.json").read_text())
        except Exception:
            return {}

    def _env_save(self, d):
        try:
            import json as _j

            (self.root / ".env.json").write_text(_j.dumps(d, indent=2))
        except Exception:
            pass

    def _auto_port(self):
        e = self._env_load()
        last = None
        try:
            last = int(e.get("port", 0))
        except Exception:
            pass
        cands = []
        if last and 22000 <= last <= 24000:
            cands.append(last)
        import random as rn

        rest = list(range(22000, 24001))
        rn.shuffle(rest)
        for x in rest:
            if x not in cands:
                cands.append(x)
        import socket as so

        for x in cands:
            try:
                q = so.socket()
                q.bind(("0.0.0.0", x))
                q.close()
                e["port"] = x
                self._env_save(e)
                return x
            except Exception:
                continue
        return 23000

    def serve(self, port=None):
        import sys as _s

        if port is None:
            port = int(_s.argv[1]) if len(_s.argv) > 1 else self._auto_port()
        aio.run(self._run(port))


if __name__ == "__main__":
    import sys

    app = Mini()
    app.static()

    @app.get("/")
    async def idx(req):
        return {"ok": True, "v": app.version}

    app.serve(int(sys.argv[1]) if len(sys.argv) > 1 else None)