From 64e0c05d6697796bae8137f124c66ba64eada2d8 Mon Sep 17 00:00:00 2001 From: Nicole Rappe Date: Sat, 18 Oct 2025 03:41:29 -0600 Subject: [PATCH] Share installer codes across agent contexts --- Data/Agent/agent.py | 103 +++++++++++++++++++++++++++- Data/Agent/security.py | 105 +++++++++++++++++++++++++++++ tests/test_agent_installer_code.py | 58 ++++++++++++++++ 3 files changed, 263 insertions(+), 3 deletions(-) create mode 100644 tests/test_agent_installer_code.py diff --git a/Data/Agent/agent.py b/Data/Agent/agent.py index b54c5f8..ab73c6d 100644 --- a/Data/Agent/agent.py +++ b/Data/Agent/agent.py @@ -81,6 +81,7 @@ def _bootstrap_log(msg: str): # Headless/service mode flag (skip Qt and interactive UI) SYSTEM_SERVICE_MODE = ('--system-service' in sys.argv) or (os.environ.get('BOREALIS_AGENT_MODE') == 'system') SERVICE_MODE = 'system' if SYSTEM_SERVICE_MODE else 'currentuser' +SERVICE_MODE_CANONICAL = SERVICE_MODE.upper() _bootstrap_log(f'agent.py loaded; SYSTEM_SERVICE_MODE={SYSTEM_SERVICE_MODE}; argv={sys.argv!r}') def _argv_get(flag: str, default: str = None): try: @@ -571,6 +572,57 @@ DEFAULT_CONFIG = { "installer_code": "" } + +def _load_installer_code_from_file(path: str) -> str: + try: + with open(path, "r", encoding="utf-8") as fh: + data = json.load(fh) + except Exception: + return "" + value = data.get("installer_code") if isinstance(data, dict) else "" + if isinstance(value, str): + return value.strip() + return "" + + +def _fallback_installer_code(current_path: str) -> str: + settings_dir = os.path.dirname(current_path) + candidates: List[str] = [] + suffix = CONFIG_SUFFIX_CANONICAL + sibling_map = { + "SYSTEM": "agent_settings_CURRENTUSER.json", + "CURRENTUSER": "agent_settings_SYSTEM.json", + } + sibling_name = sibling_map.get(suffix or "") + if sibling_name: + candidates.append(os.path.join(settings_dir, sibling_name)) + # Prefer the shared/base config next + candidates.append(os.path.join(settings_dir, "agent_settings.json")) + # Legacy location fallback + try: + project_root = _find_project_root() + legacy_dir = os.path.join(project_root, "Agent", "Settings") + if sibling_name: + candidates.append(os.path.join(legacy_dir, sibling_name)) + candidates.append(os.path.join(legacy_dir, "agent_settings.json")) + except Exception: + pass + + current_abspath = os.path.abspath(current_path) + for candidate in candidates: + if not candidate: + continue + try: + candidate_path = os.path.abspath(candidate) + except Exception: + continue + if candidate_path == current_abspath or not os.path.isfile(candidate_path): + continue + code = _load_installer_code_from_file(candidate_path) + if code: + return code + return "" + class ConfigManager: def __init__(self, path): self.path = path @@ -1077,12 +1129,50 @@ class AgentHttpClient: def _resolve_installer_code(self) -> str: if INSTALLER_CODE_OVERRIDE: - return INSTALLER_CODE_OVERRIDE + code = INSTALLER_CODE_OVERRIDE.strip() + if code: + try: + self.key_store.cache_installer_code(code, consumer=SERVICE_MODE_CANONICAL) + except Exception: + pass + return code + code = "" try: code = (CONFIG.data.get("installer_code") or "").strip() - return code except Exception: - return "" + code = "" + if code: + try: + self.key_store.cache_installer_code(code, consumer=SERVICE_MODE_CANONICAL) + except Exception: + pass + return code + try: + cached = self.key_store.load_cached_installer_code() + except Exception: + cached = None + if cached: + try: + self.key_store.cache_installer_code(cached, consumer=SERVICE_MODE_CANONICAL) + except Exception: + pass + return cached + fallback = _fallback_installer_code(CONFIG.path) + if fallback: + try: + CONFIG.data["installer_code"] = fallback + CONFIG._write() + _log_agent( + "Adopted installer code from sibling configuration", fname="agent.log" + ) + except Exception: + pass + try: + self.key_store.cache_installer_code(fallback, consumer=SERVICE_MODE_CANONICAL) + except Exception: + pass + return fallback + return "" def _consume_installer_code(self) -> None: # Avoid clearing explicit CLI/env overrides; only mutate persisted config. @@ -1096,6 +1186,13 @@ class AgentHttpClient: _log_agent("Cleared persisted installer code after successful enrollment", fname="agent.log") except Exception as exc: _log_agent(f"Failed to clear installer code after enrollment: {exc}", fname="agent.error.log") + try: + self.key_store.mark_installer_code_consumed(SERVICE_MODE_CANONICAL) + except Exception as exc: + _log_agent( + f"Failed to update shared installer code cache: {exc}", + fname="agent.error.log", + ) # ------------------------------------------------------------------ # HTTP helpers diff --git a/Data/Agent/security.py b/Data/Agent/security.py index 443a0dd..30b935d 100644 --- a/Data/Agent/security.py +++ b/Data/Agent/security.py @@ -230,6 +230,7 @@ class AgentKeyStore: self._server_certificate_path = os.path.join(self.settings_dir, "server_certificate.pem") self._server_signing_key_path = os.path.join(self.settings_dir, "server_signing_key.pub") self._identity_lock_path = os.path.join(self.settings_dir, "identity.lock") + self._installer_cache_path = os.path.join(self.settings_dir, "installer_code.shared.json") # ------------------------------------------------------------------ # Identity management @@ -455,3 +456,107 @@ class AgentKeyStore: if isinstance(value, str) and value.strip(): return value.strip() return None + + # ------------------------------------------------------------------ + # Installer code sharing helpers + # ------------------------------------------------------------------ + def _load_installer_cache(self) -> dict: + if not os.path.isfile(self._installer_cache_path): + return {} + try: + with open(self._installer_cache_path, "r", encoding="utf-8") as fh: + data = json.load(fh) + if isinstance(data, dict): + return data + except Exception: + pass + return {} + + def _store_installer_cache(self, payload: dict) -> None: + try: + with open(self._installer_cache_path, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + _restrict_permissions(self._installer_cache_path) + except Exception: + pass + + def cache_installer_code(self, code: str, consumer: Optional[str] = None) -> None: + normalized = (code or "").strip() + if not normalized: + return + payload = self._load_installer_cache() + payload["code"] = normalized + consumers = set() + existing = payload.get("consumed") + if isinstance(existing, list): + consumers = {str(item).upper() for item in existing if isinstance(item, str)} + if consumer: + consumers.add(str(consumer).upper()) + payload["consumed"] = sorted(consumers) + payload["updated_at"] = int(time.time()) + self._store_installer_cache(payload) + + def load_cached_installer_code(self) -> Optional[str]: + payload = self._load_installer_cache() + code = payload.get("code") + if isinstance(code, str): + stripped = code.strip() + if stripped: + return stripped + return None + + def mark_installer_code_consumed(self, consumer: Optional[str] = None) -> None: + payload = self._load_installer_cache() + if not payload: + return + consumers = set() + existing = payload.get("consumed") + if isinstance(existing, list): + consumers = {str(item).upper() for item in existing if isinstance(item, str)} + if consumer: + consumers.add(str(consumer).upper()) + payload["consumed"] = sorted(consumers) + payload["updated_at"] = int(time.time()) + + code_present = isinstance(payload.get("code"), str) and payload["code"].strip() + should_clear = False + if not code_present: + should_clear = True + else: + required_consumers = {"SYSTEM", "CURRENTUSER"} + if required_consumers.issubset(consumers): + should_clear = True + else: + remaining = required_consumers - consumers + if not remaining: + should_clear = True + else: + exists_other = False + for other in remaining: + if other == "SYSTEM": + cfg_name = "agent_settings_SYSTEM.json" + elif other == "CURRENTUSER": + cfg_name = "agent_settings_CURRENTUSER.json" + else: + cfg_name = None + if not cfg_name: + continue + path = os.path.join(self.settings_dir, cfg_name) + if os.path.isfile(path): + exists_other = True + break + if not exists_other: + should_clear = True + + if should_clear: + payload.pop("code", None) + payload["consumed"] = [] + + if payload.get("code") or payload.get("consumed"): + self._store_installer_cache(payload) + else: + try: + if os.path.isfile(self._installer_cache_path): + os.remove(self._installer_cache_path) + except Exception: + pass diff --git a/tests/test_agent_installer_code.py b/tests/test_agent_installer_code.py new file mode 100644 index 0000000..ef6732f --- /dev/null +++ b/tests/test_agent_installer_code.py @@ -0,0 +1,58 @@ +import json +import sys + +import pytest + + +@pytest.fixture +def agent_module(tmp_path, monkeypatch): + settings_dir = tmp_path / "Agent" / "Borealis" / "Settings" + settings_dir.mkdir(parents=True) + system_config = settings_dir / "agent_settings_SYSTEM.json" + system_config.write_text(json.dumps({ + "config_file_watcher_interval": 2, + "agent_id": "", + "regions": {}, + "installer_code": "", + }, indent=2)) + current_config = settings_dir / "agent_settings_CURRENTUSER.json" + current_config.write_text(json.dumps({ + "config_file_watcher_interval": 2, + "agent_id": "", + "regions": {}, + "installer_code": "", + }, indent=2)) + + monkeypatch.setenv("BOREALIS_ROOT", str(tmp_path)) + monkeypatch.setenv("BOREALIS_AGENT_MODE", "system") + monkeypatch.setenv("BOREALIS_AGENT_CONFIG", "") + monkeypatch.setitem(sys.modules, "PyQt5", None) + monkeypatch.setitem(sys.modules, "qasync", None) + monkeypatch.setattr(sys, "argv", ["agent.py", "--system-service", "--config", "SYSTEM"], raising=False) + + agent = pytest.importorskip( + "Data.Agent.agent", reason="agent module requires optional dependencies" + ) + return agent, system_config + + +def test_shared_installer_code_cache_allows_system_reuse(agent_module, tmp_path): + agent, system_config = agent_module + + client = agent.AgentHttpClient() + shared_code = "SHARED-CODE-1234" + client.key_store.cache_installer_code(shared_code, consumer="CURRENTUSER") + + # System agent should discover the cached code even though its config is empty. + resolved = client._resolve_installer_code() + assert resolved == shared_code + + # Config should now persist the adopted code to avoid repeated lookups. + data = json.loads(system_config.read_text()) + assert data.get("installer_code") == shared_code + + # After enrollment completes, the cache should be cleared for future runs. + client._consume_installer_code() + assert client.key_store.load_cached_installer_code() is None + data = json.loads(system_config.read_text()) + assert data.get("installer_code") == ""