diff --git a/monkey/infection_monkey/model/host.py b/monkey/infection_monkey/model/host.py index dcc6e7455..3d8c3d0e6 100644 --- a/monkey/infection_monkey/model/host.py +++ b/monkey/infection_monkey/model/host.py @@ -46,3 +46,6 @@ class VictimHost(object): def set_default_server(self, default_server): self.default_server = default_server + + def is_linux(self): + return 'linux' in self.os['type'] diff --git a/monkey/infection_monkey/post_breach/post_breach.py b/monkey/infection_monkey/post_breach/post_breach.py new file mode 100644 index 000000000..24274582a --- /dev/null +++ b/monkey/infection_monkey/post_breach/post_breach.py @@ -0,0 +1,45 @@ +import logging +import infection_monkey.config +import subprocess +from abc import abstractmethod + +LOG = logging.getLogger(__name__) + +__author__ = 'VakarisZ' + + +# Class that handles post breach action execution +class PostBreach(object): + def __init__(self, host, pba_list): + self._config = infection_monkey.config.WormConfiguration + self.pba_list = pba_list + self.host = host + + def execute(self): + for pba in self.pba_list: + if self.host.is_linux(): + pba.execute_linux() + else: + pba.execute_win() + + @staticmethod + @abstractmethod + def config_to_pba_list(config): + """ + Should return a list of PBA's generated from config + """ + raise NotImplementedError() + + +# Post Breach Action container +class PBA(object): + def __init__(self, linux_command="", windows_command=""): + self.linux_command = linux_command + self.windows_command = windows_command + + def execute_linux(self): + return subprocess.check_output(self.linux_command, shell=True) + + def execute_win(self): + return subprocess.check_output(self.windows_command, shell=True) + diff --git a/monkey/monkey_island/cc/services/config_schema.py b/monkey/monkey_island/cc/services/config_schema.py index cbcc6ba0a..e0c286065 100644 --- a/monkey/monkey_island/cc/services/config_schema.py +++ b/monkey/monkey_island/cc/services/config_schema.py @@ -298,10 +298,18 @@ SCHEMA = { }, "post_breach_actions": { "title": "Post breach actions", - "type": "array", - "uniqueItems": True, - "items": { - "$ref": "#/definitions/post_breach_acts" + "type": "object", + "properties": { + "linux": { + "title": "Linux command", + "type": "string", + "description": "Linux command to execute after breaching" + }, + "windows": { + "title": "Windows command", + "type": "string", + "description": "Windows command to execute after breaching" + } }, "default": [ ],