diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 461e2018..a1bd0c96 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -9,6 +9,7 @@ variables: FARADAY_PASSWORD: custom_pass FARADAY_EMAIL: test@test.com FARADAY_REF: white/dev + EXECUTOR_DIR: ./basic_executor.py cache: paths: @@ -37,15 +38,17 @@ unit_tests: - pip install pytest pytest-cov pytest-aiohttp - python setup.py install - mkdir run_from - - cd run_from && pytest ../tests/unittests --capture=sys -v --cov=../faraday_agent_dispatcher --color=yes --disable-warnings + - cd run_from && mkdir logs && pytest ../tests/unittests --capture=sys -v --cov=../faraday_agent_dispatcher --color=yes --disable-warnings artifacts: - when: on_failure - paths: - - dist/* + when: on_failure + paths: + - dist/* + - run_from/logs/* + expire_in: 7 days except: variables: - $INTEGRATION - - $CI_COMMIT_REF_NAME =~ /^.*ci.*$/ + - $CI_COMMIT_REF_NAME =~ /^.*ci-int.*$/ integration_faraday: stage: post_testing diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 5778337f..a021dffd 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -15,7 +15,7 @@ Types of Contributions Report Bugs ~~~~~~~~~~~ -Report bugs at https://github.com/EricHorvat/dummy_faraday_agent/issues. +Report bugs at https://github.com/faradaysec/faraday_agent_dispatcher/issues. If you are reporting a bug, please include: @@ -38,14 +38,14 @@ and "help wanted" is open to whoever wants to implement it. Write Documentation ~~~~~~~~~~~~~~~~~~~ -dummy_faraday_agent could always use more documentation, whether as part of the -official dummy_faraday_agent docs, in docstrings, or even on the web in blog posts, +faraday_agent_dispatcher could always use more documentation, whether as part of the +official faraday_agent_dispatcher docs, in docstrings, or even on the web in blog posts, articles, and such. Submit Feedback ~~~~~~~~~~~~~~~ -The best way to send feedback is to file an issue at https://github.com/EricHorvat/dummy_faraday_agent/issues. +The best way to send feedback is to file an issue at https://github.com/faraday/faraday_agent_dispatcher/issues. If you are proposing a feature: @@ -57,17 +57,17 @@ If you are proposing a feature: Get Started! ------------ -Ready to contribute? Here's how to set up `dummy_faraday_agent` for local development. +Ready to contribute? Here's how to set up `faraday_agent_dispatcher` for local development. -1. Fork the `dummy_faraday_agent` repo on GitHub. +1. Fork the `faraday_agent_dispatcher` repo on GitHub. 2. Clone your fork locally:: - $ git clone git@github.com:your_name_here/dummy_faraday_agent.git + $ git clone git@github.com:your_name_here/faraday_agent_dispatcher.git 3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: - $ mkvirtualenv dummy_faraday_agent - $ cd dummy_faraday_agent/ + $ mkvirtualenv faraday_agent_dispatcher + $ cd faraday_agent_dispatcher/ $ python setup.py develop 4. Create a branch for local development:: @@ -79,7 +79,7 @@ Ready to contribute? Here's how to set up `dummy_faraday_agent` for local develo 5. When you're done making changes, check that your changes pass flake8 and the tests, including testing other Python versions with tox:: - $ flake8 dummy_faraday_agent tests + $ flake8 faraday_agent_dispatcher tests $ python setup.py test or py.test $ tox @@ -103,7 +103,7 @@ Before you submit a pull request, check that it meets these guidelines: your new functionality into a function with a docstring, and add the feature to the list in README.rst. 3. The pull request should work for Python 2.7, 3.4, 3.5 and 3.6, and for PyPy. Check - https://travis-ci.org/EricHorvat/dummy_faraday_agent/pull_requests + https://travis-ci.org/faradaysec/faraday_agent_dispatcher/pull_requests and make sure that the tests pass for all supported Python versions. Tips @@ -111,7 +111,7 @@ Tips To run a subset of tests:: -$ py.test tests.test_dummy_faraday_agent +$ py.test tests.faraday_agent_dispatcher Deploying diff --git a/Makefile b/Makefile index ea46ec5c..1399b142 100644 --- a/Makefile +++ b/Makefile @@ -51,7 +51,7 @@ clean-test: ## remove test and coverage artifacts rm -fr .pytest_cache lint: ## check style with flake8 - flake8 dummy_faraday_agent tests + flake8 faraday_agent_dispatcher tests test: ## run tests quickly with the default Python py.test @@ -60,15 +60,15 @@ test-all: ## run tests on every Python version with tox tox coverage: ## check code coverage quickly with the default Python - coverage run --source dummy_faraday_agent -m pytest + coverage run --source faraday_agent_dispatcher -m pytest coverage report -m coverage html $(BROWSER) htmlcov/index.html docs: ## generate Sphinx HTML documentation, including API docs - rm -f docs/dummy_faraday_agent.rst + rm -f docs/faraday_agent_dispatcher.rst rm -f docs/modules.rst - sphinx-apidoc -o docs/ dummy_faraday_agent + sphinx-apidoc -o docs/ faraday_agent_dispatcher $(MAKE) -C docs clean $(MAKE) -C docs html $(BROWSER) docs/_build/html/index.html diff --git a/contrib/nmap.py b/contrib/nmap.py index a86d8ae8..3420a164 100755 --- a/contrib/nmap.py +++ b/contrib/nmap.py @@ -1,4 +1,5 @@ #!/usr/bin/env python +import os import subprocess """Yo need to clone and install faraday plugins""" @@ -6,12 +7,13 @@ cmd = [ "nmap", - "-p80,443", - "190.210.92.77", + "-p{}".format(os.environ.get('EXECUTOR_CONFIG_PORT_LIST')), + os.environ.get('EXECUTOR_CONFIG_TARGET'), "-oX", "-", ] + results = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) nmap = NmapPlugin() nmap.parseOutputString(results.stdout) diff --git a/docs/Makefile b/docs/Makefile index e85ac610..e9afa1f1 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -4,7 +4,7 @@ # You can set these variables from the command line. SPHINXOPTS = SPHINXBUILD = python -msphinx -SPHINXPROJ = dummy_faraday_agent +SPHINXPROJ = faraday_agent_dispatcher SOURCEDIR = . BUILDDIR = _build diff --git a/docs/conf.py b/docs/conf.py index 8a74d2ec..c363ce0a 100755 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,7 +1,7 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- # -# faraday_dummy_agent documentation build configuration file, created by +# faraday_agent_dispatcher documentation build configuration file, created by # sphinx-quickstart on Fri Jun 9 13:47:02 2017. # # This file is execfile()d with the current directory set to its @@ -47,7 +47,7 @@ master_doc = 'index' # General information about the project. -project = u'faraday_dummy_agent' +project = u'faraday_agent_dispatcher' copyright = u"2019, Eric Horvat" author = u"Eric Horvat" @@ -101,7 +101,7 @@ # -- Options for HTMLHelp output --------------------------------------- # Output file base name for HTML help builder. -htmlhelp_basename = 'dummy_faraday_agentdoc' +htmlhelp_basename = 'faraday_agent_dispatcher' # -- Options for LaTeX output ------------------------------------------ @@ -128,8 +128,8 @@ # (source start file, target name, title, author, documentclass # [howto, manual, or own class]). latex_documents = [ - (master_doc, 'faraday_dummy_agent.tex', - u'faraday_dummy_agent Documentation', + (master_doc, 'faraday_agent_dispatcher.tex', + u'faraday_agent_dispatcher Documentation', u'Eric Horvat', 'manual'), ] @@ -139,8 +139,8 @@ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). man_pages = [ - (master_doc, 'faraday_dummy_agent', - u'faraday_dummy_agent Documentation', + (master_doc, 'faraday_agent_dispatcher', + u'faraday_agent_dispatcher Documentation', [author], 1) ] @@ -151,10 +151,10 @@ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'faraday_dummy_agent', - u'faraday_dummy_agent Documentation', + (master_doc, 'faraday_agent_dispatcher', + u'faraday_agent_dispatcher Documentation', author, - 'faraday_dummy_agent', + 'faraday_agent_dispatcher', 'One line description of project.', 'Miscellaneous'), ] diff --git a/docs/index.rst b/docs/index.rst index 4d349063..abceba8c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,4 +1,4 @@ -Welcome to dummy_faraday_agent's documentation! +Welcome to faraday_agent_dispatcher's documentation! ====================================== .. toctree:: diff --git a/docs/installation.rst b/docs/installation.rst index c1cc67e0..c9ed0526 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -8,13 +8,13 @@ Installation Stable release -------------- -To install dummy_faraday_agent, run this command in your terminal: +To install faraday_agent_dispatcher, run this command in your terminal: .. code-block:: console - $ pip install dummy_faraday_agent + $ pip install faraday_agent_dispatcher -This is the preferred method to install dummy_faraday_agent, as it will always install the most recent stable release. +This is the preferred method to install faraday_agent_dispatcher, as it will always install the most recent stable release. If you don't have `pip`_ installed, this `Python installation guide`_ can guide you through the process. @@ -26,19 +26,19 @@ you through the process. From sources ------------ -The sources for dummy_faraday_agent can be downloaded from the `Github repo`_. +The sources for faraday_agent_dispatcher can be downloaded from the `Github repo`_. You can either clone the public repository: .. code-block:: console - $ git clone git://github.com/EricHorvat/dummy_faraday_agent + $ git clone git://github.com/faradaysec/faraday_agent_dispatcher Or download the `tarball`_: .. code-block:: console - $ curl -OL https://github.com/EricHorvat/dummy_faraday_agent/tarball/master + $ curl -OL https://github.com/faradaysec/faraday_agent_dispatcher/tarball/master Once you have a copy of the source, you can install it with: @@ -47,5 +47,5 @@ Once you have a copy of the source, you can install it with: $ python setup.py install -.. _Github repo: https://github.com/EricHorvat/dummy_faraday_agent -.. _tarball: https://github.com/EricHorvat/dummy_faraday_agent/tarball/master +.. _Github repo: https://github.com/faradaysec/faraday_agent_dispatcher +.. _tarball: https://github.com/faradaysec/faraday_agent_dispatcher/tarball/master diff --git a/docs/make.bat b/docs/make.bat index 66e06c11..9c122b2c 100644 --- a/docs/make.bat +++ b/docs/make.bat @@ -9,7 +9,7 @@ if "%SPHINXBUILD%" == "" ( ) set SOURCEDIR=. set BUILDDIR=_build -set SPHINXPROJ=dummy_faraday_agent +set SPHINXPROJ=faraday_agent_dispatcher if "%1" == "" goto help diff --git a/docs/usage.rst b/docs/usage.rst index f79ec90e..dd67e1b5 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -2,6 +2,7 @@ Usage ===== -To use dummy_faraday_agent in a project:: +To use faraday_agent_dispatcher in a project:: + + import faraday_agent_dispatcher - import dummy_faraday_agent diff --git a/faraday_agent_dispatcher/__init__.py b/faraday_agent_dispatcher/__init__.py index 6ab7ab62..50a19dab 100644 --- a/faraday_agent_dispatcher/__init__.py +++ b/faraday_agent_dispatcher/__init__.py @@ -16,8 +16,8 @@ # along with this program. If not, see . -"""Top-level package for faraday_dummy_agent.""" +"""Top-level package for faraday_agent_dispatcher.""" __author__ = """Eric Horvat""" __email__ = 'erich@infobytesec.com' -__version__ = '0.1.0' +__version__ = '1.0' diff --git a/faraday_agent_dispatcher/cli.py b/faraday_agent_dispatcher/cli.py index a5e65e85..3fbad141 100644 --- a/faraday_agent_dispatcher/cli.py +++ b/faraday_agent_dispatcher/cli.py @@ -15,7 +15,7 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -"""Console script for faraday_dummy_agent.""" +"""Console script for faraday_agent_dispatcher.""" import os import sys import shutil diff --git a/faraday_agent_dispatcher/config.py b/faraday_agent_dispatcher/config.py index 83585868..de846897 100644 --- a/faraday_agent_dispatcher/config.py +++ b/faraday_agent_dispatcher/config.py @@ -17,6 +17,7 @@ import logging import configparser from pathlib import Path +from configparser import DuplicateSectionError try: FARADAY_PATH = Path(os.environ['FARADAY_HOME']) @@ -39,8 +40,11 @@ def reset_config(filepath): instance.clear() - if not instance.read(filepath): - raise ValueError(f'Unable to read config file located at {filepath}') + try: + if not instance.read(filepath): + raise ValueError(f'Unable to read config file located at {filepath}') + except DuplicateSectionError as e: + raise ValueError(f'The config in {filepath} contains duplicated sections') def check_filepath(filepath: str = None): @@ -56,6 +60,10 @@ def save_config(filepath=None): instance.write(configfile) -TOKENS_SECTION = "tokens" -SERVER_SECTION = "server" -EXECUTOR_SECTION = "executor" +class Sections: + TOKENS = "tokens" + SERVER = "server" + AGENT = "agent" + EXECUTOR_VARENVS = "{}_varenvs" + EXECUTOR_PARAMS = "{}_params" + EXECUTOR_DATA = "{}" diff --git a/faraday_agent_dispatcher/dispatcher.py b/faraday_agent_dispatcher/dispatcher.py index 707b2e9d..027bc19c 100644 --- a/faraday_agent_dispatcher/dispatcher.py +++ b/faraday_agent_dispatcher/dispatcher.py @@ -15,6 +15,7 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . +import os import json import asyncio @@ -29,49 +30,54 @@ control_str, control_host, control_registration_token, - control_agent_token + control_agent_token, + control_list ) import faraday_agent_dispatcher.logger as logging -from faraday_agent_dispatcher.config import instance as config, \ - EXECUTOR_SECTION, SERVER_SECTION, TOKENS_SECTION, save_config +from faraday_agent_dispatcher.config import instance as config, Sections, save_config +from faraday_agent_dispatcher.executor import Executor logger = logging.get_logger() +logging.setup_logging() class Dispatcher: __control_dict = { - SERVER_SECTION: { + Sections.SERVER: { "host": control_host, - "api_port": control_int, - "websocket_port": control_int, + "api_port": control_int(), + "websocket_port": control_int(), "workspace": control_str }, - TOKENS_SECTION: { + Sections.TOKENS: { "registration": control_registration_token, "agent": control_agent_token }, - EXECUTOR_SECTION: { - "cmd": control_str, - "agent_name": control_str - } + Sections.AGENT: { + "agent_name": control_str, + "executors": control_list(can_repeat=False) + }, } def __init__(self, session, config_path=None): reset_config(filepath=config_path) self.control_config() self.config_path = config_path - self.host = config.get(SERVER_SECTION, "host") - self.api_port = config.get(SERVER_SECTION, "api_port") - self.websocket_port = config.get(SERVER_SECTION, "websocket_port") - self.workspace = config.get(SERVER_SECTION, "workspace") - self.agent_token = config[TOKENS_SECTION].get("agent", None) - self.executor_cmd = config.get(EXECUTOR_SECTION, "cmd") - self.agent_name = config.get(EXECUTOR_SECTION, "agent_name") + self.host = config.get(Sections.SERVER, "host") + self.api_port = config.get(Sections.SERVER, "api_port") + self.websocket_port = config.get(Sections.SERVER, "websocket_port") + self.workspace = config.get(Sections.SERVER, "workspace") + self.agent_token = config[Sections.TOKENS].get("agent", None) + self.agent_name = config.get(Sections.AGENT, "agent_name") self.session = session self.websocket = None self.websocket_token = None + self.executors = { + executor_name: + Executor(executor_name, config) for executor_name in config[Sections.AGENT].get("executors", []).split(",") + } async def reset_websocket_token(self): # I'm built so I ask for websocket token @@ -86,7 +92,7 @@ async def reset_websocket_token(self): async def register(self): if self.agent_token is None: - registration_token = self.agent_token = config.get(TOKENS_SECTION, "registration") + registration_token = self.agent_token = config.get(Sections.TOKENS, "registration") assert registration_token is not None, "The registration token is mandatory" token_registration_url = api_url(self.host, self.api_port, @@ -98,7 +104,7 @@ async def register(self): assert token_response.status == 201 token = await token_response.json() self.agent_token = token["token"] - config.set(TOKENS_SECTION, "agent", self.agent_token) + config.set(Sections.TOKENS, "agent", self.agent_token) save_config(self.config_path) except ClientResponseError as e: if e.status == 404: @@ -110,22 +116,30 @@ async def register(self): self.websocket_token = await self.reset_websocket_token() - async def connect(self): + async def connect(self, out_func=None): - if not self.websocket_token: + if not self.websocket_token and not out_func: return - async with websockets.connect(websocket_url(self.host, self.websocket_port)) as websocket: - await websocket.send(json.dumps({ - 'action': 'JOIN_AGENT', - 'workspace': self.workspace, - 'token': self.websocket_token, - })) + connected_data = json.dumps({ + 'action': 'JOIN_AGENT', + 'workspace': self.workspace, + 'token': self.websocket_token, + 'executors': [{"executor_name": executor.name, "args": executor.params} + for executor in self.executors.values()] + }) + + if out_func is None: + + async with websockets.connect(websocket_url(self.host, self.websocket_port)) as websocket: + await websocket.send(connected_data) - logger.info("Connection to Faraday server succeeded") - self.websocket = websocket + logger.info("Connection to Faraday server succeeded") + self.websocket = websocket - await self.run_await() # This line can we called from outside (in main) + await self.run_await() # This line can we called from outside (in main) + else: + await out_func(connected_data) async def run_await(self): while True: @@ -133,37 +147,146 @@ async def run_await(self): data = await self.websocket.recv() asyncio.create_task(self.run_once(data)) - async def run_once(self, data:str= None): - # TODO Control data - logger.info('Running executor with data: %s', data) + async def run_once(self, data:str= None, out_func=None): + out_func = out_func if out_func is not None else self.websocket.send + logger.info('Parsing data: %s', data) data_dict = json.loads(data) - if "action" in data_dict: - if data_dict["action"] == "RUN": - process = await self.create_process() + if "action" not in data_dict: + logger.info("Data not contains action to do") + await out_func(json.dumps({"error": "'action' key is mandatory in this websocket connection"})) + return + + if data_dict["action"] not in ["RUN"]: # ONLY SUPPORTED COMMAND FOR NOW + logger.info("Unrecognized action") + await out_func(json.dumps({f"{data_dict['action']}_RESPONSE": "Error: Unrecognized action"})) + return + + if data_dict["action"] == "RUN": + if "executor" not in data_dict: + logger.error("No executor selected") + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "running": False, + "message": f"No executor selected to {self.agent_name} agent" + }) + ) + return + + if data_dict["executor"] not in self.executors: + logger.error("The selected executor not exists") + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "executor_name": data_dict['executor'], + "running": False, + "message": f"The selected executor {data_dict['executor']} not exists in {self.agent_name} " + f"agent" + }) + ) + return + + executor = self.executors[data_dict["executor"]] + + params = list(executor.params.keys()).copy() + passed_params = data_dict['args'] if 'args' in data_dict else {} + [params.remove(param) for param in config.defaults()] + + all_accepted = all( + [ + any([ + param in passed_param # Control any available param + for param in params # was passed + ]) + for passed_param in passed_params # For all passed params + ]) + if not all_accepted: + logger.error("Unexpected argument passed to {} executor".format(executor.name)) + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "executor_name": executor.name, + "running": False, + "message": f"Unexpected argument(s) passed to {executor.name} executor from {self.agent_name} " + f"agent" + }) + ) + mandatory_full = all( + [ + not executor.params[param] # All params is not mandatory + or any([ + param in passed_param for passed_param in passed_params # Or was passed + ]) + for param in params + ] + ) + if not mandatory_full: + logger.error("Mandatory argument not passed to {} executor".format(executor.name)) + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "executor_name": executor.name, + "running": False, + "message": f"Mandatory argument(s) not passed to {executor.name} executor from " + f"{self.agent_name} agent" + }) + ) + + if mandatory_full and all_accepted: + running_msg = f"Running {executor.name} executor from {self.agent_name} agent" + logger.info("Running {} executor".format(executor.name)) + + process = await self.create_process(executor, passed_params) tasks = [StdOutLineProcessor(process, self.session).process_f(), StdErrLineProcessor(process).process_f(), ] - + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "executor_name": executor.name, + "running": True, + "message": running_msg + }) + ) await asyncio.gather(*tasks) await process.communicate() assert process.returncode is not None if process.returncode == 0: - logger.info("Executor finished successfully") + logger.info("Executor {} finished successfully".format(executor.name)) + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "executor_name": executor.name, + "successful": True, + "message": f"Executor {executor.name} from {self.agent_name} finished successfully" + })) else: logger.warning( - f"Executor finished with exit code {process.returncode}") - else: - logger.info("Action unrecognized") + f"Executor {executor.name} finished with exit code {process.returncode}") + await out_func( + json.dumps({ + "action": "RUN_STATUS", + "executor_name": executor.name, + "successful": False, + "message": f"Executor {executor.name} from {self.agent_name} failed" + })) + async def create_process(self, executor: Executor, args): + env = os.environ.copy() + if isinstance(args, dict): + for k in args: + env[f"EXECUTOR_CONFIG_{k.upper()}"] = str(args[k]) else: - logger.info("Data not contains action to do") - - async def create_process(self): + logger.error("Args from data received has a not supported type") + raise ValueError("Args from data received has a not supported type") + for varenv, value in executor.varenvs.items(): + env[f"{varenv.upper()}"] = value process = await asyncio.create_subprocess_shell( - self.executor_cmd, + executor.cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - limit=int(config[EXECUTOR_SECTION].get("max_size", 64 * 1024)) + env=env, + limit=executor.max_size # If the config is not set, use async.io default ) return process @@ -171,5 +294,9 @@ async def create_process(self): def control_config(self): for section in self.__control_dict: for option in self.__control_dict[section]: + if section not in config: + err = f"Section {section} is an mandatory section in the config" # TODO "run config cmd" + logger.error(err) + raise ValueError(err) value = config.get(section, option) if option in config[section] else None self.__control_dict[section][option](option, value) diff --git a/faraday_agent_dispatcher/example_config.ini b/faraday_agent_dispatcher/example_config.ini index 795bface..04682b3a 100644 --- a/faraday_agent_dispatcher/example_config.ini +++ b/faraday_agent_dispatcher/example_config.ini @@ -4,14 +4,22 @@ host = localhost api_port = 5985 websocket_port = 9000 -[executor] -; Complete the cmd option with the command you want the dispatcher to run -; cmd = +[agent] agent_name = unnamed_agent -max_size = 65536 -; 1024 * 64 +; Complete the executor option with a comma separated list of executor names +executors = ex1 [tokens] ; To get your registration token, visit http://localhost:5985/#/admin/agents, copy ; the value and uncomment the line ; registration = + +[ex1] +; Complete the cmd option with the command you want the dispatcher to run +; cmd = +max_size = 65536 +; 1024 * 64 + +[ex1_varenvs] + +[ex1_params] \ No newline at end of file diff --git a/faraday_agent_dispatcher/executor.py b/faraday_agent_dispatcher/executor.py new file mode 100644 index 00000000..77386ea2 --- /dev/null +++ b/faraday_agent_dispatcher/executor.py @@ -0,0 +1,44 @@ +from faraday_agent_dispatcher.config import Sections +from faraday_agent_dispatcher.utils.control_values_utils import ( + control_int, + control_str, + control_bool +) + + +class Executor: + __control_dict = { + Sections.EXECUTOR_DATA: { + "cmd": control_str, + "max_size": control_int(True) + } + } + + def __init__(self, name: str, config): + name = name.strip() + self.control_config(name, config) + self.name = name + executor_section = Sections.EXECUTOR_DATA.format(name) + params_section = Sections.EXECUTOR_PARAMS.format(name) + varenvs_section = Sections.EXECUTOR_VARENVS.format(name) + self.cmd = config.get(executor_section, "cmd") + self.max_size = int(config[executor_section].get("max_size", 64 * 1024)) + self.params = dict(config[params_section]) if params_section in config else {} + self.params = {key: value.lower() in ["t", "true"] for key, value in self.params.items()} + self.varenvs = dict(config[varenvs_section]) if varenvs_section in config else {} + + def control_config(self, name, config): + if " " in name: + raise ValueError(f"Executor names can't contains space character, passed name: {name}") + if Sections.EXECUTOR_DATA.format(name) not in config: + raise ValueError(f"{name} is an executor name but there is no proper section") + + for section in self.__control_dict: + for option in self.__control_dict[section]: + value = config.get(section.format(name), option) if option in config[section.format(name)] else None + self.__control_dict[section][option](option, value) + params_section = Sections.EXECUTOR_PARAMS.format(name) + if params_section in config: + for option in config[params_section]: + value = config.get(params_section, option) + control_bool(option, value) \ No newline at end of file diff --git a/faraday_agent_dispatcher/logger.py b/faraday_agent_dispatcher/logger.py index 6ee35566..19da14a7 100644 --- a/faraday_agent_dispatcher/logger.py +++ b/faraday_agent_dispatcher/logger.py @@ -52,6 +52,7 @@ def setup_console_logging(formatter): console_handler = logging.StreamHandler() console_handler.setFormatter(formatter) console_handler.setLevel(config.LOGGING_LEVEL) + console_handler.name = "CONSOLE_HANDLER" add_handler(console_handler) LVL_SETTABLE_HANDLERS.append(console_handler) @@ -62,11 +63,15 @@ def setup_file_logging(formatter): log_file(), maxBytes=MAX_LOG_FILE_SIZE, backupCount=MAX_LOG_FILE_BACKUP_COUNT) file_handler.setFormatter(formatter) file_handler.setLevel(logging.DEBUG) + file_handler.name = "FILE_HANDLER" add_handler(file_handler) def add_handler(handler): logger = logging.getLogger(ROOT_LOGGER) + for hldr in logger.handlers: + if hldr.name == handler.name: + logger.removeHandler(hldr) logger.addHandler(handler) LOGGING_HANDLERS.append(handler) diff --git a/faraday_agent_dispatcher/utils/control_values_utils.py b/faraday_agent_dispatcher/utils/control_values_utils.py index 6464c0a0..16b01786 100644 --- a/faraday_agent_dispatcher/utils/control_values_utils.py +++ b/faraday_agent_dispatcher/utils/control_values_utils.py @@ -1,11 +1,15 @@ +def control_int(nullable=False): + def control(field_name, value): + if value is None and nullable: + return + if value is None: + raise ValueError(f"Trying to parse {field_name} with None value and should be an int") + try: + int(value) + except ValueError: + raise ValueError(f"Trying to parse {field_name} with value {value} and should be an int") -def control_int(field_name,value): - if value is None: - raise ValueError(f"Trying to parse {field_name} with None value and should be an int") - try: - int(value) - except ValueError: - raise ValueError(f"Trying to parse {field_name} with value {value} and should be an int") + return control def control_str(field_name, value): @@ -17,6 +21,23 @@ def control_host(field_name, value): control_str(field_name, value) + +def control_list(can_repeat=True): + def control(field_name, value): + if not isinstance(value, str): + raise ValueError(f"Trying to parse {field_name} with value {value} and should be a list") + listt = value.split(",") + if len(listt) != len(set(listt)): + raise ValueError(f"Trying to parse {field_name} with value {value} and contains repeated values") + + return control + + +def control_bool(field_name, value): + if value.lower() not in ["true", "false", "t", "f"]: + raise ValueError(f"Trying to parse {field_name} with value {value} and should be a bool") + + def control_registration_token(field_name, value): if value is None: raise ValueError(f'"{field_name}" option is required in the configuration ' diff --git a/setup.cfg b/setup.cfg index 2c38743e..1ce9b3c0 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.1.0 +current_version = 1.0.0 commit = True tag = True @@ -7,7 +7,7 @@ tag = True search = version='{current_version}' replace = version='{new_version}' -[bumpversion:file:dummy_faraday_agent/__init__.py] +[bumpversion:file:faraday_agent_dispatcher/__init__.py] search = __version__ = '{current_version}' replace = __version__ = '{new_version}' diff --git a/setup.py b/setup.py index 11b33cbf..c5c55d65 100644 --- a/setup.py +++ b/setup.py @@ -70,6 +70,6 @@ test_suite='tests', tests_require=test_requirements, url='https://github.com/infobyte/faraday_agent_dispatcher', - version='0.1', + version='1.0', zip_safe=False, ) diff --git a/tests/data/basic_executor.py b/tests/data/basic_executor.py index ab32d8f6..a3fa208c 100644 --- a/tests/data/basic_executor.py +++ b/tests/data/basic_executor.py @@ -13,9 +13,9 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . +import os import sys import json -import argparse host_data = { @@ -37,32 +37,32 @@ } if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--out', action='store', help='if set prints by stdout') - parser.add_argument('--err', action='store_true', help='if set prints by stderr') - parser.add_argument('--fails', action='store_true', help='if true fails') - parser.add_argument('--spaced_before', action='store_true', help='if set prints by stdout') - parser.add_argument('--spare', action='store_true', help='if set prints by stdout') - parser.add_argument('--spaced_middle', action='store_true', help='if set prints by stdout') - parser.add_argument('--count', action='store', default=1, help='if set prints by stdout') - args = parser.parse_args() - - if args.out: + out = os.getenv("EXECUTOR_CONFIG_OUT") + count = os.getenv("EXECUTOR_CONFIG_COUNT", 1) + err = os.getenv("EXECUTOR_CONFIG_ERR") is not None + fails = os.getenv("EXECUTOR_CONFIG_FAILS") is not None + spaced_before = os.getenv("EXECUTOR_CONFIG_SPACED_BEFORE") is not None + spaced_middle = os.getenv("EXECUTOR_CONFIG_SPACED_MIDDLE") is not None + spare = os.getenv("EXECUTOR_CONFIG_SPARE") is not None + omit_everything = os.getenv("DO_NOTHING", None) + if out and omit_everything is None: host_data_ = host_data.copy() host_data_['vulnerabilities'] = [vuln_data] data = dict(hosts=[host_data_]) - if args.out == "json": - prefix = '\n' if args.spaced_before else '' - suffix = '\n' if args.spaced_middle else '' - suffix += ('\n' if args.spare else '').join([''] + [json.dumps(data) for _ in range(int(args.count) - 1)]) + if out == "json": + prefix = '\n' if spaced_before else '' + suffix = '\n' if spaced_middle else '' + suffix += ('\n' if spare else '').join([''] + [json.dumps(data) for _ in range(int(count) - 1)]) print(f"{prefix}{json.dumps(data)}{suffix}") - elif args.out == "str": + elif out == "str": print("NO JSON OUTPUT") - elif args.out == "bad_json": + elif out == "bad_json": del data["hosts"][0]["ip"] print(f"{json.dumps(data)}") + else: + print(omit_everything, file=sys.stderr) - if args.err: + if err: print("Print by stderr", file=sys.stderr) - if args.fails: + if fails: sys.exit(1) \ No newline at end of file diff --git a/tests/data/test_config.ini b/tests/data/test_config.ini index 2431bce0..f88c7df0 100644 --- a/tests/data/test_config.ini +++ b/tests/data/test_config.ini @@ -4,13 +4,23 @@ host = localhost api_port = 5985 websocket_port = 9000 -[executor] -; Complete the cmd option with the command you want the dispatcher to run -cmd = exit 1 +[agent] agent_name = unnamed_agent -max_size = 65536 +executors = ex1 [tokens] ; To get your registration token, visit http://localhost:5985/#/admin/agents, copy ; the value and uncomment the line registration = 1234567890123456789012345 + +[ex1] +; Complete the cmd option with the command you want the dispatcher to run +cmd = exit 1 + +[ex2] +; Complete the cmd option with the command you want the dispatcher to run +cmd = exit 1 + +[ex3] +; Complete the cmd option with the command you want the dispatcher to run +cmd = exit 1 diff --git a/tests/integration/faraday/test_execution.py b/tests/integration/faraday/test_execution.py index 8657883e..81db8042 100644 --- a/tests/integration/faraday/test_execution.py +++ b/tests/integration/faraday/test_execution.py @@ -1,5 +1,5 @@ -from faraday_agent_dispatcher.config import instance as config, reset_config, \ - EXECUTOR_SECTION, SERVER_SECTION, TOKENS_SECTION, save_config, EXAMPLE_CONFIG_FILENAME +from faraday_agent_dispatcher.config import instance as config, reset_config, Sections,\ + save_config, EXAMPLE_CONFIG_FILENAME from faraday_agent_dispatcher.utils.url_utils import api_url @@ -13,11 +13,12 @@ import time reset_config(EXAMPLE_CONFIG_FILENAME) -HOST = config.get(SERVER_SECTION, "host") -API_PORT = config.get(SERVER_SECTION, "api_port") -WS_PORT = config.get(SERVER_SECTION, "websocket_port") +HOST = config.get(Sections.SERVER, "host") +API_PORT = config.get(Sections.SERVER, "api_port") +WS_PORT = config.get(Sections.SERVER, "websocket_port") WORKSPACE = fuzzy_string(6).lower() # TODO FIX WHEN FARADAY ACCEPTS CAPITAL FIRST LETTER AGENT_NAME = fuzzy_string(6) +EXECUTOR_NAME = fuzzy_string(6) USER = os.getenv("FARADAY_USER") EMAIL = os.getenv("FARADAY_EMAIL") @@ -52,20 +53,31 @@ def test_execute_agent(): token = res.json()['token'] # Config set up - config.set(TOKENS_SECTION, "registration", token) - config.remove_option(TOKENS_SECTION, "agent") - config.set(SERVER_SECTION, "workspace", WORKSPACE) - config.set(EXECUTOR_SECTION, "agent_name", AGENT_NAME) - config.set(EXECUTOR_SECTION, "cmd", "python ./basic_executor.py --out json") + config.set(Sections.TOKENS, "registration", token) + config.remove_option(Sections.TOKENS, "agent") + config.set(Sections.SERVER, "workspace", WORKSPACE) + config.set(Sections.AGENT, "agent_name", AGENT_NAME) + config.set(Sections.AGENT, "executors", EXECUTOR_NAME) path_to_basic_executor = ( Path(__file__).parent.parent.parent / 'data' / 'basic_executor.py' ) + executor_section = Sections.EXECUTOR_DATA.format(EXECUTOR_NAME) + params_section = Sections.EXECUTOR_PARAMS.format(EXECUTOR_NAME) + for section in [executor_section, params_section]: + if section not in config: + config.add_section(section) + config.set( - EXECUTOR_SECTION, + Sections.EXECUTOR_DATA.format(EXECUTOR_NAME), "cmd", - f"python {path_to_basic_executor} --out json" + f"python {path_to_basic_executor}" ) + + config.set(params_section, "out", "True") + [config.set(params_section, param, "False") for param in [ + "count", "spare", "spaced_before", "spaced_middle", "err", "fails"]] + save_config(CONFIG_DIR) # Init dispatcher! @@ -79,6 +91,7 @@ def test_execute_agent(): res_data = res.json() assert len(res_data) == 1, p.communicate(timeout=0.1) agent = res_data[0] + agent_id = agent["id"] if agent_ok_status_keys_set != set(agent.keys()): print("Keys set from agent endpoint differ from expected ones, checking if its a superset") assert agent_ok_status_keys_set.issubset(set(agent.keys())) @@ -87,17 +100,27 @@ def test_execute_agent(): # Run executor! res = session.post(api_url(HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/agents/{agent["id"]}/run/'), - data={'csrf_token': session_res.json()['csrf_token']}) + json={ + 'csrf_token': session_res.json()['csrf_token'], + 'executorData': { + "agent_id": agent_id, + "executor": EXECUTOR_NAME, + "args": {"out": "json"} + } + }) assert res.status_code == 200, res.text time.sleep(2) # If fails check time # Test results res = session.get(api_url(HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/hosts')) host_dict = res.json() - assert host_dict["total_rows"] == 1 + assert host_dict["total_rows"] == 1, (res.text, host_dict) host = host_dict["rows"][0]["value"] for key in host_data: - assert host[key] == host_data[key] + if key == "hostnames": + assert set(host[key]) == set(host_data[key]) + else: + assert host[key] == host_data[key] assert host["vulns"] == 1 res = session.get(api_url(HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/vulns')) diff --git a/tests/unittests/test_agent_dispatcher.py b/tests/unittests/test_agent_dispatcher.py index 4b484678..eba89364 100644 --- a/tests/unittests/test_agent_dispatcher.py +++ b/tests/unittests/test_agent_dispatcher.py @@ -16,7 +16,7 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -"""Tests for `faraday_dummy_agent` package.""" +"""Tests for `faraday_agent_dispatcher` package.""" import json import os @@ -31,79 +31,141 @@ reset_config, save_config, instance as configuration, - SERVER_SECTION, - TOKENS_SECTION, - EXECUTOR_SECTION + Sections ) from tests.utils.text_utils import fuzzy_string from tests.utils.testing_faraday_server import FaradayTestConfig, test_config, tmp_custom_config, tmp_default_config, \ - test_logger_handler + test_logger_handler, test_logger_folder @pytest.mark.parametrize('config_changes_dict', - [{"remove": {SERVER_SECTION: ["host"]}, + [{"remove": {Sections.SERVER: ["host"]}, "replace": {}, "expected_exception": ValueError}, - {"remove": {SERVER_SECTION: ["api_port"]}, + {"remove": {Sections.SERVER: ["api_port"]}, "replace": {}, "expected_exception": ValueError}, {"remove": {}, - "replace": {SERVER_SECTION: {"api_port": "Not a port number"}}, + "replace": {Sections.SERVER: {"api_port": "Not a port number"}}, "expected_exception": ValueError}, {"remove": {}, - "replace": {SERVER_SECTION: {"api_port": "6000"}}}, # None error as parse int - {"remove": {SERVER_SECTION: ["websocket_port"]}, + "replace": {Sections.SERVER: {"api_port": "6000"}}}, # None error as parse int + {"remove": {Sections.SERVER: ["websocket_port"]}, "replace": {}, "expected_exception": ValueError}, {"remove": {}, - "replace": {SERVER_SECTION: {"websocket_port": "Not a port number"}}, + "replace": {Sections.SERVER: {"websocket_port": "Not a port number"}}, "expected_exception": ValueError}, {"remove": {}, - "replace": {SERVER_SECTION: {"websocket_port": "9001"}}}, # None error as parse int - {"remove": {SERVER_SECTION: ["workspace"]}, + "replace": {Sections.SERVER: {"websocket_port": "9001"}}}, # None error as parse int + {"remove": {Sections.SERVER: ["workspace"]}, "replace": {}, "expected_exception": ValueError}, - {"remove": {TOKENS_SECTION: ["registration"]}, + {"remove": {Sections.TOKENS: ["registration"]}, "replace": {}, "expected_exception": ValueError}, {"remove": {}, - "replace": {TOKENS_SECTION: {"registration": "invalid_token"}}, + "replace": {Sections.TOKENS: {"registration": "invalid_token"}}, "expected_exception": ValueError}, {"remove": {}, - "replace": {TOKENS_SECTION: {"registration": " 46aasdje446aasdje446aa"}}, + "replace": {Sections.TOKENS: {"registration": " 46aasdje446aasdje446aa"}}, "expected_exception": ValueError}, {"remove": {}, - "replace": {TOKENS_SECTION: {"registration": "QWE46aasdje446aasdje446aa"}}}, + "replace": {Sections.TOKENS: {"registration": "QWE46aasdje446aasdje446aa"}}}, {"remove": {}, - "replace": {TOKENS_SECTION: {"agent": "invalid_token"}}, + "replace": {Sections.TOKENS: {"agent": "invalid_token"}}, "expected_exception": ValueError}, {"remove": {}, - "replace": {TOKENS_SECTION: { + "replace": {Sections.TOKENS: { "agent": " 46aasdje446aasdje446aa46aasdje446aasdje446aa46aasdje446aasdje" }}, "expected_exception": ValueError}, {"remove": {}, - "replace": {TOKENS_SECTION: { + "replace": {Sections.TOKENS: { "agent": "QWE46aasdje446aasdje446aaQWE46aasdje446aasdje446aaQWE46aasdje446"}}}, - {"remove": {EXECUTOR_SECTION: ["cmd"]}, + {"remove": {Sections.EXECUTOR_DATA.format("ex1"): ["cmd"]}, "replace": {}, "expected_exception": ValueError}, - {"remove": {EXECUTOR_SECTION: ["agent_name"]}, + {"remove": {}, + "replace": {Sections.EXECUTOR_DATA.format("ex1"): {"max_size": "ASDASD"}}, + "expected_exception": ValueError}, + {"remove": {}, + "replace": {Sections.EXECUTOR_PARAMS.format("ex1"): {"param1": "ASDASD"}}, + "expected_exception": ValueError}, + {"remove": {}, + "replace": {Sections.EXECUTOR_PARAMS.format("ex1"): {"param1": "5"}}, + "expected_exception": ValueError}, + {"remove": {}, + "replace": {Sections.EXECUTOR_PARAMS.format("ex1"): {"param1": "True"}} + }, + {"remove": {Sections.AGENT: ["agent_name"]}, "replace": {}, "expected_exception": ValueError}, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": "ex1,ex1"}}, + "expected_exception": ValueError + }, + {"remove": {Sections.AGENT: ["section"]}, + "replace": {}, + "expected_exception": ValueError + }, + {"remove": {Sections.TOKENS: ["section"]}, + "replace": {}, + "expected_exception": ValueError + }, + {"remove": {Sections.SERVER: ["section"]}, + "replace": {}, + "expected_exception": ValueError + }, + {"remove": {}, + "replace": {}, + "duplicate_exception": True, + "expected_exception": ValueError + }, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": "ex1, ex2"}}, + }, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": "ex1,ex2 "}}, + }, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": " ex1,ex2"}}, + }, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": " ex1, ex2 , ex3"}}, + }, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": "ex1,ex 1"}}, + "expected_exception": ValueError + }, + {"remove": {}, + "replace": {Sections.AGENT: {"executors": "ex1,ex8"}}, + "expected_exception": ValueError + }, {"remove": {}, "replace": {}} ]) def test_basic_built(tmp_custom_config, config_changes_dict): for section in config_changes_dict["replace"]: for option in config_changes_dict["replace"][section]: + if section not in configuration: + configuration.add_section(section) configuration.set(section, option, config_changes_dict["replace"][section][option]) for section in config_changes_dict["remove"]: - for option in config_changes_dict["remove"][section]: - configuration.remove_option(section, option) + if "section" in config_changes_dict["remove"][section]: + configuration.remove_section(section) + else: + for option in config_changes_dict["remove"][section]: + configuration.remove_option(section, option) tmp_custom_config.save() if "expected_exception" in config_changes_dict: + if "duplicate_exception" in config_changes_dict and config_changes_dict["duplicate_exception"]: + with open(tmp_custom_config.config_file_path, "r") as file: + content = file.read() + with open(tmp_custom_config.config_file_path, "w") as file: + file.write(content) + file.write(content) with pytest.raises(config_changes_dict["expected_exception"]): Dispatcher(None, tmp_custom_config.config_file_path) else: @@ -112,11 +174,11 @@ def test_basic_built(tmp_custom_config, config_changes_dict): async def test_start_and_register(test_config: FaradayTestConfig, tmp_default_config): # Config - configuration.set(SERVER_SECTION, "api_port", str(test_config.client.port)) - configuration.set(SERVER_SECTION, "host", test_config.client.host) - configuration.set(SERVER_SECTION, "workspace", test_config.workspace) - configuration.set(TOKENS_SECTION, "registration", test_config.registration_token) - configuration.set(EXECUTOR_SECTION, "cmd", 'exit 1') + configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) + configuration.set(Sections.SERVER, "host", test_config.client.host) + configuration.set(Sections.SERVER, "workspace", test_config.workspace) + configuration.set(Sections.TOKENS, "registration", test_config.registration_token) + configuration.set(Sections.EXECUTOR_DATA.format("ex1"), "cmd", 'exit 1') tmp_default_config.save() # Init and register it @@ -133,11 +195,11 @@ async def test_start_and_register(test_config: FaradayTestConfig, tmp_default_co async def test_start_with_bad_config(test_config: FaradayTestConfig, tmp_default_config): # Config - configuration.set(SERVER_SECTION, "api_port", str(test_config.client.port)) - configuration.set(SERVER_SECTION, "host", test_config.client.host) - configuration.set(SERVER_SECTION, "workspace", test_config.workspace) - configuration.set(TOKENS_SECTION, "registration", "NotOk" * 5) - configuration.set(EXECUTOR_SECTION, "cmd", 'exit 1') + configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) + configuration.set(Sections.SERVER, "host", test_config.client.host) + configuration.set(Sections.SERVER, "workspace", test_config.workspace) + configuration.set(Sections.TOKENS, "registration", "NotOk" * 5) + configuration.set(Sections.EXECUTOR_DATA.format("ex1"), "cmd", 'exit 1') tmp_default_config.save() # Init and register it @@ -149,170 +211,593 @@ async def test_start_with_bad_config(test_config: FaradayTestConfig, tmp_default @pytest.mark.parametrize('executor_options', [ - { + { # 0 "data": {"agent_id": 1}, - "args": [], "logs": [ {"levelname": "INFO", "msg": "Data not contains action to do"}, + ], + "ws_responses": [ + {"error": "'action' key is mandatory in this websocket connection"} ] }, - { + { # 1 "data": {"action": "CUT", "agent_id": 1}, - "args": [], "logs": [ - {"levelname": "INFO", "msg": "Action unrecognized"}, + {"levelname": "INFO", "msg": "Unrecognized action"}, + ], + "ws_responses": [ + {"CUT_RESPONSE": "Error: Unrecognized action"} ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json"], + { # 2 + "data": {"action": "RUN", "agent_id": 1, "executor": "ex1", "args": {"out": "json"}}, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "INFO", "msg": "Data sent to bulk create"}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json", "count 5"], + { # 3 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", + "args": {"out": "json", "count": "5"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "ERROR", "msg": "JSON Parsing error: Extra data"}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json", "count 5", "spare"], + { # 4 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", + "args": {"out": "json", "count": "5", "spare": "T"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "INFO", "msg": "Data sent to bulk create", "min_count": 5}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json", "spaced_before"], + { # 5 + "data": { + "action": "RUN", + "agent_id": 1, + "executor": "ex1", + "args": {"out": "json", "spaced_before": "T"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Running ex1 executor"}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json", "spaced_middle", "count 5", "spare"], + { # 6 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", + "args": {"out": "json", "spaced_middle": "T", "count": "5", "spare": "T"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "INFO", "msg": "Data sent to bulk create", "max_count": 1}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out bad_json"], + { # 7 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", "args": {"out": "bad_json"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "ERROR", - "msg": "Invalid data supplied by the executor to the bulk create endpoint. Server responded: "}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + "msg": "Invalid data supplied by the executor to the bulk create endpoint. " + "Server responded: "}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out str"], + { # 8 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", "args": {"out": "str"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "ERROR", "msg": "JSON Parsing error: Expecting value"}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["err"], + { # 9 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", + "args": {"out": "none", "err": "T"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "DEBUG", "msg": "Print by stderr"}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["fails"], + { # 10 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", + "args": {"out": "none", "fails": "T"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, - {"levelname": "WARNING", "msg": "Executor finished with exit code 1"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, + {"levelname": "WARNING", "msg": "Executor ex1 finished with exit code 1"}, + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": False, + "message": "Executor ex1 from unnamed_agent failed" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["err", "fails"], + { # 11 + "data": { + "action": "RUN", + "agent_id": 1, + "executor": "ex1", + "args": {"out": "none", "err": "T", "fails": "T"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "DEBUG", "msg": "Print by stderr"}, - {"levelname": "WARNING", "msg": "Executor finished with exit code 1"}, + {"levelname": "WARNING", "msg": "Executor ex1 finished with exit code 1"}, + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": False, + "message": "Executor ex1 from unnamed_agent failed" + } ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json"], + { # 12 + "data": { + "action": "RUN", + "agent_id": 1, + "executor": "ex1", + "args": {"out": "json"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, + {"levelname": "INFO", "msg": "Data sent to bulk create", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "varenvs": {"DO_NOTHING": "True"}, + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } + ] + }, + { # 13 + "data": { + "action": "RUN", + "agent_id": 1, + "executor": "ex1", + "args": {"err": "T", "fails": "T"}, + }, + "logs": [ + {"levelname": "INFO", "msg": "Running ex1 executor", "max_count": 0, + "min_count": 0}, + {"levelname": "ERROR", "msg": "Mandatory argument not passed"}, + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": False, + "message": "Mandatory argument(s) not passed to ex1 executor from " + "unnamed_agent agent" + } + ] + }, + { # 14 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", + "args": {"out": "json", "WTF": "T"} + }, + "logs": [ + {"levelname": "INFO", "msg": "Running ex1 executor", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Data sent to bulk create", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully", "max_count": 0, + "min_count": 0}, + {"levelname": "ERROR", "msg": "Unexpected argument passed"}, + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": False, + "message": "Unexpected argument(s) passed to ex1 executor from unnamed_agent " + "agent" + } + ] + }, + { # 15 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", "args": {"out": "json"} + }, + "logs": [ + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "ERROR", - "msg": "Invalid data supplied by the executor to the bulk create endpoint. Server responded: "}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + "msg": "Invalid data supplied by the executor to the bulk create endpoint. " + "Server responded: "}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} ], - "workspace": "error500" + "workspace": "error500", + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } + ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json"], + { # 16 + "data": { + "action": "RUN", "agent_id": 1, "executor": "ex1", "args": {"out": "json"} + }, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "ERROR", - "msg": "Invalid data supplied by the executor to the bulk create endpoint. Server responded: "}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + "msg": "Invalid data supplied by the executor to the bulk create endpoint. " + "Server responded: "}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} ], - "workspace": "error429" + "workspace": "error429", + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } + ] }, - { - "data": {"action": "RUN", "agent_id": 1}, - "args": ["out json"], + { # 17 + "data": {"action": "RUN", "agent_id": 1, "executor": "ex1", "args": {"out": "json"}}, "logs": [ - {"levelname": "INFO", "msg": "Running executor"}, + {"levelname": "INFO", "msg": "Running ex1 executor"}, {"levelname": "ERROR", "msg": "ValueError raised processing stdout, try with " "bigger limiting size in config"}, - {"levelname": "INFO", "msg": "Executor finished successfully"} + {"levelname": "INFO", "msg": "Executor ex1 finished successfully"} + ], + "max_size": "1", + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "ex1", + "running": True, + "message": "Running ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "ex1", + "successful": True, + "message": "Executor ex1 from unnamed_agent finished successfully" + } + ] + }, + { # 18 + "data": { + "action": "RUN", "agent_id": 1, + "args": {"out": "json"} + }, + "logs": [ + {"levelname": "INFO", "msg": "Running ex1 executor", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Data sent to bulk create", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully", "max_count": 0, + "min_count": 0}, + {"levelname": "ERROR", "msg": "No executor selected"}, + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "running": False, + "message": "No executor selected to unnamed_agent agent" + } + ] + }, + { # 19 + "data": { + "action": "RUN", "agent_id": 1, "executor": "NOT_4N_CORRECT_EXECUTOR", + "args": {"out": "json"} + }, + "logs": [ + {"levelname": "INFO", "msg": "Running ex1 executor", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Data sent to bulk create", "max_count": 0, + "min_count": 0}, + {"levelname": "INFO", "msg": "Executor ex1 finished successfully", "max_count": 0, + "min_count": 0}, + {"levelname": "ERROR", "msg": "The selected executor not exists"}, + ], + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "NOT_4N_CORRECT_EXECUTOR", + "running": False, + "message": "The selected executor NOT_4N_CORRECT_EXECUTOR not exists in " + "unnamed_agent agent"} + ] + }, + { # 20 + "data": {"action": "RUN", "agent_id": 1, "executor": "add_ex1", "args": {"out": "json"}}, + "logs": [ + {"levelname": "INFO", "msg": "Running add_ex1 executor"}, + {"levelname": "INFO", "msg": "Data sent to bulk create"}, + {"levelname": "INFO", "msg": "Executor add_ex1 finished successfully"} ], - "max_size": "1" + "ws_responses": [ + { + "action": "RUN_STATUS", + "executor_name": "add_ex1", + "running": True, + "message": "Running add_ex1 executor from unnamed_agent agent" + }, { + "action": "RUN_STATUS", + "executor_name": "add_ex1", + "successful": True, + "message": "Executor add_ex1 from unnamed_agent finished successfully" + } + ], + "extra": ["add_ex1"] }, ]) -async def test_run_once(test_config: FaradayTestConfig, tmp_default_config, test_logger_handler, executor_options): +async def test_run_once(test_config: FaradayTestConfig, tmp_default_config, test_logger_handler, + test_logger_folder, executor_options): # Config workspace = test_config.workspace if "workspace" not in executor_options else executor_options["workspace"] - configuration.set(SERVER_SECTION, "api_port", str(test_config.client.port)) - configuration.set(SERVER_SECTION, "host", test_config.client.host) - configuration.set(SERVER_SECTION, "workspace", workspace) - configuration.set(TOKENS_SECTION, "registration", test_config.registration_token) - configuration.set(TOKENS_SECTION, "agent", test_config.agent_token) + configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) + configuration.set(Sections.SERVER, "host", test_config.client.host) + configuration.set(Sections.SERVER, "workspace", workspace) + configuration.set(Sections.TOKENS, "registration", test_config.registration_token) + configuration.set(Sections.TOKENS, "agent", test_config.agent_token) path_to_basic_executor = ( - Path(__file__).parent.parent / - 'data' / 'basic_executor.py' + Path(__file__).parent.parent / + 'data' / 'basic_executor.py' ) - args = ' --'.join([''] + executor_options['args']) - configuration.set(EXECUTOR_SECTION, "cmd", f"python {path_to_basic_executor} {args}") + executor_names = ["ex1"] + ([] if "extra" not in executor_options else executor_options["extra"]) + configuration.set(Sections.AGENT, "executors", ",".join(executor_names)) + for executor_name in executor_names: + executor_section = Sections.EXECUTOR_DATA.format(executor_name) + params_section = Sections.EXECUTOR_PARAMS.format(executor_name) + varenvs_section = Sections.EXECUTOR_VARENVS.format(executor_name) + for section in [executor_section, params_section, varenvs_section]: + if section not in configuration: + configuration.add_section(section) + + configuration.set(executor_section, "cmd", "python {}".format(path_to_basic_executor)) + configuration.set(params_section, "out", "True") + [configuration.set(params_section, param, "False") for param in [ + "count", "spare", "spaced_before", "spaced_middle", "err", "fails"]] + if "varenvs" in executor_options: + for varenv in executor_options["varenvs"]: + configuration.set(varenvs_section, varenv, executor_options["varenvs"][varenv]) - max_size = str(64 * 1024) if "max_size" not in executor_options else executor_options["max_size"] - configuration.set(EXECUTOR_SECTION, "max_size", max_size) + max_size = str(64 * 1024) if "max_size" not in executor_options else executor_options["max_size"] + configuration.set(executor_section, "max_size", max_size) tmp_default_config.save() + async def ws_messages_checker(msg): + msg_ = json.loads(msg) + assert msg_ in executor_options["ws_responses"] + executor_options["ws_responses"].remove(msg_) + # Init and register it dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path) - await dispatcher.run_once(json.dumps(executor_options["data"])) + await dispatcher.run_once(json.dumps(executor_options["data"]), ws_messages_checker) history = test_logger_handler.history + assert len(executor_options["ws_responses"]) == 0 for l in executor_options["logs"]: min_count = 1 if "min_count" not in l else l["min_count"] max_count = sys.maxsize if "max_count" not in l else l["max_count"] assert max_count >= \ len(list(filter(lambda x: x.levelname == l["levelname"] and l["msg"] in x.message, history))) >= \ - min_count \ No newline at end of file + min_count, l["msg"] + + +async def test_connect(test_config: FaradayTestConfig, tmp_default_config, test_logger_handler, + test_logger_folder): + configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) + configuration.set(Sections.SERVER, "host", test_config.client.host) + configuration.set(Sections.SERVER, "workspace", test_config.workspace) + configuration.set(Sections.TOKENS, "registration", test_config.registration_token) + configuration.set(Sections.TOKENS, "agent", test_config.agent_token) + path_to_basic_executor = ( + Path(__file__).parent.parent / + 'data' / 'basic_executor.py' + ) + configuration.set(Sections.AGENT, "executors", "ex1,ex2,ex3") + + for executor_name in ["ex1","ex2","ex3"]: + executor_section = Sections.EXECUTOR_DATA.format(executor_name) + params_section = Sections.EXECUTOR_PARAMS.format(executor_name) + for section in [executor_section, params_section]: + if section not in configuration: + configuration.add_section(section) + configuration.set(executor_section, "cmd", "python {}".format(path_to_basic_executor)) + + configuration.set(Sections.EXECUTOR_PARAMS.format("ex1"), "param1", "True") + configuration.set(Sections.EXECUTOR_PARAMS.format("ex1"), "param2", "False") + configuration.set(Sections.EXECUTOR_PARAMS.format("ex2"), "param3", "False") + configuration.set(Sections.EXECUTOR_PARAMS.format("ex2"), "param4", "False") + tmp_default_config.save() + dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path) + + ws_responses = [{ + 'action': 'JOIN_AGENT', + 'workspace': test_config.workspace, + 'token': None, + 'executors': [ + { + "executor_name": "ex1", + "args": { + "param1": True, + "param2": False + } + }, + { + "executor_name": "ex2", + "args": { + "param3": False, + "param4": False + } + }, + { + "executor_name": "ex3", + "args": {} + } + ] + }] + + async def ws_messages_checker(msg): + msg_ = json.loads(msg) + assert msg_ in ws_responses + ws_responses.remove(msg_) + + await dispatcher.connect(ws_messages_checker) + + assert len(ws_responses) == 0 diff --git a/tests/utils/testing_faraday_server.py b/tests/utils/testing_faraday_server.py index 09e10ddc..0663ce0f 100644 --- a/tests/utils/testing_faraday_server.py +++ b/tests/utils/testing_faraday_server.py @@ -9,7 +9,7 @@ from itsdangerous import TimestampSigner import logging from logging import StreamHandler -from faraday_agent_dispatcher.logger import get_logger +from faraday_agent_dispatcher.logger import get_logger, reset_logger from queue import Queue from faraday_agent_dispatcher.config import ( @@ -166,12 +166,18 @@ class TestLoggerHandler(StreamHandler): def __init__(self): super().__init__() self.history = [] + self.name = "TEST_HANDLER" def emit(self, record): self.history.append(record) -@pytest.fixture +@pytest.fixture(scope="session") +def test_logger_folder(): + reset_logger("./logs") + + +@pytest.fixture() def test_logger_handler(): logger_handler = TestLoggerHandler() logger = get_logger() diff --git a/tox.ini b/tox.ini index ba9723ce..0709d924 100644 --- a/tox.ini +++ b/tox.ini @@ -11,7 +11,7 @@ python = [testenv:flake8] basepython = python deps = flake8 -commands = flake8 dummy_faraday_agent +commands = flake8 faraday_agent_dispatcher [testenv] setenv =