Contrib agent Support for Machine Target Executor (#106)

* MachineTargetExecutor support for contrib agent

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
     ```

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
 

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 =

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)

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()
 

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()

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

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