"""Connection and session management for VergeOS API."""
import warnings
from base64 import b64encode
from collections.abc import Iterable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Optional, Union
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from pyvergeos.constants import (
API_VERSION,
RETRY_BACKOFF_FACTOR,
RETRY_METHODS,
RETRY_STATUS_CODES,
RETRY_TOTAL,
)
HeaderValue = Union[str, bytes]
HeaderSnapshot = tuple[str, Optional[HeaderValue]]
[docs]
class AuthMethod(Enum):
"""Authentication methods supported by VergeOS API."""
BASIC = "basic"
TOKEN = "token"
[docs]
@dataclass
class VergeConnection:
"""Manages connection state to a VergeOS system.
Attributes:
host: VergeOS hostname or IP address.
username: Username for authentication.
api_base_url: Computed API base URL.
token: Authentication token (Basic or Bearer).
token_expires: Token expiration time (if applicable).
verify_ssl: Whether to verify SSL certificates. When False, the
InsecureRequestWarning suppression is process-global — it will
silence warnings for all clients in the same process, even those
with verify_ssl=True.
retry_total: Number of retry attempts for transient failures.
retry_backoff_factor: Backoff factor for retry delay calculation.
retry_status_codes: HTTP status codes that trigger automatic retry.
session: Optional caller-supplied requests session. When supplied, the
SDK does not mount adapters, change TLS verification, or close it
on disconnect by default. Header snapshots for VergeClient-managed
authentication are populated by VergeClient.connect(). Do not share
one supplied session across multiple active VergeClient instances;
authentication is stored in session headers while connected.
close_session: Override for disconnect lifecycle. None closes
SDK-created sessions and leaves caller-supplied sessions open.
close_session=False requires a caller-supplied session.
connected_at: Timestamp when connection was established.
vergeos_version: VergeOS version from system endpoint.
is_connected: Whether connection is active.
"""
host: str
username: str = ""
api_base_url: str = field(init=False)
token: Optional[str] = None
token_expires: Optional[datetime] = None
verify_ssl: bool = True
retry_total: int = RETRY_TOTAL
retry_backoff_factor: float = RETRY_BACKOFF_FACTOR
retry_status_codes: Iterable[int] = RETRY_STATUS_CODES
connected_at: Optional[datetime] = None
vergeos_version: Optional[str] = None
os_version: Optional[str] = None
cloud_name: Optional[str] = None
is_connected: bool = False
session: Optional[requests.Session] = field(default=None, repr=False)
close_session: Optional[bool] = None
_owns_session: bool = field(init=False, default=True, repr=False)
_pre_connect_headers: Optional[dict[str, HeaderSnapshot]] = field(
init=False, default=None, repr=False
)
def __post_init__(self) -> None:
self.api_base_url = f"https://{self.host}/api/{API_VERSION}"
self._owns_session = self.session is None
self.retry_status_codes = frozenset(self.retry_status_codes)
if self._owns_session and self.close_session is False:
raise ValueError("close_session=False requires a caller-supplied session")
# Create session (done here so it can be mocked in tests)
if self.session is None:
self.session = requests.Session()
if not self._owns_session:
ignored_config = []
if self.retry_total != RETRY_TOTAL:
ignored_config.append("retry_total")
if self.retry_backoff_factor != RETRY_BACKOFF_FACTOR:
ignored_config.append("retry_backoff_factor")
if set(self.retry_status_codes) != set(RETRY_STATUS_CODES):
ignored_config.append("retry_status_codes")
if not self.verify_ssl:
ignored_config.append("verify_ssl")
if ignored_config:
warnings.warn(
"VergeConnection received a caller-supplied session; "
f"{', '.join(ignored_config)} will be ignored. Configure retry "
"and TLS behavior on the supplied session instead.",
UserWarning,
stacklevel=2,
)
return
# Configure retry strategy with configurable parameters
retry_strategy = Retry(
total=self.retry_total,
backoff_factor=self.retry_backoff_factor,
status_forcelist=list(self.retry_status_codes),
allowed_methods=list(RETRY_METHODS),
)
adapter = HTTPAdapter(
max_retries=retry_strategy,
pool_connections=1,
pool_maxsize=10,
)
self.session.mount("https://", adapter)
if not self.verify_ssl:
self.session.verify = False
# Suppress InsecureRequestWarning. Note: this is process-global —
# setting verify_ssl=False on any client silences warnings for all
# clients in the same process.
import urllib3
warnings.filterwarnings(
"ignore",
message="Unverified HTTPS request",
category=urllib3.exceptions.InsecureRequestWarning,
)
def _matching_header_key(self, key: str) -> str:
"""Return the current header spelling matching key, if present."""
if self.session is None:
raise RuntimeError("Session not initialized")
key_lower = key.lower()
for existing_key in self.session.headers:
if existing_key.lower() == key_lower:
return existing_key
return key
[docs]
def is_token_valid(self) -> bool:
"""Check if the current token/credentials are valid.
Returns:
True if connected and token not expired.
"""
if not self.is_connected:
return False
return not (self.token_expires and datetime.now(timezone.utc) >= self.token_expires)
[docs]
def disconnect(self) -> None:
"""Clear connection state and close the HTTP session.
Resets all authentication and connection state including token,
expiration time, and connection timestamp. SDK-owned requests sessions
are closed to release network resources. Caller-supplied session
headers are restored and the session is left open by default.
Note:
After calling disconnect(), the connection object can be reused
by establishing a new connection through the client.
"""
self.token = None
self.token_expires = None
self.connected_at = None
self.is_connected = False
if self.session and not self._owns_session and self._pre_connect_headers is not None:
for key, (prior_key, prior_value) in self._pre_connect_headers.items():
self.session.headers.pop(key, None)
if prior_value is not None:
self.session.headers[prior_key] = prior_value
self._pre_connect_headers = None
should_close = self.close_session if self.close_session is not None else self._owns_session
if self.session and should_close:
self.session.close()