From 2ed228f2834975f81505c5e98e1d72c244a0efa5 Mon Sep 17 00:00:00 2001 From: itay Date: Sun, 23 Jun 2019 16:01:08 +0300 Subject: [PATCH] Refactor scan,trace,tunnel,pba telems --- monkey/infection_monkey/monkey.py | 11 ++++--- monkey/infection_monkey/post_breach/pba.py | 8 ++--- .../telemetry/post_breach_telem.py | 31 +++++++++++++++++++ .../infection_monkey/telemetry/scan_telem.py | 22 +++++++++++++ .../infection_monkey/telemetry/trace_telem.py | 26 ++++++++++++++++ .../telemetry/tunnel_telem.py | 19 ++++++++++++ .../cc/resources/telemetry_feed.py | 2 +- 7 files changed, 107 insertions(+), 12 deletions(-) create mode 100644 monkey/infection_monkey/telemetry/post_breach_telem.py create mode 100644 monkey/infection_monkey/telemetry/scan_telem.py create mode 100644 monkey/infection_monkey/telemetry/trace_telem.py create mode 100644 monkey/infection_monkey/telemetry/tunnel_telem.py diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 9c1ffb419..316cb6b31 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -16,7 +16,10 @@ from infection_monkey.network.network_scanner import NetworkScanner from infection_monkey.system_info import SystemInfoCollector from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem +from infection_monkey.telemetry.scan_telem import ScanTelem from infection_monkey.telemetry.state_telem import StateTelem +from infection_monkey.telemetry.trace_telem import TraceTelem +from infection_monkey.telemetry.tunnel_telem import TunnelTelem from infection_monkey.windows_upgrader import WindowsUpgrader from infection_monkey.post_breach.post_breach_handler import PostBreach from common.utils.attack_utils import ScanStatus @@ -114,7 +117,7 @@ class InfectionMonkey(object): self._default_server = WormConfiguration.current_server LOG.debug("default server: %s" % self._default_server) - ControlClient.send_telemetry("tunnel", {'proxy': ControlClient.proxies.get('https')}) + TunnelTelem().send() if WormConfiguration.collect_system_info: LOG.debug("Calling system info collection") @@ -126,8 +129,7 @@ class InfectionMonkey(object): PostBreach().execute() if 0 == WormConfiguration.depth: - LOG.debug("Reached max depth, shutting down") - ControlClient.send_telemetry("trace", "Reached max depth, shutting down") + TraceTelem("Reached max depth, shutting down").send() return else: LOG.debug("Running with depth: %d" % WormConfiguration.depth) @@ -158,8 +160,7 @@ class InfectionMonkey(object): machine, finger.__class__.__name__) finger.get_host_fingerprint(machine) - ControlClient.send_telemetry('scan', {'machine': machine.as_dict(), - 'service_count': len(machine.services)}) + ScanTelem(machine).send() # skip machines that we've already exploited if machine in self._exploited_machines: diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 7df3693fa..3fb8b251f 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -2,6 +2,7 @@ import logging import subprocess import socket from infection_monkey.control import ControlClient +from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils import is_windows_os from infection_monkey.config import WormConfiguration @@ -45,12 +46,7 @@ class PBA(object): """ exec_funct = self._execute_default result = exec_funct() - hostname = socket.gethostname() - ControlClient.send_telemetry('post_breach', {'command': self.command, - 'result': result, - 'name': self.name, - 'hostname': hostname, - 'ip': socket.gethostbyname(hostname)}) + PostBreachTelem(self, result).send() def _execute_default(self): """ diff --git a/monkey/infection_monkey/telemetry/post_breach_telem.py b/monkey/infection_monkey/telemetry/post_breach_telem.py new file mode 100644 index 000000000..4cb5b50df --- /dev/null +++ b/monkey/infection_monkey/telemetry/post_breach_telem.py @@ -0,0 +1,31 @@ +import socket + +from infection_monkey.telemetry.base_telem import BaseTelem + +__author__ = "itay.mizeretz" + + +class PostBreachTelem(BaseTelem): + + def __init__(self, pba, result): + """ + Default post breach telemetry constructor + :param pba: Post breach action which was used + :param result: Result of PBA + """ + super(PostBreachTelem, self).__init__() + self.pba = pba + self.result = result + self.hostname = socket.gethostname() + self.ip = socket.gethostbyname(self.hostname) + + telem_category = 'post_breach' + + def get_data(self): + return { + 'command': self.pba.command, + 'result': self.result, + 'name': self.pba.name, + 'hostname': self.hostname, + 'ip': self.ip + } diff --git a/monkey/infection_monkey/telemetry/scan_telem.py b/monkey/infection_monkey/telemetry/scan_telem.py new file mode 100644 index 000000000..b1c58ab1b --- /dev/null +++ b/monkey/infection_monkey/telemetry/scan_telem.py @@ -0,0 +1,22 @@ +from infection_monkey.telemetry.base_telem import BaseTelem + +__author__ = "itay.mizeretz" + + +class ScanTelem(BaseTelem): + + def __init__(self, machine): + """ + Default scan telemetry constructor + :param machine: Scanned machine + """ + super(ScanTelem, self).__init__() + self.machine = machine + + telem_category = 'scan' + + def get_data(self): + return { + 'machine': self.machine.as_dict(), + 'service_count': len(self.machine.services) + } diff --git a/monkey/infection_monkey/telemetry/trace_telem.py b/monkey/infection_monkey/telemetry/trace_telem.py new file mode 100644 index 000000000..0782affb4 --- /dev/null +++ b/monkey/infection_monkey/telemetry/trace_telem.py @@ -0,0 +1,26 @@ +import logging + +from infection_monkey.telemetry.base_telem import BaseTelem + +__author__ = "itay.mizeretz" + +LOG = logging.getLogger(__name__) + + +class TraceTelem(BaseTelem): + + def __init__(self, msg): + """ + Default trace telemetry constructor + :param msg: Trace message + """ + super(TraceTelem, self).__init__() + self.msg = msg + LOG.debug("Trace: %s" % msg) + + telem_category = 'trace' + + def get_data(self): + return { + 'msg': self.msg + } diff --git a/monkey/infection_monkey/telemetry/tunnel_telem.py b/monkey/infection_monkey/telemetry/tunnel_telem.py new file mode 100644 index 000000000..64533a252 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tunnel_telem.py @@ -0,0 +1,19 @@ +from infection_monkey.control import ControlClient +from infection_monkey.telemetry.base_telem import BaseTelem + +__author__ = "itay.mizeretz" + + +class TunnelTelem(BaseTelem): + + def __init__(self): + """ + Default tunnel telemetry constructor + """ + super(TunnelTelem, self).__init__() + self.proxy = ControlClient.proxies.get('https') + + telem_category = 'tunnel' + + def get_data(self): + return {'proxy': self.proxy} diff --git a/monkey/monkey_island/cc/resources/telemetry_feed.py b/monkey/monkey_island/cc/resources/telemetry_feed.py index b98d650c8..db9d2e905 100644 --- a/monkey/monkey_island/cc/resources/telemetry_feed.py +++ b/monkey/monkey_island/cc/resources/telemetry_feed.py @@ -78,7 +78,7 @@ class TelemetryFeed(flask_restful.Resource): @staticmethod def get_trace_telem_brief(telem): - return 'Monkey reached max depth.' + return 'Trace: %s' % telem['data']['msg'] @staticmethod def get_post_breach_telem_brief(telem):