zerto-ai-rewind/src/zerto_rewind_mcp/recover.py

"""FLR session helpers."""

from __future__ import annotations

import asyncio
import urllib.parse
from typing import Any

from zerto_rewind_mcp.client import ZertoClient, ZertoError
from zerto_rewind_mcp.util import pick

READY_STATUSES = {
    "ready",
    "mounted",
    "available",
    "mountcompletedsuccessfully",
    "completed",
}
FAILED_STATUSES = {"failed", "error", "mountfailed"}


def session_id_from(payload: Any) -> str:
    if payload is None:
        raise ZertoError("FLR start returned empty body")
    if isinstance(payload, str):
        return payload.strip().strip('"')
    if isinstance(payload, dict):
        value = pick(
            payload,
            "FlrSessionIdentifier",
            "sessionId",
            "SessionId",
            "flrSessionIdentifier",
            "identifier",
            "Identifier",
            "id",
            "Id",
        )
        if value:
            return str(value)
    raise ZertoError(f"Could not read FLR session id from {payload!r}")


def download_token_from(payload: Any) -> str:
    if isinstance(payload, str):
        return payload.strip().strip('"')
    if isinstance(payload, dict):
        value = pick(payload, "token", "Token", "downloadToken", "DownloadToken")
        if value:
            return str(value)
    raise ZertoError(f"Could not read FLR download token from {payload!r}")


def flr_row(payload: Any) -> dict[str, Any]:
    if isinstance(payload, list):
        if not payload:
            return {}
        first = payload[0]
        return first if isinstance(first, dict) else {}
    if isinstance(payload, dict):
        return payload
    return {}


def flr_status(row: dict[str, Any]) -> str:
    value = pick(row, "FlrSessionStatus", "Status", "status", "state", "State") or ""
    return str(value)


async def wait_flr_ready(
    client: ZertoClient,
    session_id: str,
    *,
    timeout_s: float = 180.0,
    interval_s: float = 3.0,
) -> dict[str, Any]:
    deadline = asyncio.get_event_loop().time() + timeout_s
    last: Any = None
    while asyncio.get_event_loop().time() < deadline:
        last = await client.get_flr(session_id)
        row = flr_row(last)
        status = flr_status(row)
        low = status.lower()
        if low in READY_STATUSES:
            return row
        if "mountinprogress" in low or (low.endswith("inprogress") and "fail" not in low):
            await asyncio.sleep(interval_s)
            continue
        if low in FAILED_STATUSES:
            raise ZertoError(f"FLR session {session_id} failed: {last}")
        if row and not status:
            return row
        await asyncio.sleep(interval_s)
    raise ZertoError(
        f"FLR session {session_id} not ready within {timeout_s:.0f}s (last={last!r}). "
        "Retry browse after a minute, or check EJC is not running and no clone/test is active."
    )


def _decode_flr_path(value: Any) -> str:
    """Browse returns paths percent-encoded (%2f). Download wants them decoded."""
    return urllib.parse.unquote(str(value or "")).replace("\\", "/")


def path_items(payload: Any) -> list[dict[str, Any]]:
    if isinstance(payload, dict):
        items = payload.get("PathItems") or payload.get("pathItems") or []
        return [i for i in items if isinstance(i, dict)]
    if isinstance(payload, list):
        return [i for i in payload if isinstance(i, dict)]
    return []


async def browsable_partitions(client: ZertoClient, session_id: str) -> list[str]:
    """FLR is rooted at partitions (Volume2-Ext4), not the guest's /.

    Volume1-Unknown and friends report IsBrowsable false and cannot be restored.
    """
    rows = path_items(await client.browse_flr(session_id, path=""))
    return [str(r.get("Path")) for r in rows if r.get("IsBrowsable")]


async def resolve_flr_path(client: ZertoClient, session_id: str, guest_path: str) -> str:
    """Map a guest absolute path to the FLR namespace path the download API accepts.

    /home/justin/app-config.yaml -> Volume2-Ext4/home/justin/app-config.yaml
    """
    rel = guest_path.replace("\\", "/").strip("/")
    if not rel:
        raise ZertoError("Empty guest_path")
    parts = rel.split("/")
    name, parent = parts[-1], "/".join(parts[:-1])

    partitions = await browsable_partitions(client, session_id)
    if not partitions:
        raise ZertoError(
            "FLR mounted but no browsable partition. Unsupported partition type "
            "(LVM/unknown) cannot be restored by FLR."
        )
    tried = []
    for vol in partitions:
        probe = f"{vol}/{parent}" if parent else vol
        tried.append(probe)
        try:
            rows = path_items(await client.browse_flr(session_id, path=probe))
        except ZertoError:
            continue
        for row in rows:
            decoded = _decode_flr_path(row.get("Path"))
            if decoded.rsplit("/", 1)[-1] == name:
                return decoded
    raise ZertoError(
        f"{guest_path!r} not found in the FLR mount. Looked under {tried}. "
        "The file may not have replicated into that checkpoint yet."
    )


def session_rows(payload: Any) -> list[dict[str, Any]]:
    """Normalise GET /v1/flrs, which returns a list or a single object."""
    if isinstance(payload, list):
        return [r for r in payload if isinstance(r, dict)]
    if isinstance(payload, dict):
        return [payload]
    return []


def session_id_of(row: dict[str, Any]) -> str:
    value = pick(
        row,
        "FlrSessionIdentifier",
        "flrSessionIdentifier",
        "SessionId",
        "sessionId",
        "Identifier",
        "identifier",
    )
    return str(value) if value is not None else ""


def session_summary(row: dict[str, Any]) -> dict[str, Any]:
    return {
        "session_id": session_id_of(row),
        "state": flr_status(row),
        "vpg_name": pick(row, "VpgName", "vpgName"),
        "vm_name": pick(row, "VmName", "vmName"),
        "checkpoint_id": pick(row, "CheckpointIdentifier", "checkpointIdentifier"),
        "mounted_at": pick(row, "MountedTime", "mountedTime", "StartTime", "startTime"),
    }


def is_live_session(row: dict[str, Any]) -> bool:
    """A session still holding a mount on the recovery site.

    Unmounted/ended sessions linger in GET /v1/flrs as history; they hold nothing.
    """
    state = flr_status(row).lower()
    if not state:
        return False
    return "unmount" not in state and "fail" not in state and "end" not in state