448 lines
15 KiB
Python
448 lines
15 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import math
|
||
import os
|
||
import secrets
|
||
from contextlib import asynccontextmanager
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Annotated, Any
|
||
from urllib.parse import urlencode
|
||
from uuid import uuid4
|
||
|
||
from fastapi import APIRouter, Depends, FastAPI, Form, HTTPException, Request
|
||
from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
|
||
from fastapi.security import HTTPBasic, HTTPBasicCredentials
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.templating import Jinja2Templates
|
||
from PIL import Image, ImageOps
|
||
|
||
from config import Settings
|
||
from db import Database
|
||
from db import Request as HistoryRequest
|
||
|
||
BASE_DIR = Path(__file__).resolve().parent
|
||
TEMPLATES_DIR = BASE_DIR / "web" / "templates"
|
||
STATIC_DIR = BASE_DIR / "web" / "static"
|
||
|
||
STATUSES = ("queued", "running", "done", "failed")
|
||
MODES = ("t2i", "i2i")
|
||
|
||
logger = logging.getLogger(__name__)
|
||
security = HTTPBasic(auto_error=False)
|
||
|
||
|
||
def _safe_storage_path(raw_path: str, storage_root: Path) -> Path | None:
|
||
path = Path(raw_path).resolve()
|
||
try:
|
||
path.relative_to(storage_root)
|
||
except ValueError:
|
||
return None
|
||
return path
|
||
|
||
|
||
def _make_thumbnail(source: Path, destination: Path) -> None:
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path = destination.with_name(f"{destination.name}.{uuid4().hex}.tmp")
|
||
try:
|
||
with Image.open(source) as image:
|
||
image = ImageOps.exif_transpose(image)
|
||
image.thumbnail((320, 320))
|
||
if image.mode in {"RGBA", "LA"} or "transparency" in image.info:
|
||
rgba = image.convert("RGBA")
|
||
background = Image.new("RGB", rgba.size, (15, 17, 21))
|
||
background.paste(rgba, mask=rgba.getchannel("A"))
|
||
image = background
|
||
else:
|
||
image = image.convert("RGB")
|
||
image.save(temp_path, format="JPEG", quality=85, optimize=True)
|
||
os.replace(temp_path, destination)
|
||
finally:
|
||
temp_path.unlink(missing_ok=True)
|
||
|
||
|
||
def fmt_dt(value: str | None) -> str:
|
||
if not value:
|
||
return "—"
|
||
try:
|
||
parsed = datetime.fromisoformat(value)
|
||
return parsed.astimezone().strftime("%d.%m.%Y %H:%M")
|
||
except ValueError:
|
||
return value
|
||
|
||
|
||
def create_app(settings: Settings | None = None) -> FastAPI:
|
||
settings = settings or Settings.from_env(require_token=False)
|
||
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))
|
||
templates.env.filters["fmt_dt"] = fmt_dt
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
db = Database(
|
||
settings.database_url,
|
||
min_size=settings.db_pool_min_size,
|
||
max_size=settings.db_pool_max_size,
|
||
)
|
||
await db.connect()
|
||
app.state.db = db
|
||
yield
|
||
await db.close()
|
||
|
||
app = FastAPI(
|
||
title="MpakoImageAIbot history",
|
||
docs_url=None,
|
||
redoc_url=None,
|
||
lifespan=lifespan,
|
||
)
|
||
app.state.settings = settings
|
||
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
|
||
|
||
async def require_admin(
|
||
request: Request,
|
||
credentials: Annotated[
|
||
HTTPBasicCredentials | None, Depends(security)
|
||
],
|
||
) -> None:
|
||
current: Settings = request.app.state.settings
|
||
if not current.web_admin_password:
|
||
raise HTTPException(
|
||
status_code=503,
|
||
detail="Админ-панель отключена: задай WEB_ADMIN_PASSWORD",
|
||
)
|
||
if credentials is None:
|
||
raise HTTPException(
|
||
status_code=401,
|
||
detail="Требуется авторизация",
|
||
headers={"WWW-Authenticate": "Basic"},
|
||
)
|
||
username_ok = secrets.compare_digest(
|
||
credentials.username, current.web_admin_username
|
||
)
|
||
password_ok = secrets.compare_digest(
|
||
credentials.password, current.web_admin_password
|
||
)
|
||
if not (username_ok and password_ok):
|
||
raise HTTPException(
|
||
status_code=401,
|
||
detail="Неверный логин или пароль",
|
||
headers={"WWW-Authenticate": "Basic"},
|
||
)
|
||
|
||
async def history_page(
|
||
request: Request,
|
||
*,
|
||
admin: bool,
|
||
user_id: str,
|
||
status: str,
|
||
mode: str,
|
||
page: int,
|
||
):
|
||
db: Database = request.app.state.db
|
||
current: Settings = request.app.state.settings
|
||
|
||
error = ""
|
||
uid: int | None = None
|
||
raw_user_id = user_id.strip()
|
||
if raw_user_id:
|
||
if raw_user_id.isdigit():
|
||
uid = int(raw_user_id)
|
||
else:
|
||
error = "user_id должен состоять только из цифр."
|
||
|
||
status = status if status in STATUSES else ""
|
||
mode = mode if mode in MODES else ""
|
||
|
||
total = await db.count_requests(
|
||
user_id=uid,
|
||
status=status or None,
|
||
mode=mode or None,
|
||
include_hidden=admin,
|
||
)
|
||
page_size = current.web_page_size
|
||
pages = max(1, math.ceil(total / page_size))
|
||
page = max(1, min(page, pages))
|
||
offset = (page - 1) * page_size
|
||
|
||
rows = await db.list_requests(
|
||
user_id=uid,
|
||
status=status or None,
|
||
mode=mode or None,
|
||
include_hidden=admin,
|
||
limit=page_size,
|
||
offset=offset,
|
||
)
|
||
users = await db.list_users(include_hidden=admin)
|
||
stats = await db.stats(include_hidden=admin)
|
||
|
||
base_filters = {
|
||
key: value
|
||
for key, value in (
|
||
("user_id", raw_user_id),
|
||
("status", status),
|
||
("mode", mode),
|
||
)
|
||
if value
|
||
}
|
||
base_path = "/admin" if admin else ""
|
||
history_path = base_path or "/"
|
||
|
||
def page_url(number: int) -> str:
|
||
params = dict(base_filters)
|
||
if number > 1:
|
||
params["page"] = str(number)
|
||
query = urlencode(params)
|
||
return f"{history_path}?{query}" if query else history_path
|
||
|
||
context: dict[str, Any] = {
|
||
"rows": rows,
|
||
"users": users,
|
||
"stats": stats,
|
||
"statuses": STATUSES,
|
||
"modes": MODES,
|
||
"user_id": raw_user_id,
|
||
"status": status,
|
||
"mode": mode,
|
||
"error": error,
|
||
"page": page,
|
||
"pages": pages,
|
||
"total": total,
|
||
"admin": admin,
|
||
"base_path": base_path,
|
||
"history_path": history_path,
|
||
"prev_url": page_url(page - 1) if page > 1 else None,
|
||
"next_url": page_url(page + 1) if page < pages else None,
|
||
}
|
||
return templates.TemplateResponse(
|
||
request=request, name="index.html", context=context
|
||
)
|
||
|
||
async def detail_page(
|
||
request: Request, request_id: int, *, admin: bool
|
||
):
|
||
db: Database = request.app.state.db
|
||
item = await db.get_request(request_id)
|
||
if item is None or (item.hidden and not admin):
|
||
raise HTTPException(status_code=404, detail="Запрос не найден")
|
||
return templates.TemplateResponse(
|
||
request=request,
|
||
name="detail.html",
|
||
context={
|
||
"item": item,
|
||
"admin": admin,
|
||
"base_path": "/admin" if admin else "",
|
||
},
|
||
)
|
||
|
||
async def media_response(
|
||
request: Request, request_id: int, kind: str, *, admin: bool
|
||
):
|
||
if kind not in {"input", "output"}:
|
||
raise HTTPException(status_code=404, detail="Не найдено")
|
||
db: Database = request.app.state.db
|
||
current: Settings = request.app.state.settings
|
||
item: HistoryRequest | None = await db.get_request(request_id)
|
||
if item is None or (item.hidden and not admin):
|
||
raise HTTPException(status_code=404, detail="Запрос не найден")
|
||
|
||
raw_path = (
|
||
item.input_file_path if kind == "input" else item.output_file_path
|
||
)
|
||
if not raw_path:
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
|
||
storage_root = current.storage_dir.resolve()
|
||
path = _safe_storage_path(raw_path, storage_root)
|
||
if path is None or not path.is_file():
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
return FileResponse(path)
|
||
|
||
async def thumbnail_response(
|
||
request: Request, request_id: int, *, admin: bool
|
||
):
|
||
db: Database = request.app.state.db
|
||
current: Settings = request.app.state.settings
|
||
item: HistoryRequest | None = await db.get_request(request_id)
|
||
if item is None or (item.hidden and not admin):
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
if not item.output_file_path:
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
|
||
storage_root = current.storage_dir.resolve()
|
||
source = _safe_storage_path(item.output_file_path, storage_root)
|
||
if source is None or not source.is_file():
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
|
||
cache_dir = storage_root / ".thumbnails"
|
||
cache_path = cache_dir / f"{request_id}.jpg"
|
||
if (
|
||
cache_path.is_file()
|
||
and cache_path.stat().st_mtime >= source.stat().st_mtime
|
||
):
|
||
return FileResponse(
|
||
cache_path,
|
||
media_type="image/jpeg",
|
||
headers={"Cache-Control": "public, max-age=3600"},
|
||
)
|
||
|
||
try:
|
||
await asyncio.to_thread(_make_thumbnail, source, cache_path)
|
||
except Exception:
|
||
logger.exception(
|
||
"Не удалось создать превью для запроса %s", request_id
|
||
)
|
||
return FileResponse(source)
|
||
|
||
return FileResponse(
|
||
cache_path,
|
||
media_type="image/jpeg",
|
||
headers={"Cache-Control": "public, max-age=3600"},
|
||
)
|
||
|
||
@app.get("/healthz")
|
||
async def healthz() -> dict[str, bool]:
|
||
return {"ok": True}
|
||
|
||
@app.get("/users")
|
||
async def users_redirect() -> RedirectResponse:
|
||
return RedirectResponse("/admin/users", status_code=307)
|
||
|
||
@app.get("/", response_class=HTMLResponse)
|
||
async def index(
|
||
request: Request,
|
||
user_id: str = "",
|
||
status: str = "",
|
||
mode: str = "",
|
||
page: int = 1,
|
||
):
|
||
return await history_page(
|
||
request,
|
||
admin=False,
|
||
user_id=user_id,
|
||
status=status,
|
||
mode=mode,
|
||
page=page,
|
||
)
|
||
|
||
@app.get("/request/{request_id}", response_class=HTMLResponse)
|
||
async def detail(request: Request, request_id: int):
|
||
return await detail_page(request, request_id, admin=False)
|
||
|
||
@app.get("/media/{request_id}/{kind}")
|
||
async def media(request: Request, request_id: int, kind: str):
|
||
return await media_response(request, request_id, kind, admin=False)
|
||
|
||
@app.get("/thumb/{request_id}")
|
||
async def thumbnail(request: Request, request_id: int):
|
||
return await thumbnail_response(request, request_id, admin=False)
|
||
|
||
admin_router = APIRouter(
|
||
prefix="/admin", dependencies=[Depends(require_admin)]
|
||
)
|
||
|
||
@admin_router.get("", response_class=HTMLResponse)
|
||
@admin_router.get("/", response_class=HTMLResponse)
|
||
async def admin_index(
|
||
request: Request,
|
||
user_id: str = "",
|
||
status: str = "",
|
||
mode: str = "",
|
||
page: int = 1,
|
||
):
|
||
return await history_page(
|
||
request,
|
||
admin=True,
|
||
user_id=user_id,
|
||
status=status,
|
||
mode=mode,
|
||
page=page,
|
||
)
|
||
|
||
@admin_router.get("/users", response_class=HTMLResponse)
|
||
async def admin_users(request: Request, q: str = ""):
|
||
db: Database = request.app.state.db
|
||
rows = await db.list_users_admin(1000)
|
||
needle = q.strip().lower().lstrip("@")
|
||
if needle:
|
||
rows = [
|
||
user
|
||
for user in rows
|
||
if needle in str(user.user_id)
|
||
or (user.username and needle in user.username.lower())
|
||
]
|
||
return templates.TemplateResponse(
|
||
request=request,
|
||
name="users.html",
|
||
context={"users": rows, "q": q, "admin": True},
|
||
)
|
||
|
||
@admin_router.post("/users/{user_id}/credits")
|
||
async def admin_user_credits(
|
||
request: Request,
|
||
user_id: int,
|
||
credits: Annotated[int, Form()],
|
||
) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
if not await db.set_credits(user_id, max(0, credits)):
|
||
raise HTTPException(
|
||
status_code=404, detail="Пользователь не найден"
|
||
)
|
||
return RedirectResponse("/admin/users", status_code=303)
|
||
|
||
@admin_router.post("/users/{user_id}/admin")
|
||
async def admin_user_admin(
|
||
request: Request, user_id: int
|
||
) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
user = await db.get_user(user_id)
|
||
if user is None:
|
||
raise HTTPException(
|
||
status_code=404, detail="Пользователь не найден"
|
||
)
|
||
await db.set_admin(user_id, not user.is_admin)
|
||
return RedirectResponse("/admin/users", status_code=303)
|
||
|
||
@admin_router.post("/users/{user_id}/block")
|
||
async def admin_user_block(
|
||
request: Request, user_id: int
|
||
) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
user = await db.get_user(user_id)
|
||
if user is None:
|
||
raise HTTPException(
|
||
status_code=404, detail="Пользователь не найден"
|
||
)
|
||
await db.set_blocked(user_id, not user.is_blocked)
|
||
return RedirectResponse("/admin/users", status_code=303)
|
||
|
||
@admin_router.post("/requests/{request_id}/hide")
|
||
async def admin_hide_request(
|
||
request: Request, request_id: int
|
||
) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
item = await db.get_request(request_id)
|
||
if item is None:
|
||
raise HTTPException(status_code=404, detail="Запрос не найден")
|
||
if not await db.set_hidden(request_id, not item.hidden):
|
||
raise HTTPException(status_code=404, detail="Запрос не найден")
|
||
return RedirectResponse("/admin", status_code=303)
|
||
|
||
@admin_router.get("/request/{request_id}", response_class=HTMLResponse)
|
||
async def admin_detail(request: Request, request_id: int):
|
||
return await detail_page(request, request_id, admin=True)
|
||
|
||
@admin_router.get("/media/{request_id}/{kind}")
|
||
async def admin_media(request: Request, request_id: int, kind: str):
|
||
return await media_response(request, request_id, kind, admin=True)
|
||
|
||
@admin_router.get("/thumb/{request_id}")
|
||
async def admin_thumbnail(request: Request, request_id: int):
|
||
return await thumbnail_response(request, request_id, admin=True)
|
||
|
||
app.include_router(admin_router)
|
||
return app
|
||
|
||
|
||
app = create_app()
|