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
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,
)

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 and not (
                    cid.startswith("POOL:")
                    or cid.startswith("POOLIDX:")
                    or cid.startswith("WEBPOOLIDX:")
                ):
                    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()
        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 run_maintenance_service() -> None:
    logger.info("maintenance_service_bootstrap_started")
    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()

    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:
        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()
    finally:
        db.close()

    logger.info("maintenance_service_loop_started")
    cleanup_loop()


def on_startup() -> 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()
    if not try_acquire_maintenance_leader():
        logger.info("maintenance_leader_skipped")
        return

    if ENABLE_STARTUP_MAINTENANCE:
        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)
                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()
        finally:
            db.close()

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