Nexus/server/infrastructure/btpanel/bootstrap_state.py

"""Baota SSH bootstrap state in ``servers.extra_attrs.bt_panel``."""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Any

from server.infrastructure.btpanel.credentials import BT_PANEL_ATTR, bt_panel_configured

BOOTSTRAP_RETRY_SECONDS = 300
STATE_PENDING = "pending"
STATE_INSTALLING = "installing"
STATE_READY = "ready"
STATE_FAILED = "failed"
STATE_DISABLED = "disabled"
RETRYABLE_STATES = frozenset({STATE_PENDING, STATE_INSTALLING})
PERMANENT_STOP_STATES = frozenset({STATE_READY, STATE_FAILED, STATE_DISABLED})


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _iso(dt: datetime) -> str:
    return dt.replace(microsecond=0).isoformat().replace("+00:00", "Z")


def _parse_iso(value: Any) -> datetime | None:
    if not value or not isinstance(value, str):
        return None
    text = value.strip()
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        parsed = datetime.fromisoformat(text)
        if parsed.tzinfo is None:
            parsed = parsed.replace(tzinfo=timezone.utc)
        return parsed
    except ValueError:
        return None


def get_bootstrap_block(server) -> dict[str, Any]:
    from server.infrastructure.btpanel.credentials import get_bt_panel_block

    block = get_bt_panel_block(server)
    if not block:
        return {
            "auto_bootstrap": None,
            "bootstrap_state": None,
            "bootstrap_last_attempt_at": None,
            "bootstrap_next_attempt_at": None,
            "bootstrap_attempts": 0,
            "bootstrap_last_error": None,
            "bootstrap_last_actions": None,
        }
    return {
        "auto_bootstrap": block.get("auto_bootstrap", True) is not False,
        "bootstrap_state": block.get("bootstrap_state"),
        "bootstrap_last_attempt_at": block.get("bootstrap_last_attempt_at"),
        "bootstrap_next_attempt_at": block.get("bootstrap_next_attempt_at"),
        "bootstrap_attempts": int(block.get("bootstrap_attempts") or 0),
        "bootstrap_last_error": block.get("bootstrap_last_error"),
        "bootstrap_last_actions": block.get("bootstrap_last_actions"),
    }


def public_bootstrap_fields(server) -> dict[str, Any]:
    b = get_bootstrap_block(server)
    return {
        "auto_bootstrap": b["auto_bootstrap"],
        "bootstrap_state": b["bootstrap_state"],
        "bootstrap_last_attempt_at": b["bootstrap_last_attempt_at"],
        "bootstrap_next_attempt_at": b["bootstrap_next_attempt_at"],
        "bootstrap_attempts": b["bootstrap_attempts"],
        "bootstrap_last_error": b["bootstrap_last_error"],
        "bootstrap_last_actions": b["bootstrap_last_actions"],
    }


def merge_bootstrap_fields(
    extra_attrs: dict[str, Any] | None,
    **fields: Any,
) -> dict[str, Any]:
    out = dict(extra_attrs or {})
    block = dict(out.get(BT_PANEL_ATTR) or {}) if isinstance(out.get(BT_PANEL_ATTR), dict) else {}
    for key, value in fields.items():
        if value is None:
            block.pop(key, None)
        else:
            block[key] = value
    if block:
        out[BT_PANEL_ATTR] = block
    return out


def mark_bootstrap_pending(extra_attrs: dict[str, Any] | None) -> dict[str, Any]:
    if bt_panel_configured_from_attrs(extra_attrs):
        return extra_attrs or {}
    now = _iso(_utcnow())
    return merge_bootstrap_fields(
        extra_attrs,
        auto_bootstrap=True,
        bootstrap_state=STATE_PENDING,
        bootstrap_last_attempt_at=None,
        bootstrap_next_attempt_at=now,
        bootstrap_last_error=None,
    )


def bt_panel_configured_from_attrs(extra_attrs: dict[str, Any] | None) -> bool:
    if not isinstance(extra_attrs, dict):
        return False
    block = extra_attrs.get(BT_PANEL_ATTR)
    if not isinstance(block, dict):
        return False
    return bool(str(block.get("base_url") or "").strip() and block.get("api_key_enc"))


def is_eligible_for_background_bootstrap(server, *, now: datetime | None = None) -> bool:
    if bt_panel_configured(server):
        return False
    from server.infrastructure.btpanel.credentials import get_bt_panel_block

    block = get_bt_panel_block(server)
    if not block:
        return False
    if block.get("auto_bootstrap") is False:
        return False
    state = block.get("bootstrap_state")
    if state not in RETRYABLE_STATES:
        return False
    next_at = _parse_iso(block.get("bootstrap_next_attempt_at"))
    ref = now or _utcnow()
    if next_at is None:
        return False
    return next_at <= ref


def apply_bootstrap_success(
    extra_attrs: dict[str, Any] | None,
    *,
    base_url: str,
    api_key: str,
    verify_ssl: bool,
    actions: list[str],
) -> dict[str, Any]:
    from server.infrastructure.btpanel.credentials import merge_bt_panel_extra

    now = _iso(_utcnow())
    out = merge_bt_panel_extra(
        extra_attrs,
        base_url=base_url,
        api_key=api_key,
        verify_ssl=verify_ssl,
    )
    attempts = int((get_bootstrap_block_from_attrs(out)).get("bootstrap_attempts") or 0) + 1
    return merge_bootstrap_fields(
        out,
        bootstrap_state=STATE_READY,
        bootstrap_last_attempt_at=now,
        bootstrap_next_attempt_at=None,
        bootstrap_attempts=attempts,
        bootstrap_last_error=None,
        bootstrap_last_actions=actions,
        bt_installed=True,
        bt_installed_checked_at=now,
    )


def apply_bootstrap_failure(
    extra_attrs: dict[str, Any] | None,
    *,
    error_code: str,
    permanent: bool,
) -> dict[str, Any]:
    now = _utcnow()
    attempts = int((get_bootstrap_block_from_attrs(extra_attrs)).get("bootstrap_attempts") or 0) + 1
    state = STATE_FAILED if permanent else (
        STATE_INSTALLING if error_code == "bt_not_installed" else STATE_PENDING
    )
    next_attempt = None if permanent or state == STATE_READY else _iso(now + timedelta(seconds=BOOTSTRAP_RETRY_SECONDS))
    cache_fields: dict[str, Any] = {}
    if error_code == "bt_not_installed":
        cache_fields["bt_installed"] = False
        cache_fields["bt_installed_checked_at"] = _iso(now)
    return merge_bootstrap_fields(
        extra_attrs,
        bootstrap_state=state,
        bootstrap_last_attempt_at=_iso(now),
        bootstrap_next_attempt_at=next_attempt,
        bootstrap_attempts=attempts,
        bootstrap_last_error=error_code,
        **cache_fields,
    )


def get_bootstrap_block_from_attrs(extra_attrs: dict[str, Any] | None) -> dict[str, Any]:
    if not isinstance(extra_attrs, dict):
        return {}
    block = extra_attrs.get(BT_PANEL_ATTR)
    return block if isinstance(block, dict) else {}