fetch.py

from apscheduler.executors.asyncio import AsyncIOExecutor
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from contextlib import asynccontextmanager
from datetime import datetime
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, PlainTextResponse, Response
from os import environ, getenv
from pathlib import Path
from typing import Any

import httpx
import json
import subprocess

LOGFILE = Path.home() / "a.json"


class PrettyJSONResponse(Response):
    media_type = "application/json"

    def render(self, content: Any) -> bytes:
        return json.dumps(content, indent=2).encode("utf-8")


@asynccontextmanager
async def lifespan(app: FastAPI):
    scheduler.start()
    yield
    scheduler.shutdown()


app = FastAPI(lifespan=lifespan)


@app.get("/")
def read_root(request: Request):
    """Main URL returning an executable installer script.

    Query parameters can be used to install specific things, e.g.
    curl -fsSL "https://pup-py-fetch.hf.space?python=3.11&pixi=marimo&myenv=cowsay,duckdb"

    Package specs "duckdb>=1.1" are not supported.
    """

    query_params = dict(request.query_params)
    py_ver = query_params.pop("python", "3.12")
    if "Windows" in request.headers.get("user-agent"):
        pup_url = "https://raw.githubusercontent.com/liquidcarbon/puppy/main/pup.ps1"
        script = [
            f"$pup_ps1 = (iwr -useb {pup_url}).Content",
            f"& ([scriptblock]::Create($pup_ps1)) {py_ver}",
        ]
        hint = f"""iex (iwr "{request.url}").Content"""
    else:
        pup_url = "https://raw.githubusercontent.com/liquidcarbon/puppy/main/pup.sh"
        script = [f"curl -fsSL {pup_url} | bash -s {py_ver}"]
        hint = f"""curl -fsSL "{request.url}" | bash"""

    pixi_packages = query_params.pop("pixi", "")
    if pixi_packages:
        for pkg in pixi_packages.split(","):
            script.append(f"pixi add {pkg}")

    # remaining query params are venvs
    for venv, uv_packages in query_params.items():
        if venv in ["logs"]:
            continue
        for pkg in uv_packages.split(","):
            script.append(f"pup add {venv} {pkg}")

    script.extend(
        [
            "# 🐶 scripts end here; if you like what you see, copy-paste this recipe, or run like so:",
            f"# {hint}",
            "# to learn more, visit https://github.com/liquidcarbon/puppy\n",
        ]
    )
    return PlainTextResponse("\n".join(script))


@app.middleware("http")
async def log_request(request: Request, call_next: Any):
    ts = datetime.now().strftime("%y%m%d%H%M%S%f")
    data = {
        # "day": int(ts[:6]),
        "dt": int(ts[:-3]),
        "url": request.url,
        "query_params": request.query_params,
        "user-agent": request.headers.get("user-agent"),
        "client": request.headers.get("x-forwarded-for"),
        "private_ip": request.client.host,
        "method": request.method,
        "headers": str(request.headers),
    }
    output = json.dumps(obj=data, default=str, indent=None, separators=(", ", ":"))
    with open(LOGFILE, "a") as f:
        f.write(output + "\n")

    response = await call_next(request)
    return response


@app.get("/a", response_class=PrettyJSONResponse)
def get_analytics(n: int = 5):
    if n == 0:
        cmd = f"tac {LOGFILE.as_posix()}"
    else:
        cmd = f"tail -n {n} {LOGFILE.as_posix()} | tac"
    _subprocess = subprocess.run(cmd, shell=True, text=True, capture_output=True)
    json_lines, stderr = _subprocess.stdout[-1], _subprocess.stderr
    try:
        content = json.loads(f"[ {json_lines.replace("\n", ",")} ]")
        return content
    except Exception as e:
        return {"error": str(e), "stderr": stderr, "json_lines": json_lines}


@app.api_route("/qa", response_class=FileResponse, methods=["GET", "HEAD"])
def query_analytics():
    return LOGFILE.as_posix()


@app.api_route("/env", response_class=PrettyJSONResponse)
def show_env():
    return environ


@app.get("/favicon.ico")
async def favicon():
    return {"message": "woof!"}


@app.get("/ping")
async def ping():
    return {"message": "woof!"}


def self_ping():
    self_host = getenv("SPACE_HOST", "0.0.0.0:7860")
    with httpx.Client() as client:
        _ = client.get(f"http://{self_host}/ping", follow_redirects=True)


scheduler = AsyncIOScheduler(executors={"default": AsyncIOExecutor()})
scheduler.add_job(self_ping, "interval", minutes=720)


if __name__ == "__main__":
    import uvicorn

    fmt = "%(asctime)s %(levelprefix)s %(message)s"
    uvicorn_logging = uvicorn.config.LOGGING_CONFIG
    uvicorn_logging["formatters"]["access"]["datefmt"] = "%y%m%d @ %T"
    uvicorn_logging["formatters"]["access"]["fmt"] = fmt
    uvicorn_logging["formatters"]["default"]["datefmt"] = "%y%m%d @ %T"
    uvicorn_logging["formatters"]["default"]["fmt"] = fmt
    uvicorn.run(app, host="0.0.0.0", port=7860)