kelpgate-client 0.1.0__py3-none-any.whl

kelpgate/__init__.py

@@ -0,0 +1,61 @@
+"""Python client for the Kelpgate malware-scanning API.
+
+    from kelpgate import KelpgateClient
+
+    with KelpgateClient.from_env() as kg:
+        result = kg.scan_file("upload.pdf")
+        if result.infected:
+            print("Rejected:", result.signature)
+
+Use :class:`AsyncKelpgateClient` from asyncio code.
+"""
+
+from ._async import AsyncKelpgateClient
+from ._errors import (
+    AuthError,
+    ConfigError,
+    InfectedError,
+    KelpgateError,
+    NetworkError,
+    PayloadTooLargeError,
+    QuotaExceededError,
+    RateLimitedError,
+    RequestTimeout,
+    ScanError,
+    ServerError,
+    UnsupportedFileError,
+)
+from ._models import AcceptedScan, Health, PendingScan, ScanResult
+from ._sync import KelpgateClient
+from ._version import __version__
+
+
+def assert_clean(result: ScanResult) -> ScanResult:
+    """Return ``result`` if it is clean. Raise :class:`InfectedError` if it is infected."""
+    if result.infected:
+        raise InfectedError(result)
+    return result
+
+
+__all__ = [
+    "AcceptedScan",
+    "AsyncKelpgateClient",
+    "AuthError",
+    "ConfigError",
+    "Health",
+    "InfectedError",
+    "KelpgateClient",
+    "KelpgateError",
+    "NetworkError",
+    "PayloadTooLargeError",
+    "PendingScan",
+    "QuotaExceededError",
+    "RateLimitedError",
+    "RequestTimeout",
+    "ScanError",
+    "ScanResult",
+    "ServerError",
+    "UnsupportedFileError",
+    "__version__",
+    "assert_clean",
+]

kelpgate/_async.py

@@ -0,0 +1,256 @@
+"""The asyncio client. Same behavior as :class:`kelpgate.KelpgateClient`, with ``await``."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from typing import Any, Dict, Mapping, Optional, Union
+from urllib.parse import quote
+
+import httpx
+
+from ._common import (
+    REFRESH_AHEAD,
+    RETRY_STATUSES,
+    Config,
+    PreparedUpload,
+    Source,
+    accepted_outcome,
+    backoff_delay,
+    basic_auth_header,
+    build_config,
+    env_options,
+    health_outcome,
+    parse_json,
+    parse_token,
+    poll_outcome,
+    prepare_upload,
+    request_id_of,
+    scan_outcome,
+)
+from ._errors import KelpgateError, NetworkError, RequestTimeout
+from ._models import AcceptedScan, Health, PendingScan, ScanResult
+
+
+class AsyncKelpgateClient:
+    """Asyncio client for the Kelpgate malware-scanning API.
+
+    Arguments are the same as :class:`kelpgate.KelpgateClient`.
+
+    Use it as ``async with AsyncKelpgateClient(...) as kg:`` or call :meth:`aclose` when you are done.
+    """
+
+    def __init__(
+        self,
+        api_url: str,
+        *,
+        token_url: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        bearer_token: Optional[str] = None,
+        scope: str = "scan",
+        timeout: float = 60.0,
+        max_retries: int = 2,
+        retry_delay: float = 0.5,
+        max_bytes: Optional[int] = None,
+        transport: Optional[httpx.AsyncBaseTransport] = None,
+    ) -> None:
+        opts: Dict[str, Any] = {}
+        if max_bytes is not None:
+            opts["max_bytes"] = max_bytes
+        self._cfg: Config = build_config(
+            api_url=api_url,
+            token_url=token_url,
+            client_id=client_id,
+            client_secret=client_secret,
+            bearer_token=bearer_token,
+            scope=scope,
+            timeout=timeout,
+            max_retries=max_retries,
+            retry_delay=retry_delay,
+            **opts,
+        )
+        self._http = httpx.AsyncClient(timeout=timeout, transport=transport, follow_redirects=False)
+        self._lock: Optional[asyncio.Lock] = None  # created on first use, inside the running loop
+        self._token: Optional[str] = None
+        self._expires_at = 0.0
+
+    @classmethod
+    def from_env(cls, env: Optional[Mapping[str, str]] = None, **overrides: Any) -> AsyncKelpgateClient:
+        """Build a client from ``KELPGATE_*`` environment variables, like ``KelpgateClient.from_env``."""
+        opts = env_options(env)
+        opts.update(overrides)
+        return cls(**opts)
+
+    def __repr__(self) -> str:
+        return f"AsyncKelpgateClient(api_url={self._cfg.api_url!r})"
+
+    async def __aenter__(self) -> AsyncKelpgateClient:
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        await self.aclose()
+
+    async def aclose(self) -> None:
+        """Release the underlying connections."""
+        await self._http.aclose()
+
+    # ---- authentication
+
+    async def _access_token(self, force: bool = False) -> str:
+        if self._cfg.bearer_token:
+            return self._cfg.bearer_token
+        if not force and self._token and self._expires_at - time.monotonic() > REFRESH_AHEAD:
+            return self._token
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        async with self._lock:  # one renewal at a time; concurrent callers wait, then reuse the result
+            now = time.monotonic()
+            if not force and self._token and self._expires_at - now > REFRESH_AHEAD:
+                return self._token
+            try:
+                token, ttl = await self._fetch_token()
+            except KelpgateError:
+                if not force and self._token and self._expires_at > now:
+                    return self._token  # renewing early failed, the old token is still good
+                raise
+            self._token, self._expires_at = token, time.monotonic() + ttl
+            return token
+
+    async def _fetch_token(self) -> tuple[str, float]:
+        cfg = self._cfg
+        assert cfg.token_url and cfg.client_id and cfg.client_secret
+        try:
+            res = await self._http.post(
+                cfg.token_url,
+                headers={
+                    "Authorization": basic_auth_header(cfg.client_id, cfg.client_secret),
+                    "User-Agent": cfg.user_agent,
+                },
+                data={"grant_type": "client_credentials", "scope": cfg.scope},
+            )
+        except httpx.TimeoutException as exc:
+            raise RequestTimeout(f"The token request timed out after {cfg.timeout:g} s") from exc
+        except httpx.HTTPError as exc:
+            raise NetworkError(f"The token request failed: {exc}") from exc
+        return parse_token(res.status_code, parse_json(res.content), request_id_of(res.headers))
+
+    # ---- HTTP
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json_body: Any = None,
+        upload: Optional[PreparedUpload] = None,
+        form: Optional[Dict[str, str]] = None,
+        auth: bool = True,
+    ) -> tuple[int, Dict[str, Any], httpx.Headers]:
+        cfg = self._cfg
+        token_retried = False
+        attempt = 0
+        while True:
+            headers = {"User-Agent": cfg.user_agent, "Accept": "application/json"}
+            opened: Any = None
+            must_close = False
+            try:
+                if auth:
+                    headers["Authorization"] = f"Bearer {await self._access_token()}"
+                kwargs: Dict[str, Any] = {}
+                if upload is not None:
+                    opened, must_close = upload.open()
+                    kwargs["files"] = {"file": (upload.filename, opened, "application/octet-stream")}
+                    if form:
+                        kwargs["data"] = form
+                elif json_body is not None:
+                    kwargs["content"] = json.dumps(json_body).encode()
+                    headers["Content-Type"] = "application/json"
+                res = await self._http.request(method, cfg.api_url + path, headers=headers, **kwargs)
+            except KelpgateError:
+                raise
+            except httpx.HTTPError as exc:
+                err: KelpgateError = (
+                    RequestTimeout(f"The request timed out after {cfg.timeout:g} s")
+                    if isinstance(exc, httpx.TimeoutException)
+                    else NetworkError(f"The request failed: {exc}")
+                )
+                if attempt >= cfg.max_retries:
+                    raise err from exc
+                await asyncio.sleep(backoff_delay(attempt, cfg.retry_delay))
+                attempt += 1
+                continue
+            finally:
+                if must_close and opened is not None:
+                    opened.close()
+
+            data = parse_json(res.content)
+            if res.status_code == 401 and auth and not cfg.bearer_token and not token_retried:
+                token_retried = True  # the token may have been revoked or rotated: renew once
+                self._token = None
+                continue
+            if res.status_code in RETRY_STATUSES and attempt < cfg.max_retries:
+                await asyncio.sleep(backoff_delay(attempt, cfg.retry_delay, res.headers.get("retry-after")))
+                attempt += 1
+                continue
+            return res.status_code, data, res.headers
+
+    # ---- scanning
+
+    async def scan_file(self, source: Source, *, filename: Optional[str] = None) -> ScanResult:
+        """Scan a file given as a path, ``bytes`` or a binary file object, and return the verdict."""
+        upload = prepare_upload(source, filename, self._cfg.max_bytes)
+        return scan_outcome(*await self._request("POST", "/api/v1/scan/file", upload=upload))
+
+    async def scan_bytes(
+        self, data: Union[bytes, bytearray, memoryview], *, filename: Optional[str] = None
+    ) -> ScanResult:
+        """Scan bytes you already have in memory."""
+        return await self.scan_file(data, filename=filename)
+
+    async def scan_url(self, url: str) -> ScanResult:
+        """Ask the service to fetch and scan a remote file."""
+        return scan_outcome(*await self._request("POST", "/api/v1/scan/url", json_body={"url": str(url)}))
+
+    async def submit_file(
+        self, source: Source, *, filename: Optional[str] = None, webhook_url: Optional[str] = None
+    ) -> AcceptedScan:
+        """Start an asynchronous scan. Collect the result with :meth:`get_result` or :meth:`wait`."""
+        upload = prepare_upload(source, filename, self._cfg.max_bytes)
+        form = {"webhook_url": webhook_url} if webhook_url else None
+        return accepted_outcome(
+            *await self._request("POST", "/api/v1/async-scan/file", upload=upload, form=form)
+        )
+
+    async def submit_url(self, url: str, *, webhook_url: Optional[str] = None) -> AcceptedScan:
+        """Start an asynchronous scan of a remote file."""
+        body: Dict[str, str] = {"url": str(url)}
+        if webhook_url:
+            body["webhook_url"] = webhook_url
+        return accepted_outcome(*await self._request("POST", "/api/v1/async-scan/url", json_body=body))
+
+    async def get_result(self, scan_id: str, *, kind: str = "file") -> Union[ScanResult, PendingScan]:
+        """Read an asynchronous result. Returns a :class:`PendingScan` until the verdict is ready."""
+        which = "url" if kind == "url" else "file"
+        status, data, headers = await self._request(
+            "GET", f"/api/v1/async-scan/{which}?scan_id={quote(scan_id, safe='')}"
+        )
+        return poll_outcome(scan_id, status, data, headers)
+
+    async def wait(
+        self, scan_id: str, *, kind: str = "file", interval: float = 1.0, timeout: float = 120.0
+    ) -> ScanResult:
+        """Poll until an asynchronous scan finishes. Raises :class:`RequestTimeout` if it takes too long."""
+        deadline = time.monotonic() + timeout
+        while True:
+            result = await self.get_result(scan_id, kind=kind)
+            if isinstance(result, ScanResult):
+                return result
+            if time.monotonic() + interval > deadline:
+                raise RequestTimeout(f"Scan {scan_id} was not finished after {timeout:g} s")
+            await asyncio.sleep(interval)
+
+    async def health(self) -> Health:
+        """Check that the service is up. Needs no credentials."""
+        return health_outcome(*await self._request("GET", "/api/v1/health", auth=False))

kelpgate/_common.py

@@ -0,0 +1,335 @@
+"""Logic shared by the synchronous and asynchronous clients: no I/O happens here."""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import random
+import re
+import urllib.parse
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Any, BinaryIO, Dict, Mapping, Optional, Tuple, Union
+
+from ._errors import (
+    AuthError,
+    ConfigError,
+    KelpgateError,
+    PayloadTooLargeError,
+    QuotaExceededError,
+    RateLimitedError,
+    ScanError,
+    ServerError,
+    UnsupportedFileError,
+)
+from ._models import AcceptedScan, Health, PendingScan, ScanResult
+from ._version import __version__
+
+DEFAULT_MAX_BYTES = 128 * 1024 * 1024  # a little above the largest plan limit
+REFRESH_AHEAD = 120.0  # renew a token this many seconds before it expires
+RETRY_AFTER_CAP = 30.0
+RETRY_STATUSES = frozenset({429, 502, 503, 504})
+_LOOPBACK = frozenset({"localhost", "127.0.0.1", "::1"})
+# RFC 6749 asks for form-urlencoding of the client id and secret in Basic auth.
+_SAFE = "!~*'()"
+
+Source = Union[str, "os.PathLike[str]", bytes, bytearray, memoryview, BinaryIO]
+
+
+def check_url(name: str, value: Optional[str]) -> str:
+    """Return ``value`` without a trailing slash, or raise ConfigError. Plain http is only for localhost."""
+    parts = urllib.parse.urlsplit(value or "")
+    if parts.scheme not in ("http", "https") or not parts.hostname:
+        raise ConfigError(f"{name} is not a valid URL")
+    if parts.scheme == "http" and parts.hostname not in _LOOPBACK:
+        raise ConfigError(f"{name} must use https (plain http is only allowed for localhost)")
+    return (value or "").rstrip("/")
+
+
+@dataclass
+class Config:
+    api_url: str
+    token_url: Optional[str]
+    client_id: Optional[str]
+    client_secret: Optional[str]
+    bearer_token: Optional[str]
+    scope: str
+    timeout: float
+    max_retries: int
+    retry_delay: float
+    max_bytes: int
+    user_agent: str
+
+    def __repr__(self) -> str:  # never show the secret or token
+        return f"Config(api_url={self.api_url!r}, scope={self.scope!r})"
+
+
+def build_config(
+    *,
+    api_url: Optional[str],
+    token_url: Optional[str] = None,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
+    bearer_token: Optional[str] = None,
+    scope: str = "scan",
+    timeout: float = 60.0,
+    max_retries: int = 2,
+    retry_delay: float = 0.5,
+    max_bytes: int = DEFAULT_MAX_BYTES,
+) -> Config:
+    checked_api = check_url("api_url", api_url)
+    checked_token: Optional[str] = None
+    if not bearer_token:
+        checked_token = check_url("token_url", token_url)
+        if not client_id or not client_secret:
+            raise ConfigError("client_id and client_secret are required (or pass bearer_token)")
+    return Config(
+        api_url=checked_api,
+        token_url=checked_token,
+        client_id=client_id,
+        client_secret=client_secret,
+        bearer_token=bearer_token or None,
+        scope=scope,
+        timeout=float(timeout),
+        max_retries=max(0, int(max_retries)),
+        retry_delay=float(retry_delay),
+        max_bytes=int(max_bytes),
+        user_agent=f"kelpgate-python/{__version__}",
+    )
+
+
+def env_options(env: Optional[Mapping[str, str]] = None) -> Dict[str, Any]:
+    """Client options from ``KELPGATE_*`` environment variables, with the production endpoints as defaults."""
+    e = os.environ if env is None else env
+    return {
+        "api_url": e.get("KELPGATE_API_URL") or "https://api.kelpgate.cloud",
+        "token_url": e.get("KELPGATE_TOKEN_URL") or "https://auth.kelpgate.cloud/oauth2/token",
+        "client_id": e.get("KELPGATE_CLIENT_ID"),
+        "client_secret": e.get("KELPGATE_CLIENT_SECRET"),
+        "bearer_token": e.get("KELPGATE_TOKEN"),
+    }
+
+
+def basic_auth_header(client_id: str, client_secret: str) -> str:
+    raw = f"{urllib.parse.quote(client_id, safe=_SAFE)}:{urllib.parse.quote(client_secret, safe=_SAFE)}"
+    return "Basic " + base64.b64encode(raw.encode()).decode()
+
+
+def parse_json(content: bytes) -> Dict[str, Any]:
+    if not content:
+        return {}
+    try:
+        data = json.loads(content)
+    except ValueError:
+        return {"error": content[:200].decode("utf-8", "replace")}
+    return data if isinstance(data, dict) else {"value": data}
+
+
+def parse_token(status: int, data: Dict[str, Any], request_id: Optional[str]) -> Tuple[str, float]:
+    """Return ``(access_token, lifetime_seconds)`` or raise AuthError."""
+    token = data.get("access_token")
+    if status != 200 or not isinstance(token, str) or not token:
+        why = data.get("error_description") or data.get("error") or f"HTTP {status}"
+        safe_body = {k: v for k, v in data.items() if k not in ("access_token", "refresh_token")}
+        raise AuthError(
+            f"Could not get an access token: {why}",
+            code="TOKEN_REQUEST_FAILED",
+            status=status,
+            request_id=request_id,
+            body=safe_body,
+        )
+    try:
+        ttl = float(data.get("expires_in") or 0)
+    except (TypeError, ValueError):
+        ttl = 0.0
+    return token, ttl if ttl > 0 else 300.0
+
+
+def backoff_delay(attempt: int, base: float, retry_after: Optional[str] = None) -> float:
+    if retry_after is not None:
+        try:
+            seconds = float(retry_after)
+        except ValueError:
+            seconds = -1.0
+        if seconds >= 0:
+            return min(seconds, RETRY_AFTER_CAP)
+    return base * (2.0**attempt) + random.uniform(0, min(0.1, base))  # noqa: S311 - retry jitter, not security
+
+
+def _basename(name: str) -> str:
+    return re.split(r"[\\/]", name)[-1] or "upload.bin"
+
+
+class PreparedUpload:
+    """A file ready to send. :meth:`open` gives a fresh file object for each attempt, so retries work."""
+
+    def __init__(
+        self,
+        filename: str,
+        size: int,
+        *,
+        data: Optional[bytes] = None,
+        path: Optional[str] = None,
+        stream: Optional[BinaryIO] = None,
+    ) -> None:
+        self.filename = filename
+        self.size = size
+        self._data = data
+        self._path = path
+        self._stream = stream
+        self._start = stream.tell() if stream is not None else 0
+
+    def open(self) -> Tuple[Any, bool]:
+        """Return ``(content, must_close)``."""
+        if self._data is not None:
+            return self._data, False
+        if self._path is not None:
+            return open(self._path, "rb"), True
+        assert self._stream is not None
+        self._stream.seek(self._start)
+        return self._stream, False
+
+
+def prepare_upload(source: Source, filename: Optional[str], max_bytes: int) -> PreparedUpload:
+    """Turn a path, bytes, or binary file object into something we can send, enforcing ``max_bytes``."""
+
+    def too_big(size: int) -> PayloadTooLargeError:
+        return PayloadTooLargeError(
+            f"The file is {size} bytes, over this client's limit of {max_bytes} bytes", status=413
+        )
+
+    upload: PreparedUpload
+    if isinstance(source, (str, os.PathLike)):
+        path = os.fspath(source)
+        try:
+            size = os.path.getsize(path)
+            with open(path, "rb"):
+                pass
+        except OSError as exc:
+            raise ScanError(f"Cannot read {path}: {exc.strerror or exc}", code="FILE_UNREADABLE") from exc
+        upload = PreparedUpload(filename or _basename(path), size, path=path)
+    elif isinstance(source, (bytes, bytearray, memoryview)):
+        data = bytes(source)
+        upload = PreparedUpload(filename or "upload.bin", len(data), data=data)
+    elif hasattr(source, "read"):
+        stream: BinaryIO = source
+        name = filename
+        if not name:
+            attr = getattr(stream, "name", None)
+            name = _basename(attr) if isinstance(attr, str) else "upload.bin"
+        seekable = False
+        try:
+            seekable = bool(stream.seekable())
+        except (AttributeError, ValueError):
+            seekable = False
+        if seekable:
+            start = stream.tell()
+            end = stream.seek(0, os.SEEK_END)
+            stream.seek(start)
+            if end - start > max_bytes:
+                raise too_big(end - start)
+            upload = PreparedUpload(name, end - start, stream=stream)
+        else:  # a pipe or socket: buffer it, but never more than the limit
+            buf = BytesIO()
+            while True:
+                chunk = stream.read(1024 * 1024)
+                if not chunk:
+                    break
+                buf.write(chunk)
+                if buf.tell() > max_bytes:
+                    raise too_big(buf.tell())
+            data = buf.getvalue()
+            upload = PreparedUpload(name, len(data), data=data)
+    else:
+        raise ConfigError("source must be a file path, bytes, or a binary file object")
+    if upload.size > max_bytes:
+        raise too_big(upload.size)
+    return upload
+
+
+def request_id_of(headers: Mapping[str, str]) -> Optional[str]:
+    return headers.get("x-request-id")
+
+
+def to_error(status: int, data: Dict[str, Any], headers: Mapping[str, str]) -> KelpgateError:
+    """Map a failed response to the right exception."""
+    message = str(data.get("error") or data.get("message") or f"Kelpgate returned HTTP {status}")
+    kw: Dict[str, Any] = {"status": status, "request_id": request_id_of(headers), "body": data}
+    if status in (401, 403):
+        return AuthError(message, code="UNAUTHORIZED" if status == 401 else "FORBIDDEN", **kw)
+    if status == 402:
+        return QuotaExceededError(message, **kw)
+    if status == 413:
+        return PayloadTooLargeError(message, **kw)
+    if status == 415:
+        return UnsupportedFileError(message, **kw)
+    if status == 429:
+        try:
+            retry_after: Optional[float] = float(headers.get("retry-after", ""))
+        except ValueError:
+            retry_after = None
+        return RateLimitedError(message, retry_after=retry_after, **kw)
+    if status == 404:
+        return KelpgateError(message, code="NOT_FOUND", **kw)
+    if status in (400, 412):
+        return ScanError(message, **kw)
+    if status >= 500:
+        return ServerError(message, **kw)
+    return KelpgateError(message, **kw)
+
+
+def to_result(data: Dict[str, Any], headers: Mapping[str, str]) -> ScanResult:
+    signature = str(data.get("av-signature") or "")
+    authors = data.get("rule_authors")
+    return ScanResult(
+        infected=data.get("av-status") == "INFECTED",
+        status=str(data.get("av-status")),
+        signature=signature,
+        sha256=data.get("sha256"),
+        rule_authors=tuple(str(a) for a in authors) if isinstance(authors, list) else (),
+        filename=data.get("filename"),
+        scan_id=data.get("scan_id"),
+        timestamp=data.get("av-timestamp"),
+        s3_path=data.get("s3_path"),
+        download_url=data.get("download_url"),
+        request_id=request_id_of(headers),
+        policy_rejected=signature.startswith("POLICY:"),
+        degraded="scan coverage reduced" in signature,
+        raw=data,
+    )
+
+
+def scan_outcome(status: int, data: Dict[str, Any], headers: Mapping[str, str]) -> ScanResult:
+    """A scan answers 200 (clean) or 406 (infected). Both are verdicts. Anything else is an error."""
+    if status in (200, 406) and data.get("av-status"):
+        return to_result(data, headers)
+    raise to_error(status, data, headers)
+
+
+def accepted_outcome(status: int, data: Dict[str, Any], headers: Mapping[str, str]) -> AcceptedScan:
+    if status == 202 and data.get("scan_id"):
+        return AcceptedScan(
+            scan_id=str(data["scan_id"]),
+            filename=data.get("filename"),
+            message=data.get("message"),
+            request_id=request_id_of(headers),
+        )
+    raise to_error(status, data, headers)
+
+
+def poll_outcome(
+    scan_id: str, status: int, data: Dict[str, Any], headers: Mapping[str, str]
+) -> Union[ScanResult, PendingScan]:
+    if status == 202 or data.get("av-status") == "PROCESSING":
+        return PendingScan(scan_id=scan_id, request_id=request_id_of(headers))
+    return scan_outcome(status, data, headers)
+
+
+def health_outcome(status: int, data: Dict[str, Any], headers: Mapping[str, str]) -> Health:
+    return Health(
+        ok=status == 200,
+        status=status,
+        signature_version=data.get("signature_version"),
+        request_id=request_id_of(headers),
+    )