agent: Add ITelem interface

Create a telemetry interface that sits above the BaseTelem abstract
class to allow telemetries to be extended without inheritance.
This commit is contained in:
Mike Salvatore 2021-06-24 12:07:14 -04:00
parent 77e3c8a257
commit 46da0b7b1f
2 changed files with 31 additions and 19 deletions

View File

@ -3,6 +3,7 @@ import json
import logging
from infection_monkey.control import ControlClient
from infection_monkey.telemetry.i_telem import ITelem
logger = logging.getLogger(__name__)
LOGGED_DATA_LENGTH = 300 # How many characters of telemetry data will be logged
@ -24,14 +25,11 @@ __author__ = "itay.mizeretz"
# logging and sending them.
class BaseTelem(object, metaclass=abc.ABCMeta):
class BaseTelem(ITelem, metaclass=abc.ABCMeta):
"""
Abstract base class for telemetry.
"""
def __init__(self):
pass
def send(self, log_data=True):
"""
Sends telemetry to island
@ -41,13 +39,6 @@ class BaseTelem(object, metaclass=abc.ABCMeta):
self._log_telem_sending(serialized_data, log_data)
ControlClient.send_telemetry(self.telem_category, serialized_data)
@abc.abstractmethod
def get_data(self) -> dict:
"""
:return: Data of telemetry (should be dict)
"""
pass
@property
def json_encoder(self):
return json.JSONEncoder
@ -57,14 +48,6 @@ class BaseTelem(object, metaclass=abc.ABCMeta):
if log_data:
logger.debug(f"Telemetry contents: {BaseTelem._truncate_data(serialized_data)}")
@property
@abc.abstractmethod
def telem_category(self):
"""
:return: Telemetry type
"""
pass
@staticmethod
def _truncate_data(data: str):
if len(data) <= LOGGED_DATA_LENGTH:

View File

@ -0,0 +1,29 @@
import abc
class ITelem(metaclass=abc.ABCMeta):
@abc.abstractmethod
def send(self, log_data=True):
"""
Sends telemetry to island
"""
@abc.abstractmethod
def get_data(self) -> dict:
"""
:return: Data of telemetry (should be dict)
"""
pass
@property
@abc.abstractmethod
def json_encoder(self):
pass
@property
@abc.abstractmethod
def telem_category(self):
"""
:return: Telemetry type
"""
pass