diff --git a/dreadnode_cli/agent/cli.py b/dreadnode_cli/agent/cli.py index 8a6b473..c9c9f91 100644 --- a/dreadnode_cli/agent/cli.py +++ b/dreadnode_cli/agent/cli.py @@ -97,7 +97,19 @@ def push( env = {env_var.split("=")[0]: env_var.split("=")[1] for env_var in env_vars or []} agent_config = AgentConfig.read(directory) - server_config = UserConfig.read().get_server_config() + user_config = UserConfig.read() + + if not user_config.active_profile_name: + raise Exception("No server profile is set, use [bold]dreadnode login[/] to authenticate") + + if not agent_config.is_linked_to_profile(user_config.active_profile_name): + linked_profiles = ", ".join(agent_config.linked_profiles) + plural = "s" if len(agent_config.linked_profiles) > 1 else "" + raise Exception( + f"This agent is linked to the [magenta]{linked_profiles}[/] server profile{plural}, but the current server profile is [yellow]{user_config.active_profile_name}[/], create the agent again." + ) + + server_config = user_config.get_server_config() registry = get_registry(server_config) @@ -124,7 +136,7 @@ def push( notes = notes or Prompt.ask("Notes?") agent = client.create_strike_agent(container, name, strike=agent_config.strike, notes=notes) - agent_config.add_link(agent.key, agent.id).write(directory) + agent_config.add_link(agent.key, agent.id, user_config.active_profile_name).write(directory) else: active_agent_id = agent_config.active if active_agent_id is None: @@ -134,7 +146,16 @@ def push( print(":robot: Creating a new version ...") notes = notes or Prompt.ask("Notes?") - agent = client.create_strike_agent_version(str(active_agent_id), container, notes) + try: + agent = client.create_strike_agent_version(str(active_agent_id), container, notes) + except Exception as e: + # 404 is expected if the agent was created on a different server profile + if str(e).startswith("404"): + raise Exception( + f"Agent '{active_agent_id}' not found for the current server profile, create the agent again." + ) from e + else: + raise e print(format_agent(agent)) diff --git a/dreadnode_cli/agent/config.py b/dreadnode_cli/agent/config.py index 76d83cd..aaf4875 100644 --- a/dreadnode_cli/agent/config.py +++ b/dreadnode_cli/agent/config.py @@ -8,6 +8,7 @@ class AgentLink(pydantic.BaseModel): + profile: str id: UUID runs: list[UUID] = [] @@ -42,12 +43,19 @@ def write(self, directory: pathlib.Path = pathlib.Path(".")) -> None: with (directory / AGENT_CONFIG_FILENAME).open("w") as f: YAML().dump(self.model_dump(mode="json"), f) - def add_link(self, key: str, id: UUID) -> "AgentConfig": + def add_link(self, key: str, id: UUID, profile: str) -> "AgentConfig": if key not in self.links: - self.links[key] = AgentLink(id=id) + self.links[key] = AgentLink(id=id, profile=profile) self.active = key return self + @property + def linked_profiles(self) -> list[str]: + return list({link.profile for link in self.links.values()}) + + def is_linked_to_profile(self, profile: str) -> bool: + return any(link.profile == profile for link in self.links.values()) + def add_run(self, id: UUID) -> "AgentConfig": self.active_link.runs.append(id) return self diff --git a/dreadnode_cli/agent/tests/test_config.py b/dreadnode_cli/agent/tests/test_config.py index f0c9277..5da9de3 100644 --- a/dreadnode_cli/agent/tests/test_config.py +++ b/dreadnode_cli/agent/tests/test_config.py @@ -21,11 +21,15 @@ def test_agent_config_add_link() -> None: config = AgentConfig(project_name="test") id = UUID("00000000-0000-0000-0000-000000000000") - config.add_link("test", id) + config.add_link("test", id, "test") assert config.active == "test" assert config.links["test"].id == id assert config.links["test"].runs == [] + assert config.links["test"].profile == "test" + assert config.linked_profiles == ["test"] + assert config.is_linked_to_profile("test") + assert not config.is_linked_to_profile("other") def test_agent_config_add_run() -> None: @@ -33,7 +37,7 @@ def test_agent_config_add_run() -> None: agent_id = UUID("00000000-0000-0000-0000-000000000000") run_id = UUID("11111111-1111-1111-1111-111111111111") - config.add_link("test", agent_id) + config.add_link("test", agent_id, "test") config.add_run(run_id) assert config.links["test"].runs == [run_id] @@ -44,7 +48,7 @@ def test_agent_config_write_read(tmp_path: Path) -> None: agent_id = UUID("00000000-0000-0000-0000-000000000000") run_id = UUID("11111111-1111-1111-1111-111111111111") - config.add_link("test", agent_id) + config.add_link("test", agent_id, "test") config.add_run(run_id) config.write(tmp_path) @@ -53,6 +57,7 @@ def test_agent_config_write_read(tmp_path: Path) -> None: assert loaded.active == "test" assert loaded.links["test"].id == agent_id assert loaded.links["test"].runs == [run_id] + assert loaded.links["test"].profile == "test" def test_agent_config_update_active() -> None: @@ -61,11 +66,11 @@ def test_agent_config_update_active() -> None: id2 = UUID("11111111-1111-1111-1111-111111111111") # Add first link - config.add_link("test1", id1) + config.add_link("test1", id1, "test1") assert config.active == "test1" # Add second link - config.add_link("test2", id2) + config.add_link("test2", id2, "test2") assert config.active == "test2" # Remove active link diff --git a/dreadnode_cli/config.py b/dreadnode_cli/config.py index 2434de4..b5c5f5f 100644 --- a/dreadnode_cli/config.py +++ b/dreadnode_cli/config.py @@ -50,6 +50,12 @@ def write(self) -> None: with USER_CONFIG_PATH.open("w") as f: YAML().dump(self.model_dump(mode="json"), f) + @property + def active_profile_name(self) -> str | None: + """Get the name of the active profile.""" + self._update_active() + return self.active + def get_server_config(self, profile: str | None = None) -> ServerConfig: """Get the server configuration for the given profile or None if not set."""