workshop/routers.py

import io
import os
from pathlib import PurePosixPath
from uuid import uuid4
from zipfile import BadZipFile, ZipFile

import boto3
from botocore.config import Config
from botocore.exceptions import BotoCoreError, ClientError
from fastapi import APIRouter, File, Form, HTTPException, UploadFile, status
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel, ConfigDict
from sqlalchemy import or_
from sqlalchemy.orm import Session

from database import Service, engine


CATEGORIES = (
    "Web",
    "Binary",
    "Network",
    "Database",
    "Authentication",
    "Storage",
    "Crypto",
    "Other",
)
MAX_ARCHIVE_BYTES = 5_000_000
MAX_FILES = 512
MAX_UNPACKED_BYTES = 20_000_000
REQUIRED_ARCHIVE_FILES = {
    "checker/checker.py",
    "service/docker-compose.yml",
    "service/Dockerfile.api",
}

router = APIRouter(prefix="/api", tags=["workshop"])
s3 = boto3.client(
    "s3",
    endpoint_url=os.getenv("S3_ENDPOINT_URL"),
    aws_access_key_id=os.getenv("S3_ACCESS_KEY"),
    aws_secret_access_key=os.getenv("S3_SECRET_KEY"),
    region_name=os.getenv("S3_REGION", "us-east-1"),
    config=Config(
        proxies={},
        signature_version="s3v4",
        s3={"addressing_style": "path"},
    ),
)


def remove_expect_header(request, **_):
    request.headers.pop("Expect", None)


s3.meta.events.register("before-send.s3.PutObject", remove_expect_header)
S3_BUCKET = os.getenv("S3_BUCKET", "workshop")


class ServiceOut(BaseModel):
    id: int
    name: str
    description: str
    category: str
    author: str
    selected: bool

    model_config = ConfigDict(from_attributes=True)


class SelectionIn(BaseModel):
    selected: bool


def get_service(db: Session, service_id: int) -> Service:
    service = db.get(Service, service_id)
    if not service:
        raise HTTPException(status_code=404, detail="Сервис не найден")
    return service


def safe_archive_path(path: str) -> bool:
    normalized = path.replace("\\", "/")
    return bool(path) and not normalized.startswith("/") and ".." not in PurePosixPath(normalized).parts


@router.get("/categories", response_model=list[str])
def list_categories():
    return CATEGORIES


@router.get("/services", response_model=list[ServiceOut])
def list_services(
    category: str | None = None,
    q: str | None = None,
    selected: bool | None = None,
):
    with Session(engine) as db:
        query = db.query(Service)
        if category:
            query = query.filter(Service.category == category)
        if q and (query_text := q.strip()):
            pattern = "%" + query_text + "%"
            query = query.filter(
                or_(
                    Service.name.ilike(pattern),
                    Service.description.ilike(pattern),
                    Service.author.ilike(pattern),
                )
            )
        if selected is not None:
            query = query.filter(Service.selected == selected)
        return query.order_by(Service.id.desc()).all()


@router.post("/services", response_model=ServiceOut, status_code=status.HTTP_201_CREATED)
async def create_service(
    file: UploadFile = File(...),
    name: str = Form(..., max_length=100),
    description: str = Form("", max_length=2000),
    category: str = Form("Other", max_length=50),
    author: str = Form("anonymous", max_length=100),
):
    name, description, category, author = (
        name.strip(),
        description.strip(),
        category.strip(),
        author.strip() or "anonymous",
    )
    if not name:
        raise HTTPException(status_code=422, detail="Укажите название сервиса")
    if category not in CATEGORIES:
        raise HTTPException(status_code=422, detail="Выберите категорию из списка")
    if not file.filename or not file.filename.lower().endswith(".zip"):
        raise HTTPException(status_code=415, detail="Нужен ZIP-архив сервиса")

    archive_data = await file.read(MAX_ARCHIVE_BYTES + 1)
    await file.close()
    if len(archive_data) > MAX_ARCHIVE_BYTES:
        raise HTTPException(status_code=413, detail="ZIP больше 5 МБ")

    try:
        with ZipFile(io.BytesIO(archive_data)) as archive:
            entries = [entry for entry in archive.infolist() if not entry.is_dir()]
            if not entries:
                raise HTTPException(status_code=422, detail="В ZIP нет файлов")
            if len(entries) > MAX_FILES:
                raise HTTPException(status_code=413, detail="В ZIP больше 512 файлов")
            if sum(entry.file_size for entry in entries) > MAX_UNPACKED_BYTES:
                raise HTTPException(status_code=413, detail="Распакованный сервис больше 20 МБ")
            if any(not safe_archive_path(entry.filename) for entry in entries):
                raise HTTPException(status_code=422, detail="ZIP содержит небезопасный путь")

            paths = {entry.filename.replace("\\", "/") for entry in entries}
            missing = REQUIRED_ARCHIVE_FILES - paths
            if missing:
                raise HTTPException(
                    status_code=422,
                    detail="В ZIP нет обязательных файлов: " + ", ".join(sorted(missing)),
                )
            if any(not path.startswith(("checker/", "service/")) for path in paths):
                raise HTTPException(
                    status_code=422,
                    detail="Файлы ZIP должны лежать только в checker/ или service/",
                )

            archive_key = f"services/{uuid4()}/service.zip"
            try:
                await run_in_threadpool(
                    s3.put_object,
                    Bucket=S3_BUCKET,
                    Key=archive_key,
                    Body=archive_data,
                    ContentType="application/zip",
                )
            except (BotoCoreError, ClientError) as error:
                print(f"S3 upload error: {error}")
                raise HTTPException(status_code=503, detail="S3-хранилище недоступно")

            with Session(engine) as db:
                service = Service(
                    name=name,
                    description=description,
                    category=category,
                    author=author,
                    archive_key=archive_key,
                )
                db.add(service)
                db.commit()
                db.refresh(service)
                return service
    except BadZipFile:
        raise HTTPException(status_code=422, detail="Файл не является корректным ZIP-архивом")


@router.patch("/services/{service_id}/selection", response_model=ServiceOut)
def update_selection(service_id: int, payload: SelectionIn):
    with Session(engine) as db:
        service = get_service(db, service_id)
        service.selected = payload.selected
        db.commit()
        db.refresh(service)
        return service