broker/tests/test_reference_tools.py

import ast
import base64
import contextlib
import importlib.util
import io
import json
import os
import sys
import tempfile
import time
import unittest
import urllib.error
import urllib.request
import uuid as uuidlib

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)))

import coreci_broker

ADMIN_USER = "admin"
ADMIN_PASS = "correct-horse-battery-staple"
APPS_PATH = "/api/broker/app_registrations"
USERS_PATH = "/api/broker/user_allowlist"
PROXMOX_SECRET = {"url": "https://pve.coreci.dev:8006", "token": "pve-tok"}
REFERENCE_DIR = os.path.abspath(
    os.path.join(os.path.dirname(__file__), os.pardir, "reference_tools")
)
REFERENCE_TITLES = {
    "proxmox_readonly_reference": "Proxmox Read-Only Reference [CoreCI]",
    "minimal_auth_showcase": "Minimal Auth Showcase [CoreCI]",
}
ALLOWED_REFERENCE_IMPORTS = {"coreci_broker_client", "datetime"}
FORBIDDEN_SOURCE_MARKERS = ("urllib", "httpx", "socket", "proxmoxer")
BROKER_ONLY_AUDIT_ACTIONS = {"add_app", "add_user", "register", "read"}


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


def _load_reference(name):
    spec = importlib.util.spec_from_file_location(
        name, os.path.join(REFERENCE_DIR, name + ".py")
    )
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _reference_source(name):
    with open(os.path.join(REFERENCE_DIR, name + ".py"), "r", encoding="utf-8") as handle:
        return handle.read()


PROXMOX_REFERENCE = _load_reference("proxmox_readonly_reference")
MINIMAL_SHOWCASE = _load_reference("minimal_auth_showcase")


class _ReferenceCase(unittest.TestCase):
    extra_env = {}

    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.tmp = tmp.name
        self.creds_path = os.path.join(self.tmp, "creds.json")
        with open(self.creds_path, "w", encoding="utf-8") as handle:
            handle.write(json.dumps({"coreci/proxmox": PROXMOX_SECRET}))
        self.env = {
            "BROKER_ENV_PATH": os.path.join(self.tmp, ".env"),
            "BROKER_DB_PATH": os.path.join(self.tmp, "broker.db"),
            "BROKER_AUDIT_LOG_PATH": os.path.join(self.tmp, "broker.audit.log"),
            "BROKER_ADMIN_USER": ADMIN_USER,
            "BROKER_ADMIN_PASS": ADMIN_PASS,
            "BROKER_CREDS_SOURCE": "file",
            "BROKER_CREDS_FILE": self.creds_path,
        }
        self.env.update(self.extra_env)
        with contextlib.redirect_stdout(io.StringIO()):
            self.broker = coreci_broker.boot(self.env)
        self.addCleanup(self.broker.audit.close)
        self.addCleanup(self.broker.store.close)
        self.server = coreci_broker.serve("127.0.0.1", 0, self.broker)
        self.addCleanup(self.server.shutdown)
        self.addCleanup(self.server.server_close)
        self.base = "http://127.0.0.1:%d" % self.server.server_address[1]
        self.admin = {"Authorization": _basic_header(ADMIN_USER, ADMIN_PASS)}
        self.app_uuid = str(uuidlib.uuid4())
        self.user_id = "ceo"
        self._seed()

    def _seed(self):
        status, _ = self._request("POST", APPS_PATH, {"uuid": self.app_uuid}, self.admin)
        self.assertEqual(201, status)
        status, _ = self._request("POST", USERS_PATH, {"user_id": self.user_id}, self.admin)
        self.assertEqual(201, status)

    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 + path, data=data, method=method)
        for key, value in (headers or {}).items():
            request.add_header(key, value)
        try:
            with urllib.request.urlopen(request, timeout=10) as response:
                raw = response.read().decode("utf-8")
                return response.status, json.loads(raw) if raw else None
        except urllib.error.HTTPError as error:
            raw = error.read().decode("utf-8")
            return error.code, json.loads(raw) if raw else None

    def _audit(self, action=None, result=None):
        with open(self.env["BROKER_AUDIT_LOG_PATH"], "r", encoding="utf-8") as handle:
            entries = [json.loads(raw) for raw in handle if raw.strip()]
        return [
            entry for entry in entries
            if (action is None or entry["action"] == action)
            and (result is None or entry["result"] == result)
        ]

    def _tools(self, module, agent_id="proxmox"):
        tools = module.Tools()
        tools.valves.base_url = self.base
        tools.valves.agent_id = agent_id
        tools.valves.app_uuid = self.app_uuid
        tools.valves.user_id = self.user_id
        return tools

    def _persisted_state(self, agent_id="proxmox"):
        return self.broker.store.get_agent_state(agent_id, self.app_uuid)


class ProxmoxReferenceFlowTests(_ReferenceCase):
    def test_first_call_registers_persists_state_reads_secret_dry_runs(self):
        tools = self._tools(PROXMOX_REFERENCE)
        response = tools.list_nodes()
        self.assertIs(True, response["dry_run"])
        self.assertEqual("list_nodes", response["action"])
        self.assertIs(True, response["credential_present"])
        self.assertTrue(str(response["note"]).strip())
        self.assertNotIn("data", response)
        self.assertNotIn("pve-tok", json.dumps(response))
        self.assertEqual(1, len(self._audit(action="register")))
        document = self._persisted_state()
        self.assertIsNotNone(document)
        self.assertIn("current_token", document)
        self.assertEqual(tools._state["current_token"], document["current_token"])
        self.assertEqual([], self._audit(result="denied"))
        self.assertTrue(
            {entry["action"] for entry in self._audit()} <= BROKER_ONLY_AUDIT_ACTIONS
        )

    def test_second_call_reuses_persisted_state_without_new_register(self):
        tools = self._tools(PROXMOX_REFERENCE)
        tools.list_nodes()
        token = tools._state["current_token"]
        second = tools.list_nodes()
        self.assertIs(True, second["dry_run"])
        self.assertEqual("list_nodes", second["action"])
        self.assertIs(True, second["credential_present"])
        self.assertEqual(1, len(self._audit(action="register")))
        self.assertEqual(token, tools._state["current_token"])
        document = self._persisted_state()
        self.assertIn("current_token", document)
        self.assertEqual(token, document["current_token"])

    def test_vm_status_is_dry_run_with_credential_present(self):
        tools = self._tools(PROXMOX_REFERENCE)
        response = tools.vm_status("101")
        self.assertIs(True, response["dry_run"])
        self.assertEqual("vm_status", response["action"])
        self.assertIs(True, response["credential_present"])
        self.assertNotIn("pve-tok", json.dumps(response))
        self.assertEqual(1, len(self._audit(action="register")))


class ProxmoxReferenceTtlTests(_ReferenceCase):
    extra_env = {"BROKER_TOKEN_TTL": "1"}

    def test_expired_subsequent_call_reregisters_and_persists_new_expiry(self):
        tools = self._tools(PROXMOX_REFERENCE)
        self.assertTrue(tools.list_nodes()["credential_present"])
        first_expires = tools._state["token_expires_at"]
        time.sleep(1.2)
        response = tools.list_nodes()
        self.assertIs(True, response["dry_run"])
        self.assertIs(True, response["credential_present"])
        self.assertEqual(2, len(self._audit(action="register")))
        document = self._persisted_state()
        self.assertEqual(tools._state["current_token"], document["current_token"])
        self.assertEqual(tools._state["token_expires_at"], document["token_expires_at"])
        self.assertNotEqual(first_expires, document["token_expires_at"])


class MinimalShowcaseTests(_ReferenceCase):
    def test_demo_authenticates_then_reads_secret(self):
        tools = self._tools(MINIMAL_SHOWCASE)
        response = tools.demo()
        self.assertEqual(
            {"registered": True, "secret_read": True,
             "path": "coreci/proxmox", "note": "auth first, then use"},
            response,
        )
        self.assertNotIn("pve-tok", json.dumps(response))
        self.assertEqual(1, len(self._audit(action="register")))
        self.assertEqual([], self._audit(result="denied"))


class ReferenceHygieneTests(unittest.TestCase):
    def test_sources_carry_no_secret_literals(self):
        for name in REFERENCE_TITLES:
            source = _reference_source(name)
            self.assertNotRegex(source, r"(?im)^\s*(password|pepper)\s*=\s*[\"']", name)
            for value in PROXMOX_SECRET.values():
                self.assertNotIn(value, source, name)

    def test_sources_use_no_direct_network_or_infra_clients(self):
        for name in REFERENCE_TITLES:
            source = _reference_source(name)
            for marker in FORBIDDEN_SOURCE_MARKERS:
                self.assertNotIn(marker, source, "%s: %s" % (name, marker))

    def test_imports_limited_to_client_module_and_stdlib(self):
        for name in REFERENCE_TITLES:
            tree = ast.parse(_reference_source(name))
            modules = set()
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    modules.update(alias.name for alias in node.names)
                elif isinstance(node, ast.ImportFrom):
                    modules.add(node.module or "")
            for module in modules:
                self.assertIn(module, ALLOWED_REFERENCE_IMPORTS, "%s: %s" % (name, module))

    def test_owui_frontmatter_docstrings(self):
        for name, title in REFERENCE_TITLES.items():
            doc = ast.get_docstring(ast.parse(_reference_source(name)))
            self.assertTrue(doc.startswith("title: " + title), name)
            for line in ("author: coreci", "description: ", "version: 0.1.0",
                         "requirements: []", "licence: MIT"):
                self.assertIn(line, doc, "%s: %s" % (name, line))

    def test_tools_surface_valves_state_and_docstring_discipline(self):
        for name in REFERENCE_TITLES:
            module = _load_reference(name)
            tools = module.Tools()
            self.assertIsInstance(tools.valves, module.Tools.Valves)
            for attr in ("base_url", "agent_id", "app_uuid", "user_id"):
                self.assertEqual("", getattr(module.Tools.Valves, attr), attr)
                self.assertEqual("", getattr(tools.valves, attr), attr)
            self.assertEqual({}, tools._state)
            public = []
            for node in ast.walk(ast.parse(_reference_source(name))):
                if not isinstance(node, ast.FunctionDef):
                    continue
                doc = ast.get_docstring(node)
                if node.name.startswith("_"):
                    self.assertIsNone(doc, "%s.%s" % (name, node.name))
                else:
                    self.assertIsNotNone(doc, "%s.%s" % (name, node.name))
                    public.append(node.name)
            self.assertTrue(public, name)

    def test_no_comments_in_sources(self):
        for name in REFERENCE_TITLES:
            for number, line in enumerate(_reference_source(name).splitlines(), 1):
                self.assertFalse(line.lstrip().startswith("#"), "%s:%d" % (name, number))


if __name__ == "__main__":
    unittest.main()