graylog-ops/graylog_ops/http.py

"""Minimal HTTP layer for the Graylog REST API.

Deliberately built on urllib so the tool has no third-party dependencies and
can be copied onto old jump hosts. Handles the three things Graylog's API
demands that plain urllib does not do for you:

* HTTP Basic auth (API token used as username with the literal password
  ``token``, or a real username/password pair),
* the ``X-Requested-By`` header required by Graylog's CSRF protection on all
  non-GET requests,
* bounded retries with backoff for 429/5xx and transient connection errors.
"""

import base64
import json
import random
import socket
import ssl
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

USER_AGENT = "graylog-ops/0.1 (read-only recon)"
RETRY_STATUS = frozenset((429, 500, 502, 503, 504))


class GraylogError(Exception):
    """Any failure talking to the Graylog API."""

    def __init__(self, message, status=None, body=None, path=None, method=None):
        super().__init__(message)
        self.message = message
        self.status = status
        self.body = body
        self.path = path
        self.method = method

    def __str__(self):
        parts = []
        if self.status is not None:
            parts.append("HTTP %s" % self.status)
        if self.method and self.path:
            parts.append("%s %s" % (self.method, self.path))
        parts.append(self.message)
        return " | ".join(parts)


def _basic_auth(username, password):
    raw = ("%s:%s" % (username, password)).encode("utf-8")
    return "Basic " + base64.b64encode(raw).decode("ascii")


def _error_detail(body):
    """Pull a human readable message out of a Graylog/ES error body."""
    if not body:
        return ""
    try:
        data = json.loads(body)
    except (TypeError, ValueError):
        return body.strip()[:300]
    if isinstance(data, dict):
        for key in ("message", "error", "description"):
            value = data.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
        if isinstance(data.get("error"), dict):
            reason = data["error"].get("reason")
            if isinstance(reason, str):
                return reason
    return body.strip()[:300]


class HttpClient:
    def __init__(
        self,
        base_url,
        username=None,
        password=None,
        timeout=30.0,
        retries=2,
        insecure=False,
        verbose=False,
    ):
        self.base_url = base_url.rstrip("/")
        self.timeout = float(timeout)
        self.retries = max(0, int(retries))
        self.verbose = verbose
        self.auth = None
        if username is not None:
            self.auth = _basic_auth(username, "" if password is None else password)

        handlers = []
        if insecure:
            context = ssl.create_default_context()
            context.check_hostname = False
            context.verify_mode = ssl.CERT_NONE
            handlers.append(urllib.request.HTTPSHandler(context=context))
        # build_opener keeps the default ProxyHandler, so HTTP(S)_PROXY env vars
        # keep working for environments that need a proxy to reach Graylog.
        self.opener = urllib.request.build_opener(*handlers)

    # -- url helpers -----------------------------------------------------
    def url(self, path, params=None):
        if not path.startswith("/"):
            path = "/" + path
        target = self.base_url + path
        if params:
            clean = {}
            for key, value in params.items():
                if value is None:
                    continue
                if isinstance(value, bool):
                    value = "true" if value else "false"
                clean[key] = value
            if clean:
                sep = "&" if "?" in target else "?"
                target += sep + urllib.parse.urlencode(clean, doseq=True)
        return target

    # -- request core ----------------------------------------------------
    def request(self, method, path, params=None, body=None, expect_json=True):
        method = method.upper()
        url = self.url(path, params)
        payload = None
        headers = {
            "Accept": "application/json",
            "Accept-Encoding": "identity",
            "User-Agent": USER_AGENT,
            # Graylog rejects state changing requests without this header.
            "X-Requested-By": "graylog-ops",
        }
        if body is not None:
            payload = json.dumps(body).encode("utf-8")
            headers["Content-Type"] = "application/json"
        if self.auth:
            headers["Authorization"] = self.auth

        started = time.time()
        last_error = None
        for attempt in range(self.retries + 1):
            request = urllib.request.Request(
                url, data=payload, headers=headers, method=method
            )
            try:
                with self.opener.open(request, timeout=self.timeout) as response:
                    status = response.getcode()
                    raw = response.read()
                    self._debug(method, url, status, started)
                    if not raw.strip():
                        return None
                    text = raw.decode("utf-8", "replace")
                    if not expect_json:
                        return text
                    try:
                        return json.loads(text)
                    except ValueError:
                        raise GraylogError(
                            "response was not JSON",
                            status=status,
                            body=text[:500],
                            path=path,
                            method=method,
                        )
            except urllib.error.HTTPError as exc:
                text = exc.read().decode("utf-8", "replace")
                detail = _error_detail(text)
                self._debug(method, url, exc.code, started)
                last_error = GraylogError(
                    detail or ("request rejected" if exc.code != 401 else "authentication failed"),
                    status=exc.code,
                    body=text[:2000],
                    path=path,
                    method=method,
                )
                if exc.code in RETRY_STATUS and attempt < self.retries:
                    self._sleep(attempt, exc.headers.get("Retry-After"))
                    continue
                raise last_error
            except (urllib.error.URLError, socket.timeout, ConnectionError, OSError) as exc:
                reason = getattr(exc, "reason", exc)
                self._debug(method, url, "ERR", started)
                last_error = GraylogError(
                    "connection failed: %s" % reason,
                    path=path,
                    method=method,
                )
                if attempt < self.retries:
                    self._sleep(attempt, None)
                    continue
                raise last_error
        raise last_error  # pragma: no cover - loop always raises or returns

    # -- internals -------------------------------------------------------
    def _sleep(self, attempt, retry_after):
        delay = (1.5 ** attempt) + random.uniform(0, 0.3)
        try:
            if retry_after is not None:
                delay = max(delay, float(retry_after))
        except (TypeError, ValueError):
            pass
        time.sleep(min(delay, 30.0))

    def _debug(self, method, url, status, started):
        if not self.verbose:
            return
        elapsed = int((time.time() - started) * 1000)
        sys.stderr.write("[http] %s %s -> %s (%dms)\n" % (method, url, status, elapsed))
        sys.stderr.flush()

    # -- convenience -----------------------------------------------------
    def get(self, path, params=None):
        return self.request("GET", path, params=params)