fetch.py

from apscheduler.executors.asyncio import AsyncIOExecutor
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from contextlib import asynccontextmanager
from datetime import datetime, timedelta
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"
PUP_URLS = {
    "Linux": "https://raw.githubusercontent.com/liquidcarbon/puppy/main/pup.sh",
    "Windows": "https://raw.githubusercontent.com/liquidcarbon/puppy/main/pup.ps1",
}


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)  # disable for now
app = FastAPI()


@app.get("/")
def read_root(request: Request) -> Response:
    """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"

    A slash ("/") in package name is interpreted as a GitHub repo.
    Package specs "duckdb>=1.1" are not supported.
    """

    query_params = dict(request.query_params)

    # exclude internal endpoints
    _ = query_params.pop("logs", "")

    # python version
    py_ver = query_params.pop("python", "3.12")
    if "Windows" in request.headers.get("user-agent"):
        pup_url = PUP_URLS.get("Windows")
        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 = PUP_URLS.get("Linux")
        script = [f"curl -fsSL {pup_url} | bash -s {py_ver}"]
        hint = f"""curl -fsSL "{request.url}" | bash"""

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

    # repos to be cloned
    to_clone = query_params.pop("clone", "")
    if to_clone:
        for repo in to_clone.split(","):
            if "/" in repo:
                pkg = f"https://github.com/{repo}.git"
                script.append(f"pup clone {pkg}")
            else:
                script.append(f"# can't clone `{repo}`: expected <username>/<reponame>")

    # remaining query params are venvs
    for venv, uv_packages in query_params.items():
        for pkg in uv_packages.split(","):
            if "/" in pkg:  # slash implies GitHub repo
                pkg = f"https://github.com/{pkg}.git"
            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_host1 = getenv("SPACE_HOST", "0.0.0.0:7860")
    self_host2 = "https://huggingface.co/spaces/pup-py/fetch"
    with httpx.Client() as client:
        _ = client.get(f"http://{self_host1}/ping", follow_redirects=True)
        _ = client.get(self_host2, follow_redirects=True)


scheduler = AsyncIOScheduler(executors={"default": AsyncIOExecutor()})
scheduler.add_job(self_ping, next_run_time=datetime.now() + timedelta(seconds=30))
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)