openvpn-endpoint-server/client/services/api_client.py

"""REST API client for server communication."""

import httpx
from typing import Optional
from dataclasses import dataclass


@dataclass
class Gateway:
    """Gateway data class."""
    id: int
    name: str
    description: Optional[str]
    location: Optional[str]
    router_type: str
    is_online: bool
    vpn_ip: Optional[str]


@dataclass
class Endpoint:
    """Endpoint data class."""
    id: int
    gateway_id: int
    name: str
    description: Optional[str]
    internal_ip: str
    port: int
    protocol: str
    application_name: Optional[str]


class APIClient:
    """REST API client for mGuard VPN Server."""

    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip('/')
        self.access_token: Optional[str] = None
        self.refresh_token: Optional[str] = None
        self.client = httpx.Client(timeout=30.0, follow_redirects=True)

    def _headers(self) -> dict:
        """Get request headers with auth token."""
        headers = {"Content-Type": "application/json"}
        if self.access_token:
            headers["Authorization"] = f"Bearer {self.access_token}"
        return headers

    def _request(self, method: str, endpoint: str, **kwargs) -> dict:
        """Make authenticated request."""
        url = f"{self.base_url}{endpoint}"
        response = self.client.request(method, url, headers=self._headers(), **kwargs)
        response.raise_for_status()
        return response.json() if response.text else {}

    def login(self, username: str, password: str) -> bool:
        """Login and store tokens."""
        try:
            response = self.client.post(
                f"{self.base_url}/api/auth/login",
                json={"username": username, "password": password}
            )
            response.raise_for_status()
            data = response.json()
            self.access_token = data["access_token"]
            self.refresh_token = data["refresh_token"]
            print(f"Login successful, token: {self.access_token[:50]}...")
            return True
        except httpx.HTTPError as e:
            print(f"Login error: {e}")
            return False

    def logout(self):
        """Clear stored tokens."""
        self.access_token = None
        self.refresh_token = None

    def refresh_access_token(self) -> bool:
        """Refresh access token using refresh token."""
        if not self.refresh_token:
            return False
        try:
            response = self.client.post(
                f"{self.base_url}/api/auth/refresh",
                params={"refresh_token": self.refresh_token}
            )
            response.raise_for_status()
            data = response.json()
            self.access_token = data["access_token"]
            self.refresh_token = data["refresh_token"]
            return True
        except httpx.HTTPError:
            return False

    def get_current_user(self) -> Optional[dict]:
        """Get current user information."""
        try:
            return self._request("GET", "/api/auth/me")
        except httpx.HTTPError:
            return None

    def get_gateways(self) -> list[Gateway]:
        """Get list of accessible gateways."""
        try:
            print(f"Fetching gateways with token: {self.access_token[:30] if self.access_token else 'NONE'}...")
            data = self._request("GET", "/api/gateways/")
            return [
                Gateway(
                    id=g["id"],
                    name=g["name"],
                    description=g.get("description"),
                    location=g.get("location"),
                    router_type=g["router_type"],
                    is_online=g["is_online"],
                    vpn_ip=g.get("vpn_ip")
                )
                for g in data
            ]
        except httpx.HTTPError as e:
            print(f"Error fetching gateways: {e}")
            return []

    def get_gateways_status(self) -> list[dict]:
        """Get online status of all gateways."""
        try:
            return self._request("GET", "/api/gateways/status")
        except httpx.HTTPError:
            return []

    def get_endpoints(self, gateway_id: int) -> list[Endpoint]:
        """Get endpoints for a gateway."""
        try:
            data = self._request("GET", f"/api/endpoints/gateway/{gateway_id}")
            return [
                Endpoint(
                    id=e["id"],
                    gateway_id=e["gateway_id"],
                    name=e["name"],
                    description=e.get("description"),
                    internal_ip=e["internal_ip"],
                    port=e["port"],
                    protocol=e["protocol"],
                    application_name=e.get("application_name")
                )
                for e in data
            ]
        except httpx.HTTPError:
            return []

    def connect(self, gateway_id: int, endpoint_id: int) -> dict:
        """Request connection to endpoint."""
        try:
            return self._request(
                "POST", "/api/connections/connect",
                json={"gateway_id": gateway_id, "endpoint_id": endpoint_id}
            )
        except httpx.HTTPError as e:
            return {"success": False, "message": str(e)}

    def disconnect(self, connection_id: int) -> dict:
        """Disconnect from endpoint."""
        try:
            return self._request(
                "POST", "/api/connections/disconnect",
                json={"connection_id": connection_id}
            )
        except httpx.HTTPError as e:
            return {"message": str(e)}

    def get_active_connections(self) -> list[dict]:
        """Get list of active connections."""
        try:
            return self._request("GET", "/api/connections/active")
        except httpx.HTTPError:
            return []

    def close(self):
        """Close HTTP client."""
        self.client.close()