UGOS 1.0.0__py3-none-any.whl

ugos/__about__.py

@@ -0,0 +1,6 @@
+"""Package metadata."""
+
+__version__ = "1.0.0"
+
+SUPPORTED_MAJOR_VERSION = 7
+SUPPORTED_NGFW = f"{SUPPORTED_MAJOR_VERSION}.x"

ugos/__init__.py

@@ -0,0 +1,92 @@
+"""Python client for the UGOS NGFW web-console API."""
+
+from .__about__ import SUPPORTED_MAJOR_VERSION, SUPPORTED_NGFW, __version__
+from .catalog import Catalog, load_catalog
+from .client import UGOSClient
+from .exceptions import (
+    AccessDeniedError,
+    APIError,
+    AuthenticationError,
+    CatalogError,
+    CleanupError,
+    ConfigurationError,
+    ConflictError,
+    CSRFError,
+    DecodeError,
+    DownloadError,
+    DuplicateError,
+    HTTPStatusError,
+    NotFoundError,
+    ProtocolError,
+    ReadOnlyListError,
+    RequestTimeoutError,
+    ServerError,
+    SessionExpiredError,
+    TransportError,
+    UnknownActionError,
+    UnknownControllerError,
+    UGOSError,
+    ValidationError,
+)
+from .models.common import DownloadedFile, ListResult, SessionInfo
+from .models.entities import (
+    Administrator,
+    Certificate,
+    FirewallRule,
+    Group,
+    Interface,
+    LicenseInfo,
+    NetworkList,
+    PcapFilter,
+    PcapRule,
+    Tag,
+    Zone,
+)
+from .resources.base import Resource
+
+__all__ = [
+    "SUPPORTED_MAJOR_VERSION",
+    "SUPPORTED_NGFW",
+    "APIError",
+    "AccessDeniedError",
+    "Administrator",
+    "AuthenticationError",
+    "CSRFError",
+    "Catalog",
+    "CatalogError",
+    "Certificate",
+    "CleanupError",
+    "ConfigurationError",
+    "ConflictError",
+    "DecodeError",
+    "DownloadError",
+    "DownloadedFile",
+    "DuplicateError",
+    "FirewallRule",
+    "Group",
+    "HTTPStatusError",
+    "Interface",
+    "LicenseInfo",
+    "ListResult",
+    "NetworkList",
+    "NotFoundError",
+    "PcapFilter",
+    "PcapRule",
+    "ProtocolError",
+    "ReadOnlyListError",
+    "RequestTimeoutError",
+    "Resource",
+    "ServerError",
+    "SessionExpiredError",
+    "SessionInfo",
+    "Tag",
+    "TransportError",
+    "UnknownActionError",
+    "UnknownControllerError",
+    "UGOSClient",
+    "UGOSError",
+    "ValidationError",
+    "Zone",
+    "__version__",
+    "load_catalog",
+]

ugos/_auth.py

@@ -0,0 +1,293 @@
+"""Console authentication, CSRF handling and session lifecycle."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from collections.abc import Mapping
+from typing import Any, cast
+
+import httpx
+
+from ._codec import is_error_payload, parse_response
+from ._encoding import encode_params
+from ._redaction import redact_mapping
+from ._transport import FORM_HEADERS, HTML_HEADERS, HttpTransport
+from .exceptions import (
+    AuthenticationError,
+    CSRFError,
+    DecodeError,
+    HTTPStatusError,
+    ProtocolError,
+    SessionExpiredError,
+    UGOSError,
+)
+from .models.common import SessionInfo
+
+_CSRF_RE = re.compile(r"CSRF_TOKEN\s*=\s*['\"]([^'\"]+)['\"]")
+_CSRF_INPUT_RE = re.compile(
+    r"name=['\"]csrfmiddlewaretoken['\"][^>]*value=['\"]([^'\"]+)['\"]",
+    re.IGNORECASE,
+)
+_CSRF_META_RE = re.compile(
+    r"name=['\"]csrf-token['\"][^>]*content=['\"]([^'\"]+)['\"]",
+    re.IGNORECASE,
+)
+
+
+class AuthManager:
+    """Owns login/logout, cookies and the CSRF token."""
+
+    def __init__(
+        self,
+        transport: HttpTransport,
+        *,
+        username: str | None = None,
+        password: str | None = None,
+        login_path: str = "utm/login",
+        logout_path: str = "utm/CoreLogout",
+        auto_reauth: bool = True,
+    ) -> None:
+        self._transport = transport
+        self._username = username
+        self._password = password
+        self._login_path = login_path
+        self._logout_path = logout_path
+        self._auto_reauth = auto_reauth
+        self._session: SessionInfo | None = None
+        self._csrf: str | None = None
+        self._generation = 0
+
+    @property
+    def session(self) -> SessionInfo | None:
+        return self._session
+
+    @property
+    def authenticated(self) -> bool:
+        return self._session is not None
+
+    @property
+    def csrf_token(self) -> str | None:
+        return self._csrf
+
+    @property
+    def generation(self) -> int:
+        return self._generation
+
+    def invalidate(self) -> None:
+        self._session = None
+        self._csrf = None
+        self._generation += 1
+        self._transport.clear_cookies(("_console_session", "csrf_token"))
+
+    @property
+    def auto_reauth(self) -> bool:
+        return self._auto_reauth
+
+    def ensure_csrf(self) -> str:
+        if self._csrf:
+            return self._csrf
+        token = self._transport.get_cookie("csrf_token")
+        if token:
+            self._csrf = token
+            return token
+        return self.refresh_csrf()
+
+    @staticmethod
+    def _is_login_redirect(status_code: int, location: str) -> bool:
+        if status_code not in {301, 302, 303, 307, 308}:
+            return False
+        return "utm/login" in location or location.rstrip("/").endswith("/login")
+
+    def _raise_csrf_status(self, response: httpx.Response) -> None:
+        location = response.headers.get("location", "")
+        login_redirect = self._is_login_redirect(response.status_code, location)
+        if response.status_code == 401 or login_redirect:
+            raise SessionExpiredError(
+                "session expired while refreshing CSRF",
+                status_code=response.status_code,
+            )
+        if response.status_code >= 400:
+            raise CSRFError(
+                f"CSRF refresh failed with HTTP {response.status_code}",
+                status_code=response.status_code,
+            )
+
+    def _extract_csrf_token(self, response: httpx.Response) -> str:
+        token = self._transport.get_cookie("csrf_token") or _extract_csrf(response.text)
+        if not token:
+            raise CSRFError("the login page did not contain a CSRF token")
+        return token
+
+    def refresh_csrf(self) -> str:
+        path = "/" if self.authenticated else self._login_path
+        try:
+            response = self._transport.request(
+                "GET",
+                path,
+                headers=HTML_HEADERS,
+                idempotent=True,
+            )
+        except UGOSError:
+            raise
+        except Exception as exc:
+            raise CSRFError(f"failed to obtain a CSRF token: {exc}") from exc
+        self._raise_csrf_status(response)
+        self.note_response(response)
+        self._csrf = self._extract_csrf_token(response)
+        return self._csrf
+
+    def _reset_session(self) -> None:
+        self._transport.clear_cookies(("_console_session", "csrf_token"))
+        self._session = None
+        self._csrf = None
+
+    def _login_page_token(self) -> str:
+        page = self._transport.request(
+            "GET", self._login_path, headers=HTML_HEADERS, idempotent=True
+        )
+        self.note_response(page)
+        return self._extract_csrf_token(page)
+
+    def _post_login(self, token: str) -> httpx.Response:
+        body = encode_params(
+            {
+                "username": self._username,
+                "password": self._password,
+                "csrfmiddlewaretoken": token,
+            }
+        )
+        response = self._transport.request(
+            "POST",
+            self._login_path,
+            data=body,
+            headers=FORM_HEADERS,
+            idempotent=False,
+            follow_redirects=False,
+        )
+        self.note_response(response)
+        return response
+
+    def _parse_login_response(self, response: httpx.Response) -> Mapping[str, Any]:
+        try:
+            payload = parse_response(response, expect="json")
+        except DecodeError as exc:
+            raise AuthenticationError(
+                "login returned a non-JSON response",
+                status_code=response.status_code,
+            ) from exc
+        except (HTTPStatusError, ProtocolError) as exc:
+            raise AuthenticationError(
+                f"login failed: {exc}",
+                status_code=getattr(exc, "status_code", response.status_code),
+                payload=getattr(exc, "payload", None),
+            ) from exc
+        if not isinstance(payload, Mapping) or not payload.get("auth_token"):
+            raise AuthenticationError("login was rejected by the appliance")
+        return payload
+
+    def _build_session(self, payload: Mapping[str, Any], token: str) -> SessionInfo:
+        session_cookie = self._transport.get_cookie("_console_session")
+        if not session_cookie:
+            raise ProtocolError("login succeeded but _console_session cookie is absent")
+        self._csrf = self._transport.get_cookie("csrf_token") or token
+        self._session = SessionInfo(
+            username=cast(str, self._username),
+            admin_id=_as_int(payload.get("admin_id")),
+            product=_as_str(payload.get("product")),
+            version=_as_str(payload.get("version")),
+            auth_token=_as_bool(payload.get("auth_token")),
+            fingerprint=hashlib.sha256(session_cookie.encode()).hexdigest()[:12],
+            raw=redact_mapping(payload),
+        )
+        self._generation += 1
+        return self._session
+
+    def _has_reusable_session(self, force: bool) -> bool:
+        return self._session is not None and not force
+
+    def _credentials_missing(self) -> bool:
+        return not self._username or self._password is None
+
+    def _should_reset_session(self, force: bool) -> bool:
+        return force or self._session is not None
+
+    def login(self, *, force: bool = False) -> SessionInfo:
+        if self._has_reusable_session(force):
+            return self._session  # type: ignore[return-value]
+        if self._credentials_missing():
+            raise AuthenticationError("username and password are required for login")
+        if self._should_reset_session(force):
+            self._reset_session()
+        token = self._login_page_token()
+        response = self._post_login(token)
+        return self._build_session(self._parse_login_response(response), token)
+
+    def _logout_request(self) -> httpx.Response:
+        token = self._csrf or self.ensure_csrf()
+        response = self._transport.request(
+            "POST",
+            self._logout_path,
+            data=encode_params({"csrfmiddlewaretoken": token}),
+            headers=FORM_HEADERS,
+            idempotent=False,
+        )
+        self.note_response(response)
+        return response
+
+    def _logout_result(self, response: httpx.Response) -> bool:
+        try:
+            payload = parse_response(response, expect="auto")
+        except SessionExpiredError:
+            return True
+        return not is_error_payload(payload)
+
+    def logout(self, *, best_effort: bool = True) -> bool:
+        if self._session is None:
+            self._generation += 1
+            self._transport.clear_cookies(("_console_session", "csrf_token"))
+            return True
+        ok = False
+        try:
+            ok = self._logout_result(self._logout_request())
+        except UGOSError:
+            if not best_effort:
+                raise
+        finally:
+            self._generation += 1
+            self._reset_session()
+        return ok
+
+    def note_response(self, response: httpx.Response) -> None:
+        token = self._transport.get_cookie("csrf_token")
+        if token:
+            self._csrf = token
+        # Session state is invalidated by the guarded re-auth path, not here:
+        # a stale response from another thread must not wipe a fresh session.
+
+
+def _extract_csrf(html: str) -> str | None:
+    for pattern in (_CSRF_RE, _CSRF_INPUT_RE, _CSRF_META_RE):
+        match = pattern.search(html)
+        if match:
+            return match.group(1)
+    return None
+
+
+def _as_int(value: Any) -> int | None:
+    try:
+        return int(value) if value is not None else None
+    except (TypeError, ValueError):
+        return None
+
+
+def _as_str(value: Any) -> str | None:
+    return str(value) if value is not None else None
+
+
+def _as_bool(value: Any) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return bool(value)

ugos/_client_base.py

@@ -0,0 +1,52 @@
+"""Shared typing surface for the client mixins."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, TypeVar
+
+from ._auth import AuthManager
+from ._client_helpers import _RetryState
+from ._config import ClientConfig
+from ._dynamic import ControllerProxy
+from ._transport import HttpTransport
+from .catalog import Catalog
+
+CONTROLLER_RE = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9_-]*$")
+_T = TypeVar("_T", bound=ControllerProxy)
+
+
+class _ClientBase:
+    _config: ClientConfig
+    _transport: HttpTransport
+    _auth: AuthManager
+    _lock: Any
+    _catalog: Catalog
+    _resources: dict[str, ControllerProxy]
+    _auto_login: bool
+    _logout_on_exit: bool
+    _catalog_check: bool
+
+    def _prepare(self) -> None:  # pragma: no cover - overridden by the client
+        raise NotImplementedError
+
+    def _refresh_session(self, expected_generation: int) -> None:  # pragma: no cover
+        raise NotImplementedError
+
+    def _validate(
+        self, controller: str, action: str, *, strict: bool | None = None
+    ) -> None:  # pragma: no cover
+        raise NotImplementedError
+
+    def _retry_csrf(self, state: _RetryState) -> bool:  # pragma: no cover
+        raise NotImplementedError
+
+    def _retry_session(  # pragma: no cover
+        self,
+        state: _RetryState,
+        generation: int,
+        action: str,
+        idempotent: bool,
+        exc: Exception,
+    ) -> bool:
+        raise NotImplementedError

ugos/_client_calls.py

@@ -0,0 +1,200 @@
+"""Generic action call implementation."""
+
+from __future__ import annotations
+
+from collections.abc import Collection, Mapping
+from typing import Any
+
+import httpx
+
+from ._client_base import _ClientBase
+from ._client_helpers import _merge_params, _retry_allowed, _RetryState
+from ._codec import parse_response
+from ._encoding import encode_pairs, encode_params
+from ._transport import FORM_HEADERS
+from ._transport import UNSET as _UNSET
+from .exceptions import AccessDeniedError, CSRFError, SessionExpiredError
+
+
+class _CallsMixin(_ClientBase):
+    def _resolve_method(self, contract: Any, method: str | None) -> str:
+        return (method or (contract.method if contract else "POST")).upper()
+
+    def _resolve_expect(self, contract: Any, expect: str | None) -> str:
+        return expect or (contract.response_kind if contract else "auto")
+
+    def _resolve_json_fields(
+        self,
+        contract: Any,
+        json_fields: Collection[str],
+    ) -> tuple[str, ...]:
+        if contract and contract.json_fields:
+            json_fields = (*json_fields, *contract.json_fields)
+        return tuple(dict.fromkeys(json_fields))
+
+    def _resolve_idempotent(
+        self,
+        contract: Any,
+        method: str,
+        idempotent: bool | None,
+    ) -> bool:
+        if idempotent is not None:
+            return idempotent
+        if contract is not None and contract.idempotent is not None:
+            return bool(contract.idempotent)
+        return method in {"GET", "HEAD", "OPTIONS"}
+
+    def _resolve_call_options(
+        self,
+        controller: str,
+        action: str,
+        method: str | None,
+        expect: str | None,
+        idempotent: bool | None,
+        json_fields: Collection[str],
+        contract: Any,
+    ) -> tuple[str, str, tuple[str, ...], bool]:
+        resolved_method = self._resolve_method(contract, method)
+        return (
+            resolved_method,
+            self._resolve_expect(contract, expect),
+            self._resolve_json_fields(contract, json_fields),
+            self._resolve_idempotent(contract, resolved_method, idempotent),
+        )
+
+    def _send_call(
+        self,
+        controller: str,
+        action: str,
+        method: str,
+        params: Mapping[str, Any],
+        json_fields: Collection[str],
+        timeout: Any,
+        idempotent: bool,
+    ) -> tuple[httpx.Response, int]:
+        request_params = dict(params)
+        with self._lock:
+            generation = self._auth.generation
+        if method == "GET":
+            query = encode_pairs(request_params, json_fields=json_fields)
+            body = None
+        else:
+            with self._lock:
+                request_params["csrfmiddlewaretoken"] = self._auth.ensure_csrf()
+            body = encode_params(request_params, json_fields=json_fields)
+            query = None
+        response = self._transport.request(
+            method,
+            f"{controller}/{action}",
+            params=query,
+            data=body,
+            headers=FORM_HEADERS if body is not None else None,
+            idempotent=idempotent,
+            timeout=timeout,
+        )
+        with self._lock:
+            self._auth.note_response(response)
+        return response, generation
+
+    def _retry_csrf(self, state: _RetryState) -> bool:
+        if state.csrf >= 1:
+            return False
+        state.csrf += 1
+        with self._lock:
+            self._auth.refresh_csrf()
+        return True
+
+    def _retry_session(
+        self,
+        state: _RetryState,
+        generation: int,
+        action: str,
+        idempotent: bool,
+        exc: Exception,
+    ) -> bool:
+        if state.reauth >= 1:
+            return False
+        self._refresh_session(generation)
+        if not self._auth.auto_reauth or not _retry_allowed(action, idempotent, exc):
+            return False
+        state.reauth += 1
+        return True
+
+    def _retry_access_denied(
+        self,
+        state: _RetryState,
+        generation: int,
+        action: str,
+        idempotent: bool,
+        exc: Exception,
+    ) -> bool:
+        if state.reauth >= 1 or self._transport.get_cookie("_console_session") is not None:
+            return False
+        return self._retry_session(state, generation, action, idempotent, exc)
+
+    def _handle_call_retry(
+        self,
+        exc: Exception,
+        state: _RetryState,
+        generation: int,
+        action: str,
+        idempotent: bool,
+    ) -> bool:
+        if isinstance(exc, CSRFError):
+            return self._retry_csrf(state)
+        if isinstance(exc, SessionExpiredError):
+            return self._retry_session(state, generation, action, idempotent, exc)
+        return self._retry_access_denied(state, generation, action, idempotent, exc)
+
+    def call(
+        self,
+        controller: str,
+        action: str,
+        /,
+        *,
+        payload: Mapping[str, Any] | None = None,
+        json_fields: Collection[str] = (),
+        method: str | None = None,
+        expect: str | None = None,
+        idempotent: bool | None = None,
+        timeout: float | httpx.Timeout | object | None = _UNSET,
+        strict: bool | None = None,
+        **params: Any,
+    ) -> Any:
+        """Call one controller action and return the decoded response."""
+
+        self._validate(controller, action, strict=strict)
+        contract = self._catalog.action_contract(controller, action)
+        merged = _merge_params(payload, params)
+        method, expect, json_fields, idempotent = self._resolve_call_options(
+            controller,
+            action,
+            method,
+            expect,
+            idempotent,
+            json_fields,
+            contract,
+        )
+        with self._lock:
+            self._prepare()
+        state = _RetryState()
+        while True:
+            response, generation = self._send_call(
+                controller,
+                action,
+                method,
+                merged,
+                json_fields,
+                timeout,
+                idempotent,
+            )
+            try:
+                return parse_response(
+                    response,
+                    expect=expect,
+                    controller=controller,
+                    action=action,
+                )
+            except (CSRFError, SessionExpiredError, AccessDeniedError) as exc:
+                if not self._handle_call_retry(exc, state, generation, action, idempotent):
+                    raise