pjdev-armis-sdk 5.1.4__py3-none-any.whl

pjdev_armis_sdk/__about__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: 2026-present Chris O'Neill <chris@purplejay.io>
+#
+# SPDX-License-Identifier: MIT
+__version__ = "5.1.4"

pjdev_armis_sdk/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: 2026-present Chris O'Neill <chris@purplejay.io>
+#
+# SPDX-License-Identifier: MIT
+
+from pjdev_armis_sdk import (
+    access_token,
+    api_utilities,
+    boundaries,
+    config_service,
+    devices,
+    http_client,
+    integrations,
+    models,
+    search,
+    sites,
+    users,
+)
+
+__all__ = [
+    "access_token",
+    "api_utilities",
+    "boundaries",
+    "config_service",
+    "devices",
+    "http_client",
+    "integrations",
+    "models",
+    "search",
+    "sites",
+    "users",
+]

pjdev_armis_sdk/_internal.py

@@ -0,0 +1,32 @@
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Awaitable, Callable, Optional, TypeVar
+
+import httpx
+
+from pjdev_armis_sdk.http_client import http_client
+
+T = TypeVar("T")
+
+
+@asynccontextmanager
+async def _resolve_client(
+    client: Optional[httpx.AsyncClient],
+) -> AsyncIterator[httpx.AsyncClient]:
+    if client is not None:
+        yield client
+        return
+    async with http_client() as _client:
+        yield _client
+
+
+async def with_client(
+    client: Optional[httpx.AsyncClient],
+    fn: Callable[[httpx.AsyncClient], Awaitable[T]],
+) -> T:
+    """Run `fn` against either the caller-provided client or a fresh one."""
+    async with _resolve_client(client) as c:
+        return await fn(c)
+
+
+def drop_none(params: dict[str, Any]) -> dict[str, Any]:
+    return {k: v for k, v in params.items() if v is not None}

pjdev_armis_sdk/access_token.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import httpx
+
+from pjdev_armis_sdk.api_utilities import async_retry_http
+from pjdev_armis_sdk.config_service import get_config
+from pjdev_armis_sdk.models import ArmisAccessTokenResponse, ArmisTokenData
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403])
+async def get_access_token(
+    secret_key: Optional[str] = None,
+) -> ArmisTokenData:
+    """
+    Exchange a secret_key for a temporary Armis access token.
+
+    This is normally handled automatically by `ArmisAccessTokenAuth` on the shared
+    `httpx.AsyncClient`. Call this directly only when you need the raw token (e.g.
+    for inspection, or to pass to a non-SDK consumer).
+    """
+    config = get_config()
+    key = secret_key or config.secret_key
+    if not key:
+        raise ValueError("secret_key not provided and not configured")
+    if not config.instance_url:
+        raise ValueError("instance_url not configured")
+
+    async with httpx.AsyncClient(
+        base_url=config.instance_url,
+        timeout=config.request_timeout_seconds,
+    ) as client:
+        r = await client.post(
+            "/api/v1/access_token/",
+            data={"secret_key": key},
+        )
+        r.raise_for_status()
+        return ArmisAccessTokenResponse.model_validate(r.json()).data

pjdev_armis_sdk/api_utilities.py

@@ -0,0 +1,110 @@
+import asyncio
+import time
+from functools import wraps
+from typing import Any, Awaitable, Callable, List, Optional, ParamSpec, TypeVar
+
+import httpx
+from httpx import ConnectError, HTTPStatusError
+from loguru import logger
+
+from pjdev_armis_sdk.config_service import get_config
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+async def log_request_headers(request: httpx.Request) -> None:
+    logger.debug(f"Request: {request.method} {request.url}")
+
+
+async def log_response_headers(response: httpx.Response) -> None:
+    logger.debug(
+        f"Response: {response.request.method} {response.request.url} -> {response.status_code}"
+    )
+
+
+def async_retry_http(
+    default_value: Optional[Any] = None,
+    status_codes_to_ignore: Optional[List[int]] = None,
+) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+    def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
+        @wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            config = get_config()
+            max_attempts = config.http_retry_max_count
+            delay_seconds = config.http_retry_delay_seconds
+            attempts = 0
+            exceptions: List[Exception] = []
+
+            while attempts < max_attempts:
+                try:
+                    return await func(*args, **kwargs)
+                except HTTPStatusError as e:
+                    logger.warning(
+                        f"{e.response.status_code}: {e.response.reason_phrase} - {e.response.text}"
+                    )
+                    exceptions.append(e)
+                    if (
+                        status_codes_to_ignore
+                        and e.response.status_code in status_codes_to_ignore
+                    ):
+                        break
+                except ConnectError as e:
+                    logger.warning(f"{e.request.url} not reachable")
+                    exceptions.append(e)
+
+                attempts += 1
+                if attempts == max_attempts:
+                    break
+                total_delay = delay_seconds**attempts
+                logger.warning(
+                    f"Attempt {attempts}/{max_attempts} failed. Retrying in {total_delay} seconds..."
+                )
+                await asyncio.sleep(total_delay)
+
+            if default_value is None:
+                raise ExceptionGroup(
+                    f"Failed after {max_attempts} attempts", exceptions
+                )
+            logger.error(f"Failed after {max_attempts} attempts")
+            return default_value
+
+        return wrapper
+
+    return decorator
+
+
+def record_time(
+    log: bool = False,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        if asyncio.iscoroutinefunction(func):
+
+            @wraps(func)
+            async def wrap_func_async(*args: P.args, **kwargs: P.kwargs) -> Any:
+                t1 = time.time()
+                result = await func(*args, **kwargs)
+                t2 = time.time()
+                if log:
+                    logger.info(
+                        f"Function {func.__name__!r} executed in {(t2 - t1):.2f}s"
+                    )
+                return result
+
+            return wrap_func_async  # type: ignore[return-value]
+
+        @wraps(func)
+        def wrap_func(*args: P.args, **kwargs: P.kwargs) -> R:
+            t1 = time.time()
+            result = func(*args, **kwargs)
+            t2 = time.time()
+            if log:
+                logger.info(
+                    f"Function {func.__name__!r} executed in {(t2 - t1):.2f}s"
+                )
+
+            return result
+
+        return wrap_func
+
+    return decorator

pjdev_armis_sdk/boundaries.py

@@ -0,0 +1,104 @@
+from typing import Any, Dict, Optional, Sequence
+
+import httpx
+
+from pjdev_armis_sdk._internal import drop_none, with_client
+from pjdev_armis_sdk.api_utilities import async_retry_http
+from pjdev_armis_sdk.models import Boundary
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403])
+async def list_boundaries(
+    paging_from: int = 0,
+    paging_length: Optional[int] = None,
+    affected_sites: Optional[Sequence[str]] = None,
+    boundary_ids: Optional[Sequence[int]] = None,
+    include_total: bool = False,
+    fields: Optional[Sequence[str]] = None,
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """GET /api/v1/boundaries/"""
+    params: Dict[str, Any] = drop_none(
+        {
+            "from": paging_from,
+            "length": paging_length,
+            "affectedSites": ",".join(affected_sites) if affected_sites else None,
+            "boundaryIds": ",".join(str(b) for b in boundary_ids)
+            if boundary_ids
+            else None,
+            "includeTotal": include_total,
+            "fields": ",".join(fields) if fields else None,
+        }
+    )
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.get("/api/v1/boundaries/", params=params)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def get_boundary(
+    boundary_id: int,
+    fields: Optional[Sequence[str]] = None,
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """GET /api/v1/boundaries/{boundary_id}/"""
+    params: Dict[str, Any] = drop_none(
+        {"fields": ",".join(fields) if fields else None}
+    )
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.get(f"/api/v1/boundaries/{boundary_id}/", params=params)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403])
+async def create_boundary(
+    boundary: Boundary,
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """POST /api/v1/boundaries/"""
+    payload = boundary.model_dump(by_alias=True, exclude_none=True, exclude={"id"})
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.post("/api/v1/boundaries/", json=payload)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def update_boundary(
+    boundary_id: int,
+    payload: Dict[str, Any],
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """PATCH /api/v1/boundaries/{boundary_id}/"""
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.patch(f"/api/v1/boundaries/{boundary_id}/", json=payload)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def delete_boundary(
+    boundary_id: int,
+    client: Optional[httpx.AsyncClient] = None,
+) -> None:
+    """DELETE /api/v1/boundaries/{boundary_id}/"""
+
+    async def _exec(c: httpx.AsyncClient) -> None:
+        r = await c.delete(f"/api/v1/boundaries/{boundary_id}/")
+        r.raise_for_status()
+
+    await with_client(client, _exec)

pjdev_armis_sdk/config_service.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from pjdev_armis_sdk.models import Config
+
+__ctx: Dict[str, Config] = {}
+
+
+def get_config() -> Config:
+    if "config" not in __ctx:
+        __ctx["config"] = Config()
+    return __ctx["config"]
+
+
+def init(
+    env_path: Optional[Path] = None,
+    instance_url: Optional[str] = None,
+    secret_key: Optional[str] = None,
+    http_retry_max_count: Optional[int] = None,
+    http_retry_delay_seconds: Optional[int] = None,
+    request_timeout_seconds: Optional[float] = None,
+) -> None:
+    if env_path is not None:
+        Config.model_config["env_file"] = env_path
+
+    kwargs: Dict[str, Any] = {}
+    if instance_url is not None:
+        kwargs["instance_url"] = instance_url
+    if secret_key is not None:
+        kwargs["secret_key"] = secret_key
+    if http_retry_max_count is not None:
+        kwargs["http_retry_max_count"] = http_retry_max_count
+    if http_retry_delay_seconds is not None:
+        kwargs["http_retry_delay_seconds"] = http_retry_delay_seconds
+    if request_timeout_seconds is not None:
+        kwargs["request_timeout_seconds"] = request_timeout_seconds
+
+    __ctx["config"] = Config(**kwargs)

pjdev_armis_sdk/devices.py

@@ -0,0 +1,156 @@
+from typing import Any, Dict, List, Optional, Sequence
+
+import httpx
+
+from pjdev_armis_sdk._internal import drop_none, with_client
+from pjdev_armis_sdk.api_utilities import async_retry_http
+from pjdev_armis_sdk.models import Device, DeviceUpsert
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def get_device(
+    id: Optional[int] = None,
+    ip: Optional[str] = None,
+    mac: Optional[str] = None,
+    search: Optional[str] = None,
+    tag: Optional[str] = None,
+    fields: Optional[Sequence[str]] = None,
+    include_total: bool = True,
+    include_network_interfaces: bool = False,
+    paging_from: int = 0,
+    paging_length: Optional[int] = None,
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """
+    GET /api/v1/devices/ — at least one of `id`, `ip`, `mac`, `search`, or `tag` is required.
+    Returns the raw envelope as a dict (`{success, data: {count, total, results, ...}}`).
+    """
+    if not any([id, ip, mac, search, tag]):
+        raise ValueError("Must provide one of: id, ip, mac, search, tag")
+
+    params: Dict[str, Any] = drop_none(
+        {
+            "id": id,
+            "ip": ip,
+            "mac": mac,
+            "search": search,
+            "tag": tag,
+            "fields": ",".join(fields) if fields else None,
+            "includeTotal": include_total,
+            "includeNetworkInterfaces": include_network_interfaces,
+            "from": paging_from,
+            "length": paging_length,
+        }
+    )
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.get("/api/v1/devices/", params=params)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def patch_device(
+    device_id: int,
+    payload: Dict[str, Any],
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """PATCH /api/v1/devices/{device_id}/"""
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.patch(f"/api/v1/devices/{device_id}/", json=payload)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def delete_device(
+    device_id: int,
+    client: Optional[httpx.AsyncClient] = None,
+) -> None:
+    """DELETE /api/v1/devices/{device_id}/"""
+
+    async def _exec(c: httpx.AsyncClient) -> None:
+        r = await c.delete(f"/api/v1/devices/{device_id}/")
+        r.raise_for_status()
+
+    await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403])
+async def bulk_upsert_devices(
+    operations: List[Dict[str, Any]],
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """
+    POST /api/v1/devices/_bulk/
+
+    Each operation is `{"upsert": Device}` or `{"delete": {"deviceId": int}}`.
+    Returns the multi-status envelope as a dict.
+    """
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.post("/api/v1/devices/_bulk/", json=operations)
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+def upsert_op(device: DeviceUpsert | Dict[str, Any]) -> Dict[str, Any]:
+    """Build a `{"upsert": ...}` entry for `bulk_upsert_devices`."""
+    if isinstance(device, DeviceUpsert):
+        return {"upsert": device.model_dump(by_alias=True, exclude_none=True)}
+    return {"upsert": device}
+
+
+def delete_op(device_id: int) -> Dict[str, Any]:
+    """Build a `{"delete": ...}` entry for `bulk_upsert_devices`."""
+    return {"delete": {"deviceId": device_id}}
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def add_device_tags(
+    device_id: int,
+    tags: Sequence[str],
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """POST /api/v1/devices/{device_id}/tags/"""
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.post(
+            f"/api/v1/devices/{device_id}/tags/", json={"tags": list(tags)}
+        )
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+@async_retry_http(status_codes_to_ignore=[400, 401, 403, 404])
+async def remove_device_tags(
+    device_id: int,
+    tags: Sequence[str],
+    client: Optional[httpx.AsyncClient] = None,
+) -> Dict[str, Any]:
+    """DELETE /api/v1/devices/{device_id}/tags/"""
+
+    async def _exec(c: httpx.AsyncClient) -> Dict[str, Any]:
+        r = await c.request(
+            "DELETE",
+            f"/api/v1/devices/{device_id}/tags/",
+            json={"tags": list(tags)},
+        )
+        r.raise_for_status()
+        return r.json()
+
+    return await with_client(client, _exec)
+
+
+def parse_device(raw: Dict[str, Any]) -> Device:
+    """Helper: turn one raw `results[]` item into a typed `Device`."""
+    return Device.model_validate(raw)

pjdev_armis_sdk/http_client.py

@@ -0,0 +1,112 @@
+from contextlib import asynccontextmanager
+from datetime import datetime, timedelta, timezone
+from typing import AsyncIterator, Generator, Optional
+
+import httpx
+from loguru import logger
+
+from pjdev_armis_sdk.api_utilities import log_request_headers, log_response_headers
+from pjdev_armis_sdk.config_service import get_config
+from pjdev_armis_sdk.models import ArmisAccessTokenResponse, ArmisTokenData
+
+
+class ArmisAuthError(Exception):
+    """Raised when token acquisition fails."""
+
+
+class ArmisAccessTokenAuth(httpx.Auth):
+    """
+    httpx auth flow that maintains a cached Armis access token.
+
+    - Fetches a new token from `/api/v1/access_token/` when no token is cached,
+      or when the cached token is within `refresh_leeway` of expiry.
+    - On a 401 response, clears the cache, fetches a new token, and retries once.
+    """
+
+    requires_response_body = True
+
+    def __init__(
+        self,
+        secret_key: str,
+        refresh_leeway: timedelta = timedelta(minutes=1),
+    ) -> None:
+        self._secret_key = secret_key
+        self._refresh_leeway = refresh_leeway
+        self._token: Optional[ArmisTokenData] = None
+
+    def _token_is_fresh(self) -> bool:
+        if self._token is None:
+            return False
+        expiry = self._token.expiration_utc
+        if expiry.tzinfo is None:
+            expiry = expiry.replace(tzinfo=timezone.utc)
+        return expiry - self._refresh_leeway > datetime.now(timezone.utc)
+
+    def _build_token_request(self, base_url: httpx.URL) -> httpx.Request:
+        return httpx.Request(
+            "POST",
+            base_url.join("/api/v1/access_token/"),
+            data={"secret_key": self._secret_key},
+        )
+
+    def _apply_token(self, request: httpx.Request) -> None:
+        if self._token is None:
+            raise ArmisAuthError("No access token available")
+        request.headers["Authorization"] = self._token.access_token
+
+    def _parse_token_response(self, response: httpx.Response) -> None:
+        if response.status_code >= 400:
+            raise ArmisAuthError(
+                f"Failed to acquire Armis access token: "
+                f"{response.status_code} {response.reason_phrase} - {response.text}"
+            )
+        parsed = ArmisAccessTokenResponse.model_validate(response.json())
+        self._token = parsed.data
+        logger.debug(
+            f"Acquired Armis access token (expires {parsed.data.expiration_utc.isoformat()})"
+        )
+
+    def auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response, None]:
+        if not self._token_is_fresh():
+            token_response = yield self._build_token_request(request.url)
+            self._parse_token_response(token_response)
+
+        self._apply_token(request)
+        response = yield request
+
+        if response.status_code == 401:
+            self._token = None
+            token_response = yield self._build_token_request(request.url)
+            self._parse_token_response(token_response)
+            self._apply_token(request)
+            yield request
+
+
+def get_http_client() -> httpx.AsyncClient:
+    config = get_config()
+    if not config.instance_url:
+        raise ValueError(
+            "Armis instance_url not configured; call config_service.init(instance_url=..., secret_key=...) first"
+        )
+    if not config.secret_key:
+        raise ValueError(
+            "Armis secret_key not configured; call config_service.init(instance_url=..., secret_key=...) first"
+        )
+
+    return httpx.AsyncClient(
+        base_url=config.instance_url,
+        auth=ArmisAccessTokenAuth(secret_key=config.secret_key),
+        timeout=config.request_timeout_seconds,
+        event_hooks={
+            "response": [log_response_headers],
+            "request": [log_request_headers],
+        },
+    )
+
+
+@asynccontextmanager
+async def http_client() -> AsyncIterator[httpx.AsyncClient]:
+    async with get_http_client() as client:
+        yield client