From c3651caed02146b6b9e044f94726cee063417340 Mon Sep 17 00:00:00 2001 From: "Alfin S. Thomas" Date: Thu, 13 Jan 2022 17:12:57 +0530 Subject: [PATCH] Contrib agent Support for Machine Target Executor (#106) * MachineTargetExecutor support for contrib agent --- docs/agents/contrib.md | 18 ++++--- docs/releases.md | 5 ++ setup.cfg | 2 +- src/ychaos/agents/contrib.py | 10 ++-- .../core/executor/MachineTargetExecutor.py | 51 ++++++++++++++----- .../executor/test_MachineTargetExecutor.py | 21 ++++++++ .../resources/contrib_agent/awesome_agent.py | 40 +++++++++++++++ tests/resources/contrib_agent/testplan.yaml | 14 +++++ 8 files changed, 135 insertions(+), 26 deletions(-) create mode 100644 tests/resources/contrib_agent/awesome_agent.py create mode 100644 tests/resources/contrib_agent/testplan.yaml diff --git a/docs/agents/contrib.md b/docs/agents/contrib.md index 3ad7de09..a24b545b 100644 --- a/docs/agents/contrib.md +++ b/docs/agents/contrib.md @@ -17,10 +17,9 @@ with a sample structure given below. from ychaos.agents.agent import Agent, AgentConfig from ychaos.agents.utils.annotations import log_agent_lifecycle from queue import LifoQueue -from pydantic import BaseModel -class MyAwesomeAgentConfig(BaseModel): +class MyAwesomeAgentConfig(AgentConfig): name: str = "my_awesome_agent" description: str = "This is an invincible agent" @@ -79,10 +78,12 @@ schema for more details. "agents": [ { "type": "contrib", - "path": "/tmp/awesome_agent.py", "config": { - "key1": "value1", - "key2": "value2" + "path": "/tmp/awesome_agent.py", + "contrib_agent_config": { + "key1": "value1", + "key2": "value2" + } } } ] @@ -98,8 +99,9 @@ schema for more details. target_type: self # Your preferred target type agents: - type: contrib - path: "/tmp/awesome_agent.py" config: - key1: value1 - key2: value2 + path: "/tmp/awesome_agent.py" + contrib_agent_config: + key1: value1 + key2: value2 ``` diff --git a/docs/releases.md b/docs/releases.md index 26cecb56..f13699f7 100644 --- a/docs/releases.md +++ b/docs/releases.md @@ -2,10 +2,15 @@ ## Version 0.x.x +### Version 0.5.0 +1. Add `MachineTargetExecutor` support for contrib agents by [Alfin S Thomas](https://github.com/AlfinST) + ### Version 0.4.0 1. Add `SelfTargetExecutor`, that runs the agents on the same machine from where YChaos is triggered by [Alfin S Thomas](https://github.com/AlfinST) +2. Publish docker image with ychaos pre-installed by [Vijay Babu](https://github.com/vijaybabu4589) +3. Add support to SSH common args in TestPlan, make SSH config optional by [Shashank Sharma](https://github.com/shashankrnr32) ### Version 0.3.0 diff --git a/setup.cfg b/setup.cfg index 99cfb936..0e481ce6 100644 --- a/setup.cfg +++ b/setup.cfg @@ -25,7 +25,7 @@ project_urls = CI Pipeline = https://cd.screwdriver.cd/pipelines/7419 Download = https://pypi.org/project/ychaos/#files url = https://github.com/yahoo/ychaos -version = 0.4.0 +version = 0.5.0 [options] namespace_packages = diff --git a/src/ychaos/agents/contrib.py b/src/ychaos/agents/contrib.py index 22c6e96e..c0febaf3 100644 --- a/src/ychaos/agents/contrib.py +++ b/src/ychaos/agents/contrib.py @@ -2,7 +2,7 @@ # Licensed under the terms of the Apache 2.0 license. See the LICENSE file in the project root for terms import importlib.util from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Optional from pydantic import Field, PrivateAttr @@ -21,7 +21,7 @@ class ContribAgentConfig(AgentConfig): description="The class name of the contributed Agent config", ) - config: Dict[Any, Any] = Field( + contrib_agent_config: Optional[Dict[Any, Any]] = Field( default=dict(), description="The configuration that will be passed to the community agent", ) @@ -35,7 +35,9 @@ def __init__(self, **kwargs): self._import_module() # Validate that `config` adheres to the schema of the contrib agent - self.config = self.get_agent_config_class()(**self.config) + self.contrib_agent_config = self.get_agent_config_class()( + **self.contrib_agent_config + ) def _import_module(self): specification = importlib.util.spec_from_file_location( @@ -58,4 +60,4 @@ def get_agent_config_class(self) -> Any: return agent_config_klass def get_agent(self): - return self.get_agent_class()(self.config) + return self.get_agent_class()(self.contrib_agent_config) diff --git a/src/ychaos/core/executor/MachineTargetExecutor.py b/src/ychaos/core/executor/MachineTargetExecutor.py index 8e0f8dd8..1a744e40 100644 --- a/src/ychaos/core/executor/MachineTargetExecutor.py +++ b/src/ychaos/core/executor/MachineTargetExecutor.py @@ -3,8 +3,9 @@ import json import random from types import SimpleNamespace - +from pathlib import Path from ...app_logger import AppLogger +from ...agents.index import AgentType from ...testplan.attack import MachineTargetDefinition from ...testplan.schema import TestPlan from ...utils.dependency import DependencyUtils @@ -152,7 +153,6 @@ def prepare(self): hosts = ",".join(self.testplan.attack.get_target_config().get_effective_hosts()) if len(self.testplan.attack.get_target_config().get_effective_hosts()) == 1: hosts += "," - self.ansible_context.inventory = InventoryManager( loader=self.ansible_context.loader, sources=hosts ) @@ -241,17 +241,7 @@ def prepare(self): mode="0755", ), ), - dict( - name="Copy testplan from local to remote", - register="result_testplan_file", - action=dict( - module="copy", - content=json.dumps( - self.testplan.to_serialized_dict(), indent=4 - ), - dest="{{result_create_workspace.path}}/testplan.json", - ), - ), + *self.get_file_transfer_tasks(), dict( name="Run YChaos Agent", ignore_errors="yes", @@ -306,6 +296,41 @@ def prepare(self): ], ) + def get_file_transfer_tasks(self): + task_list = list() + testplan = self.testplan.copy() + for i, agent in enumerate(self.testplan.attack.agents): + if agent.type == AgentType.CONTRIB: + filename = Path(testplan.attack.agents[i].config["path"]) + task = dict( + name=f"Copy {filename.name} to remote", + register="copy_contrib_agent_" + filename.stem, + action=dict( + module="copy", + src=str(filename.absolute()), + dest="{{result_create_workspace.path}}/" + filename.name, + ), + ) + task_list.append(task) + testplan.attack.agents[i].config["path"] = "./ychaos_ws/{}".format( + filename.name + ) + + # testplan will not have any changes from original if there are no contrib agents present + testplan_task = [ + dict( + name="Copy testplan from local to remote", + register="result_testplan_file", + action=dict( + module="copy", + content=json.dumps(testplan.to_serialized_dict(), indent=4), + dest="{{result_create_workspace.path}}/testplan.json", + ), + ) + ] + + return testplan_task + task_list + def execute(self) -> None: self.prepare() diff --git a/tests/core/executor/test_MachineTargetExecutor.py b/tests/core/executor/test_MachineTargetExecutor.py index 6071dd35..164e1467 100644 --- a/tests/core/executor/test_MachineTargetExecutor.py +++ b/tests/core/executor/test_MachineTargetExecutor.py @@ -1,6 +1,7 @@ # Copyright 2021, Yahoo # Licensed under the terms of the Apache 2.0 license. See the LICENSE file in the project root for terms import json +import yaml from pathlib import Path from unittest import TestCase @@ -289,5 +290,25 @@ def __call__(self, *args, **kwargs): executor.execute() self.assertTrue(mock_hook_target_unreachable.test_value) + def test_machine_executor_when_agent_type_is_contrib(self): + with open( + Path(__file__) + .joinpath("../../../resources/contrib_agent/testplan.yaml") + .resolve() + ) as f: + tp = yaml.safe_load(f) + tp["attack"]["agents"][0]["config"]["path"] = ( + Path(__file__) + .joinpath("../../../resources/contrib_agent/awesome_agent.py") + .resolve() + ) + mock_valid_testplan = TestPlan(**tp) + executor = MachineTargetExecutor(mock_valid_testplan) + executor.prepare() + playbook_tasks = list( + map(lambda x: x["name"], executor.ansible_context.play_source["tasks"]) + ) + self.assertIn("Copy awesome_agent.py to remote", playbook_tasks) + def tearDown(self) -> None: unstub() diff --git a/tests/resources/contrib_agent/awesome_agent.py b/tests/resources/contrib_agent/awesome_agent.py new file mode 100644 index 00000000..aad9d675 --- /dev/null +++ b/tests/resources/contrib_agent/awesome_agent.py @@ -0,0 +1,40 @@ +from ychaos.agents.agent import Agent, AgentConfig +from ychaos.agents.utils.annotations import log_agent_lifecycle + + +class AwesomeAgentConfig(AgentConfig): + name: str = "awesomeagent" + description: str = "This is an dummy agent" + + +class MyAwesomeAgent(Agent): + @log_agent_lifecycle + def __init__(self, config) -> None: + print("in method: init") + assert isinstance(config, AgentConfig) + super(MyAwesomeAgent, self).__init__(config) + + @log_agent_lifecycle + def monitor(self) -> None: + print("in method: monitor") + super(MyAwesomeAgent, self).monitor() + return self._status + + @log_agent_lifecycle + def setup(self) -> None: + print("in method: setup") + super(MyAwesomeAgent, self).setup() + + @log_agent_lifecycle + def run(self) -> None: + print("in method: run") + super(MyAwesomeAgent, self).run() + + @log_agent_lifecycle + def teardown(self) -> None: + print("in method: teardown") + super(MyAwesomeAgent, self).teardown() + + +AgentClass = MyAwesomeAgent +AgentConfigClass = AwesomeAgentConfig diff --git a/tests/resources/contrib_agent/testplan.yaml b/tests/resources/contrib_agent/testplan.yaml new file mode 100644 index 00000000..4b8074fa --- /dev/null +++ b/tests/resources/contrib_agent/testplan.yaml @@ -0,0 +1,14 @@ +attack: + target_type: machine + target_config: + blast_radius: 100 + hostnames: + - mockhost.resilience.yahoo.cloud + ssh_config: {} + agents: + - type: contrib + config: + path: path/to/awesome_agent.py + contrib_agent_config: + key1 : value1 + key2 : value2