broker/coreci_broker_client.py

"""Public stdlib-only client for the CoreCI Broker: idempotent tool registration, agent-state persistence, and path-scoped secret reads."""
import json
import time
import urllib.parse
import urllib.request

REGISTER_PATH = "/api/broker/agents/register"
SECRETS_PREFIX = "/api/broker/secrets/"
CLIENT_MODULE_PATH = "/static/coreci_broker_client.py"


class BrokerError(Exception):
    def __init__(self, status, message):
        super().__init__(message)
        self.status = status
        self.message = message


def _broker_error(status, body):
    message = "broker returned HTTP %s" % status
    if isinstance(body, dict) and body.get("message"):
        message = str(body["message"])
    return BrokerError(status, message)


def _epoch_from_iso(value):
    try:
        year, month, day = int(value[0:4]), int(value[5:7]), int(value[8:10])
        hour, minute, second = int(value[11:13]), int(value[14:16]), int(value[17:19])
        rest = value[19:]
        micro = 0
        if rest[:1] == ".":
            digits = ""
            for char in rest[1:]:
                if not char.isdigit():
                    break
                digits += char
            micro = int((digits + "000000")[:6])
            rest = rest[1 + len(digits):]
        offset = 0
        if rest[:1] in ("+", "-"):
            sign = 1 if rest[0] == "+" else -1
            offset = sign * (int(rest[1:3]) * 3600 + int(rest[4:6]) * 60)
        if not 1 <= month <= 12 or not 1 <= day <= 31:
            return None
        if hour > 23 or minute > 59 or second > 59:
            return None
        year -= month <= 2
        era = (year if year >= 0 else year - 399) // 400
        yoe = year - era * 400
        doy = (153 * (month - 3 if month > 2 else month + 9) + 2) // 5 + day - 1
        doe = yoe * 365 + yoe // 4 - yoe // 100 + doy
        return (era * 146097 + doe - 719468) * 86400 + hour * 3600 + minute * 60 + second - offset + micro / 1000000.0
    except (ValueError, TypeError, IndexError):
        return None


def _token_live(state):
    token = state.get("current_token")
    expires = _epoch_from_iso(state.get("token_expires_at"))
    return isinstance(token, str) and bool(token) and expires is not None and expires > time.time()


class BrokerClient:
    def __init__(self, base_url, agent_id, app_uuid, user_id, tags=None, timeout=10.0):
        self.base_url = str(base_url).rstrip("/")
        self.agent_id = agent_id
        self.app_uuid = app_uuid
        self.user_id = user_id
        self.tags = list(tags or [])
        self.timeout = timeout

    def _app_headers(self, state=None):
        token = (state or {}).get("current_token")
        if not isinstance(token, str) or not token:
            return {}
        return {"Authorization": "App %s %s" % (self.app_uuid, token)}

    def _request(self, method, path, payload=None, headers=None):
        data = None if payload is None else json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(self.base_url + path, data=data, method=method)
        if data is not None:
            request.add_header("Content-Type", "application/json")
        for key, value in (headers or {}).items():
            request.add_header(key, value)
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                status, raw = response.status, response.read()
        except urllib.request.HTTPError as error:
            status, raw = error.code, error.read()
        body = None
        text = raw.decode("utf-8", "replace")
        if text:
            try:
                body = json.loads(text)
            except ValueError:
                body = None
        return status, body

    def _authorized_request(self, method, path, state, payload=None):
        status, body = self._request(method, path, payload, self._app_headers(state))
        if status == 401:
            self.register(state)
            status, body = self._request(method, path, payload, self._app_headers(state))
        return status, body

    def _state_path(self):
        return "/api/broker/agents/%s/state" % urllib.parse.quote(self.agent_id, safe="")

    def register(self, state):
        if state is None:
            state = {}
        payload = {"agent_id": self.agent_id, "app_registration_uuid": self.app_uuid,
                   "user_id": self.user_id, "tags": list(self.tags)}
        status, data = self._request("POST", REGISTER_PATH, payload, self._app_headers(state))
        if status not in (200, 201) or not isinstance(data, dict):
            raise _broker_error(status, data)
        state["agent_id"] = self.agent_id
        state["app_registration_uuid"] = self.app_uuid
        state["user_id"] = self.user_id
        state["tags"] = list(self.tags)
        state["current_token"] = data["token"]
        state["token_expires_at"] = data["token_expires_at"]
        pepper = data.get("pepper")
        if pepper is not None:
            state["pepper"] = pepper
        return state

    def ensure_registered(self, state):
        if state is None:
            state = {}
        if _token_live(state):
            return state
        return self.register(state)

    def get_state(self, state):
        status, body = self._authorized_request("GET", self._state_path(), state)
        if status == 200:
            return body
        if status == 404:
            return None
        raise _broker_error(status, body)

    def put_state(self, state):
        status, body = self._authorized_request("PUT", self._state_path(), state, state)
        if status == 200:
            return True
        raise _broker_error(status, body)

    def read_secret(self, path, state):
        suffix = urllib.parse.quote(str(path), safe="/")
        status, body = self._authorized_request("GET", SECRETS_PREFIX + suffix, state)
        if status == 200:
            return body
        raise _broker_error(status, body)


def fetch_client_module(base_url, timeout=10.0):
    url = str(base_url).rstrip("/") + CLIENT_MODULE_PATH
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return response.read().decode("utf-8")