"""ZVM / ZCA REST client. Same paths on both (zerto_api_lessons).""" from __future__ import annotations import time from typing import Any from urllib.parse import urljoin import httpx from zerto_rewind_mcp.util import pick class ZertoError(Exception): def __init__(self, message: str, status_code: int | None = None, body: str | None = None): super().__init__(message) self.status_code = status_code self.body = body class ZertoClient: def __init__( self, base_url: str, username: str, password: str, client_id: str = "zerto-client", verify_tls: bool = False, timeout: float = 60.0, ): self.base_url = base_url.rstrip("/") self.username = username self.password = password self.client_id = client_id self._token: str | None = None self._token_exp = 0.0 self._http = httpx.AsyncClient(verify=verify_tls, timeout=timeout) async def aclose(self) -> None: await self._http.aclose() def _url(self, path: str) -> str: if path.startswith("http"): return path return urljoin(self.base_url + "/", path.lstrip("/")) async def ensure_token(self) -> str: if self._token and time.time() < self._token_exp - 30: return self._token url = self._url("/auth/realms/zerto/protocol/openid-connect/token") response = await self._http.post( url, data={ "grant_type": "password", "username": self.username, "password": self.password, "client_id": self.client_id, "scope": "openid", }, headers={"Content-Type": "application/x-www-form-urlencoded"}, ) if response.status_code >= 400: raise ZertoError( f"Keycloak token failed HTTP {response.status_code}. " "Check username/password/client_id " "(10.x uses zerto-client; 9.x may use zerto-api).", status_code=response.status_code, body=response.text[:500], ) payload = response.json() self._token = payload["access_token"] self._token_exp = time.time() + float(payload.get("expires_in") or 60) return self._token async def request( self, method: str, path: str, *, params: dict[str, Any] | None = None, json_body: Any = None, ) -> httpx.Response: token = await self.ensure_token() headers = {"Authorization": f"Bearer {token}"} if json_body is not None: headers["Content-Type"] = "application/json" response = await self._http.request( method, self._url(path), params=params, json=json_body, headers=headers, ) if response.status_code == 401: self._token = None token = await self.ensure_token() headers["Authorization"] = f"Bearer {token}" response = await self._http.request( method, self._url(path), params=params, json=json_body, headers=headers, ) return response async def json( self, method: str, path: str, *, params: dict[str, Any] | None = None, json_body: Any = None, ) -> Any: response = await self.request(method, path, params=params, json_body=json_body) if response.status_code >= 400: raise ZertoError( f"{method} {path} HTTP {response.status_code}: {response.text[:400]}", status_code=response.status_code, body=response.text[:1000], ) if response.status_code == 204 or not response.content: return None try: return response.json() except ValueError: return response.text async def get_vms( self, *, vm_name: str | None = None, vm_identifier: str | None = None, ) -> list[dict[str, Any]]: params: dict[str, Any] = {} if vm_identifier: params["vmIdentifier"] = vm_identifier if vm_name: params["vmName"] = vm_name data = await self.json("GET", "/v1/vms", params=params or None) if data is None: return [] if isinstance(data, list): return data if isinstance(data, dict): inner = pick(data, "value", "items", "vms") if isinstance(inner, list): return inner raise ZertoError(f"GET /v1/vms returned unexpected shape: {type(data).__name__}") async def get_vpg(self, vpg_identifier: str) -> dict[str, Any]: data = await self.json("GET", f"/v1/vpgs/{vpg_identifier}") if not isinstance(data, dict): raise ZertoError(f"GET /v1/vpgs/{vpg_identifier} did not return an object") return data async def list_checkpoints(self, vpg_identifier: str) -> list[dict[str, Any]]: data = await self.json("GET", f"/v1/vpgs/{vpg_identifier}/checkpoints") if data is None: return [] if isinstance(data, list): return data raise ZertoError("GET checkpoints returned unexpected shape") async def insert_checkpoint(self, vpg_identifier: str, tag: str) -> Any: """POST tagged checkpoint. Body key is CheckpointName on documented 9.x API.""" path = f"/v1/vpgs/{vpg_identifier}/checkpoints" try: return await self.json("POST", path, json_body={"CheckpointName": tag}) except ZertoError as exc: if exc.status_code != 400: raise return await self.json("POST", path, json_body={"checkpointName": tag}) async def get_task(self, task_id: str) -> dict[str, Any]: data = await self.json("GET", f"/v1/tasks/{task_id}") if not isinstance(data, dict): raise ZertoError("GET /v1/tasks/{id} did not return an object") return data async def start_flr( self, vpg_identifier: str, vm_identifier: str, checkpoint_identifier: str, initial_download_path: str, ) -> Any: return await self.json( "POST", "/v1/flrs", json_body={ "jflr": { "vpgIdentifier": vpg_identifier, "vmIdentifier": vm_identifier, "CheckpointIdentifier": checkpoint_identifier, "initialDownloadPath": initial_download_path, } }, ) async def get_flr(self, session_id: str) -> Any: return await self.json("GET", f"/v1/flrs/{session_id}") async def browse_flr(self, session_id: str, path: str = "", recursive: bool = False) -> Any: return await self.json( "POST", f"/v1/flrs/{session_id}/browse", json_body={"path": path, "recursive": recursive}, ) async def download_flr(self, session_id: str, path_list: list[str]) -> Any: return await self.json( "POST", f"/v1/flrs/{session_id}/download", json_body={"pathList": path_list}, ) async def fetch_download(self, token: str) -> bytes: token = str(token).strip().strip('"') if token.startswith("v1/") or token.startswith("/v1/"): paths = [token if token.startswith("/") else f"/{token}"] else: paths = [f"/v1/downloads/{token}", f"/v1/flrs/{token}"] last = None for path in paths: last = await self.request("GET", path) if last.status_code < 400: return last.content raise ZertoError( f"FLR download HTTP {last.status_code if last else '?'}: {(last.text[:300] if last else '')}", status_code=last.status_code if last else None, ) async def get_localsite(self) -> dict[str, Any]: data = await self.json("GET", "/v1/localsite") return data if isinstance(data, dict) else {} async def get_peersites(self) -> list[dict[str, Any]]: data = await self.json("GET", "/v1/peersites") return [r for r in data if isinstance(r, dict)] if isinstance(data, list) else [] async def list_flrs(self) -> Any: """GET /v1/flrs. Every FLR session the ZVM currently knows about.""" return await self.json("GET", "/v1/flrs") async def end_flr(self, session_id: str) -> None: try: await self.json("DELETE", f"/v1/flrs/{session_id}") except ZertoError: await self.json("POST", f"/v1/flrs/{session_id}") async def start_failover_test( self, vpg_identifier: str, checkpoint_identifier: str | None = None, vm_identifiers: list[str] | None = None, ) -> Any: body: dict[str, Any] = {} if checkpoint_identifier: body["checkpointIdentifier"] = checkpoint_identifier if vm_identifiers: body["vmIdentifiers"] = vm_identifiers return await self.json( "POST", f"/v1/vpgs/{vpg_identifier}/FailoverTest", json_body=body or None, ) async def stop_failover_test(self, vpg_identifier: str, success: bool, summary: str) -> Any: return await self.json( "POST", f"/v1/vpgs/{vpg_identifier}/FailoverTestStop", json_body={"failoverTestSuccess": success, "failoverTestSummary": summary}, ) async def start_clone( self, vpg_identifier: str, checkpoint_identifier: str, datastore_identifier: str | None = None, vm_identifiers: list[str] | None = None, ) -> Any: body: dict[str, Any] = {"checkpointIdentifier": checkpoint_identifier} if datastore_identifier: body["datastoreIdentifier"] = datastore_identifier if vm_identifiers: body["vmIdentifiers"] = vm_identifiers return await self.json( "POST", f"/v1/vpgs/{vpg_identifier}/CloneStart", json_body=body, )