app.py - master - Befun/FamalyControlle

app.py
454 строки18 KB
BeFun
Init project
07 май 2026, 10:43
07 май 2026, 10:437e96900
"""
API-сервер (FastAPI).
Предоставляет REST-эндпоинты для управления ограничениями,
расписаниями и пользователями, а также WebSocket для real-time.
"""

import asyncio
import json
import logging
from datetime import datetime, timedelta
from typing import Optional

from fastapi import FastAPI, HTTPException, Header, WebSocket, WebSocketDisconnect, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from core.config import Config
from core.permissions import Role, PermissionChecker
from server.database import Database

logger = logging.getLogger("server")

# ═══════════════════════════════════════════════════════════════
#  Инициализация
# ═══════════════════════════════════════════════════════════════

Config.load()
db = Database(Config.database_url)

app = FastAPI(
    title="Family PC Control — API",
    description="Сервер управления семейным контролем ПК",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# WebSocket менеджер
class ConnectionManager:
    def __init__(self):
        self.active: list[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.active.append(ws)

    def disconnect(self, ws: WebSocket):
        if ws in self.active:
            self.active.remove(ws)

    async def broadcast(self, message: dict):
        dead = []
        for ws in self.active:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)

ws_manager = ConnectionManager()

# ═══════════════════════════════════════════════════════════════
#  Pydantic схемы
# ═══════════════════════════════════════════════════════════════

class RegisterRequest(BaseModel):
    max_user_id: str
    display_name: str = ""
    role: str = "child"

class IdentifyRequest(BaseModel):
    max_user_id: str

class RestrictionCreate(BaseModel):
    category: str      # games, browsers, downloads, app
    target: str = "*"  # * = все в категории, или конкретное имя процесса
    expires_minutes: Optional[int] = None

class ScheduleCreate(BaseModel):
    name: str
    schedule_type: str  # sleep, rest
    start_time: str     # "22:00"
    end_time: str       # "07:00"
    days: list[int]     # [0,1,2,3,4]

class ScheduleUpdate(BaseModel):
    name: Optional[str] = None
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    days: Optional[list[int]] = None
    is_active: Optional[bool] = None

class RoleUpdate(BaseModel):
    role: str

class SystemCommand(BaseModel):
    delay_seconds: int = 60

class AgentEvent(BaseModel):
    event: str
    details: str = ""
    timestamp: str = ""

# ═══════════════════════════════════════════════════════════════
#  Вспомогательные функции
# ═══════════════════════════════════════════════════════════════

def get_current_user(x_user_id: Optional[str] = Header(None)) -> int:
    """Извлечь user_id из заголовка X-User-ID."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="Заголовок X-User-ID обязателен")
    try:
        return int(x_user_id)
    except ValueError:
        raise HTTPException(status_code=400, detail="X-User-ID должен быть числом")


def require_role(user_id: int, *allowed_roles: Role) -> dict:
    """Проверить роль пользователя, вернув словарь с данными."""
    user = db.get_user_by_id(user_id)
    if not user:
        raise HTTPException(status_code=404, detail="Пользователь не найден")
    user_role = PermissionChecker.get_role(user.role)
    if user_role not in allowed_roles:
        raise HTTPException(
            status_code=403,
            detail=f"Недостаточно прав. Требуется: {', '.join(r.value for r in allowed_roles)}",
        )
    return {"user": user, "role": user_role}


def admin_required(user_id: int = Depends(get_current_user)) -> dict:
    return require_role(user_id, Role.ADMIN)


def family_or_admin(user_id: int = Depends(get_current_user)) -> dict:
    return require_role(user_id, Role.FAMILY, Role.ADMIN)


# ═══════════════════════════════════════════════════════════════
#  Auth endpoints
# ═══════════════════════════════════════════════════════════════

@app.post("/api/auth/register")
def register(req: RegisterRequest):
    """Регистрация нового пользователя."""
    existing = db.get_user_by_max_id(req.max_user_id)
    if existing:
        return {
            "ok": True,
            "message": "Пользователь уже существует",
            "user": existing.to_dict(),
        }

    # Первый зарегистрированный пользователь становится админом
    all_users = db.get_all_users()
    role = "admin" if len(all_users) == 0 else req.role

    # Проверяем, совпадает ли с админом из конфига
    if req.max_user_id == Config.admin_max_user_id:
        role = "admin"

    user = db.create_user(req.max_user_id, req.display_name, role)
    logger.info("Зарегистрирован пользователь: %s (роль: %s)", req.max_user_id, role)
    return {"ok": True, "user": user.to_dict()}


@app.post("/api/auth/identify")
def identify(req: IdentifyRequest):
    """Идентификация по max_user_id."""
    user = db.get_user_by_max_id(req.max_user_id)
    if not user:
        return {"ok": False, "user": None}
    return {"ok": True, "user": user.to_dict()}


# ═══════════════════════════════════════════════════════════════
#  Users endpoints
# ═══════════════════════════════════════════════════════════════

@app.get("/api/users")
def list_users():
    """Список всех активных пользователей."""
    users = db.get_all_users()
    return {"users": [u.to_dict() for u in users]}


@app.put("/api/users/{user_id}/role")
def change_role(user_id: int, body: RoleUpdate, _=Depends(admin_required)):
    """Изменить роль пользователя (только админ)."""
    if body.role not in PermissionChecker.get_all_roles():
        raise HTTPException(status_code=400, detail=f"Недопустимая роль: {body.role}")

    user = db.update_user_role(user_id, body.role)
    if not user:
        raise HTTPException(status_code=404, detail="Пользователь не найден")
    return {"ok": True, "user": user.to_dict()}


# ═══════════════════════════════════════════════════════════════
#  Restrictions endpoints
# ═══════════════════════════════════════════════════════════════

@app.get("/api/restrictions/active")
def get_active_restrictions():
    """Все активные ограничения."""
    restrictions = db.get_active_restrictions()
    return {"restrictions": [r.to_dict() for r in restrictions]}


@app.post("/api/restrictions")
def create_restriction(body: RestrictionCreate, info=Depends(family_or_admin)):
    """Создать ограничение."""
    valid_categories = {"games", "browsers", "downloads", "app"}
    if body.category not in valid_categories:
        raise HTTPException(status_code=400, detail=f"Категория должна быть одна из: {valid_categories}")

    r = db.create_restriction(
        category=body.category,
        target=body.target,
        created_by=info["user"].id,
        expires_minutes=body.expires_minutes,
    )

    # Уведомить через WebSocket
    asyncio.create_task(ws_manager.broadcast({
        "type": "restriction_created",
        "data": r.to_dict(),
    }))

    return {"ok": True, "restriction": r.to_dict()}


@app.delete("/api/restrictions/{restriction_id}")
def remove_restriction(restriction_id: int, info=Depends(family_or_admin)):
    """Отменить ограничение (с проверкой прав)."""
    restriction = db.get_restriction(restriction_id)
    if not restriction:
        raise HTTPException(status_code=404, detail="Ограничение не найдено")

    if not restriction.is_active:
        raise HTTPException(status_code=400, detail="Ограничение уже снято")

    # Проверяем, кто создал ограничение
    creator = db.get_user_by_id(restriction.created_by)
    if creator:
        creator_role = PermissionChecker.get_role(creator.role)
        actor_role = info["role"]
        if not PermissionChecker.can_unblock(actor_role, creator_role):
            raise HTTPException(
                status_code=403,
                detail=f"Вы не можете отменить ограничение, созданное {creator_role.value}",
            )

    r = db.remove_restriction(restriction_id, info["user"].id)
    if not r:
        raise HTTPException(status_code=500, detail="Ошибка снятия ограничения")

    asyncio.create_task(ws_manager.broadcast({
        "type": "restriction_removed",
        "data": {"id": restriction_id},
    }))

    return {"ok": True, "message": "Ограничение снято"}


@app.get("/api/restrictions/history")
def get_history(limit: int = 50, offset: int = 0):
    """История действий (пагинация)."""
    entries = db.get_history(limit=limit, offset=offset)
    return {
        "history": [e.to_dict() for e in entries],
        "total": len(entries),
        "offset": offset,
        "limit": limit,
    }


# ═══════════════════════════════════════════════════════════════
#  Schedules endpoints
# ═══════════════════════════════════════════════════════════════

@app.get("/api/schedules")
def list_schedules():
    """Список активных расписаний."""
    schedules = db.get_schedules()
    return {"schedules": [s.to_dict() for s in schedules]}


@app.post("/api/schedules")
def create_schedule(body: ScheduleCreate, info=Depends(family_or_admin)):
    """Создать расписание."""
    if body.schedule_type not in ("sleep", "rest"):
        raise HTTPException(status_code=400, detail="Тип расписания: sleep или rest")
    if not body.days:
        raise HTTPException(status_code=400, detail="Укажите дни расписания")

    sch = db.create_schedule(
        name=body.name,
        schedule_type=body.schedule_type,
        start_time=body.start_time,
        end_time=body.end_time,
        days=body.days,
        created_by=info["user"].id,
    )

    asyncio.create_task(ws_manager.broadcast({
        "type": "schedule_created",
        "data": sch.to_dict(),
    }))

    return {"ok": True, "schedule": sch.to_dict()}


@app.put("/api/schedules/{schedule_id}")
def update_schedule(schedule_id: int, body: ScheduleUpdate, info=Depends(family_or_admin)):
    """Обновить расписание."""
    updates = body.model_dump(exclude_none=True)
    if not updates:
        raise HTTPException(status_code=400, detail="Нет данных для обновления")

    sch = db.update_schedule(schedule_id, **updates)
    if not sch:
        raise HTTPException(status_code=404, detail="Расписание не найдено")

    asyncio.create_task(ws_manager.broadcast({
        "type": "schedule_updated",
        "data": sch.to_dict(),
    }))

    return {"ok": True, "schedule": sch.to_dict()}


@app.delete("/api/schedules/{schedule_id}")
def delete_schedule(schedule_id: int, info=Depends(family_or_admin)):
    """Удалить расписание."""
    ok = db.remove_schedule(schedule_id, info["user"].id)
    if not ok:
        raise HTTPException(status_code=404, detail="Расписание не найдено")

    asyncio.create_task(ws_manager.broadcast({
        "type": "schedule_removed",
        "data": {"id": schedule_id},
    }))

    return {"ok": True, "message": "Расписание удалено"}


# ═══════════════════════════════════════════════════════════════
#  System endpoints
# ═══════════════════════════════════════════════════════════════

@app.get("/api/system/status")
def system_status():
    """Статус системы."""
    return {
        "status": "running",
        "server_time": datetime.utcnow().isoformat() + "Z",
        "version": "1.0.0",
        "restrictions_count": len(db.get_active_restrictions()),
        "schedules_count": len(db.get_schedules()),
        "users_count": len(db.get_all_users()),
    }


@app.post("/api/system/shutdown")
def system_shutdown(body: SystemCommand, info=Depends(admin_required)):
    """Выключить ПК (только админ)."""
    from utils.helpers import shutdown_pc
    ok = shutdown_pc(body.delay_seconds)
    if not ok:
        raise HTTPException(status_code=500, detail="Не удалось инициировать выключение")

    asyncio.create_task(ws_manager.broadcast({
        "type": "system_shutdown",
        "data": {"delay": body.delay_seconds, "initiated_by": info["user"].id},
    }))

    return {"ok": True, "message": f"Выключение через {body.delay_seconds} сек"}


@app.post("/api/system/restart")
def system_restart(body: SystemCommand, info=Depends(family_or_admin)):
    """Перезагрузить ПК."""
    from utils.helpers import restart_pc
    ok = restart_pc(body.delay_seconds)
    if not ok:
        raise HTTPException(status_code=500, detail="Не удалось инициировать перезагрузку")

    asyncio.create_task(ws_manager.broadcast({
        "type": "system_restart",
        "data": {"delay": body.delay_seconds, "initiated_by": info["user"].id},
    }))

    return {"ok": True, "message": f"Перезагрузка через {body.delay_seconds} сек"}


# ═══════════════════════════════════════════════════════════════
#  Agent events endpoint
# ═══════════════════════════════════════════════════════════════

@app.post("/api/agent/event")
def agent_event(body: AgentEvent):
    """Получить событие от десктоп-агента и разослать через WebSocket."""
    asyncio.create_task(ws_manager.broadcast({
        "type": "agent_event",
        "data": body.model_dump(),
    }))
    return {"ok": True}


# ═══════════════════════════════════════════════════════════════
#  WebSocket
# ═══════════════════════════════════════════════════════════════

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    """WebSocket для real-time уведомлений."""
    await ws_manager.connect(ws)
    try:
        while True:
            data = await ws.receive_text()
            # Пинг-понг для поддержания соединения
            if data == "ping":
                await ws.send_text("pong")
    except WebSocketDisconnect:
        ws_manager.disconnect(ws)


# ═══════════════════════════════════════════════════════════════
#  Запуск сервера
# ═══════════════════════════════════════════════════════════════

def run_server():
    """Точка входа для API-сервера."""
    import uvicorn
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    logger.info("API-сервер запускается на %s:%d", Config.server_host, Config.server_port)
    uvicorn.run(
        app,
        host=Config.server_host,
        port=Config.server_port,
        log_level="info",
    )