Stend_mont/app/maintenance.py

import datetime as dt
import fcntl
import logging
import os
import threading
import time

import docker
from sqlalchemy import select

from config import ENABLE_STARTUP_MAINTENANCE, SESSION_IDLE_SECONDS, WEB_POOL_SIZE
from database import Base, SessionLocal, engine
from models import RdpSlot, Service, ServiceType, SessionModel, SessionStatus, User, UserServiceAccess
from utils import ensure_icons_dir, now_utc
from auth import hash_password
from runtime import (
    _rdp_slot_container_name,
    disconnect_rdp_slot,
    docker_client,
    ensure_schema_compatibility,
    ensure_universal_pool,
    ensure_warm_pool,
    ensure_web_pool,
    start_rdp_slot_container,
    stop_runtime_container,
    web_pool_container_name,
)

logger = logging.getLogger("portal")
maintenance_lock_file = None


def cleanup_loop():
    while True:
        time.sleep(60)
        db = SessionLocal()
        try:
            ensure_universal_pool()
            ensure_web_pool()
            for svc in db.scalars(
                select(Service).where(
                    Service.active == True,
                    Service.type.in_([ServiceType.WEB, ServiceType.RDP]),
                )
            ).all():
                if svc.type == ServiceType.WEB and WEB_POOL_SIZE <= 0:
                    ensure_warm_pool(svc)
            cutoff = now_utc() - dt.timedelta(seconds=SESSION_IDLE_SECONDS)
            q = select(SessionModel).where(
                SessionModel.status == SessionStatus.ACTIVE,
                SessionModel.last_access_at < cutoff,
            )
            stale = db.scalars(q).all()
            rdp_slots_to_restart: list[int] = []
            for sess in stale:
                cid = sess.container_id or ""
                if cid.startswith("RDPSLOT:"):
                    try:
                        rdp_slots_to_restart.append(int(cid.split(":", 1)[1]))
                    except Exception:
                        pass
                elif cid.startswith("WEBPOOLIDX:"):
                    try:
                        slot = int(cid.split(":", 1)[1])
                        stop_runtime_container(web_pool_container_name(slot))
                    except Exception:
                        logger.exception("webpool_slot_recycle_failed cid=%s", cid)
                elif cid and not (
                    cid.startswith("POOL:")
                    or cid.startswith("POOLIDX:")
                    or cid.startswith("RDPSLOT:")
                ):
                    stop_runtime_container(cid)
                sess.status = SessionStatus.EXPIRED
            if stale:
                db.commit()
            for slot_id in rdp_slots_to_restart:
                threading.Thread(target=disconnect_rdp_slot, args=(slot_id,), daemon=True).start()

            # Pilots: a slot assigned to a user whose grant for that pilot
            # service has since been revoked or has expired (per-grant
            # UserServiceAccess.expires_at, or the row is just gone) goes
            # back into the pool so an admin can hand it to someone else.
            assigned_slots = db.scalars(
                select(RdpSlot).where(RdpSlot.assigned_user_id.is_not(None))
            ).all()
            if assigned_slots:
                now = now_utc()
                freed = 0
                for slot in assigned_slots:
                    access = db.scalar(
                        select(UserServiceAccess).where(
                            UserServiceAccess.user_id == slot.assigned_user_id,
                            UserServiceAccess.service_id == slot.service_id,
                        )
                    )
                    if access is None or (access.expires_at is not None and access.expires_at <= now):
                        slot.assigned_user_id = None
                        freed += 1
                if freed:
                    db.commit()
                    logger.info("pilot_slots_released count=%s", freed)
        except Exception:
            db.rollback()
            logger.exception("cleanup_loop_failed")
        finally:
            db.close()


def bootstrap_admin():
    admin_user = os.getenv("ADMIN_USERNAME", "admin")
    admin_password = os.getenv("ADMIN_PASSWORD", "change_me")
    ttl_days = int(os.getenv("ADMIN_TTL_DAYS", "3650"))

    db = SessionLocal()
    try:
        existing = db.scalar(select(User).where(User.username == admin_user))
        if not existing:
            db.add(
                User(
                    username=admin_user,
                    password_hash=hash_password(admin_password),
                    active=True,
                    is_admin=True,
                    expires_at=now_utc() + dt.timedelta(days=ttl_days),
                )
            )
            db.commit()
    finally:
        db.close()


def try_acquire_maintenance_leader() -> bool:
    global maintenance_lock_file
    if maintenance_lock_file is not None:
        return True
    lock_file = open("/tmp/portal-maintenance.lock", "w")
    try:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        lock_file.close()
        return False
    maintenance_lock_file = lock_file
    return True


def _init_schema() -> None:
    with open("/tmp/portal-schema.lock", "w") as lock_file:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        Base.metadata.create_all(bind=engine)
        ensure_schema_compatibility()
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
    ensure_icons_dir()
    bootstrap_admin()


def _startup_pools(db) -> None:
    ensure_universal_pool()
    ensure_web_pool()
    for svc in db.scalars(
        select(Service).where(
            Service.active == True,
            Service.type.in_([ServiceType.WEB, ServiceType.RDP]),
        )
    ).all():
        if svc.type == ServiceType.WEB and WEB_POOL_SIZE <= 0:
            ensure_warm_pool(svc)
        elif svc.type == ServiceType.RDP:
            slots = db.scalars(select(RdpSlot).where(RdpSlot.service_id == svc.id)).all()
            for slot in slots:
                try:
                    cname = _rdp_slot_container_name(svc.slug, slot.id)
                    try:
                        c = docker_client().containers.get(cname)
                        if c.status != "running":
                            c.start()
                    except docker.errors.NotFound:
                        start_rdp_slot_container(slot, svc)
                        slot.container_name = cname
                except Exception:
                    logger.exception("startup_rdp_slot_start_failed slot_id=%s", slot.id)
            if slots:
                db.commit()


def run_maintenance_service() -> None:
    logger.info("maintenance_service_bootstrap_started")
    _init_schema()

    maintenance_lock = open("/tmp/portal-maintenance.lock", "w")
    fcntl.flock(maintenance_lock.fileno(), fcntl.LOCK_EX)
    logger.info("maintenance_service_leader_acquired")

    db = SessionLocal()
    try:
        _startup_pools(db)
    finally:
        db.close()

    logger.info("maintenance_service_loop_started")
    cleanup_loop()


def on_startup() -> None:
    _init_schema()
    if not try_acquire_maintenance_leader():
        logger.info("maintenance_leader_skipped")
        return

    if ENABLE_STARTUP_MAINTENANCE:
        db = SessionLocal()
        try:
            _startup_pools(db)
        finally:
            db.close()

    thread = threading.Thread(target=cleanup_loop, daemon=True)
    thread.start()
    logger.info("maintenance_leader_started")