Refactor scan,trace,tunnel,pba telems

This commit is contained in:
itay 2019-06-23 16:01:08 +03:00
parent 27ca921dbc
commit 2ed228f283
7 changed files with 107 additions and 12 deletions

View File

@ -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:

View File

@ -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):
"""

View File

@ -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
}

View File

@ -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)
}

View File

@ -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
}

View File

@ -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}

View File

@ -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):