island: Wrap services/post_breach_files.py functions in a static class

This commit is contained in:
Mike Salvatore 2021-05-11 07:47:40 -04:00
parent ba86ba0395
commit a7f2e023b8
2 changed files with 40 additions and 38 deletions

View File

@ -6,10 +6,10 @@ import logging
from jsonschema import Draft4Validator, validators
import monkey_island.cc.environment.environment_singleton as env_singleton
import monkey_island.cc.services.post_breach_files
from monkey_island.cc.database import mongo
from monkey_island.cc.server_utils.encryptor import get_encryptor
from monkey_island.cc.services.config_schema.config_schema import SCHEMA
from monkey_island.cc.services.post_breach_files import PostBreachFilesService
from monkey_island.cc.services.utils.network_utils import local_ip_addresses
__author__ = "itay.mizeretz"
@ -191,7 +191,7 @@ class ConfigService:
# PBA file upload happens on pba_file_upload endpoint and corresponding config options
# are set there
config_json = ConfigService._filter_none_values(config_json)
monkey_island.cc.services.post_breach_files.set_config_PBA_files(config_json)
PostBreachFilesService.set_config_PBA_files(config_json)
if should_encrypt:
try:
ConfigService.encrypt_config(config_json)
@ -229,7 +229,7 @@ class ConfigService:
@staticmethod
def reset_config():
monkey_island.cc.services.post_breach_files.remove_PBA_files()
PostBreachFilesService.remove_PBA_files()
config = ConfigService.get_default_config(True)
ConfigService.set_server_ips_in_config(config)
ConfigService.update_config(config, should_encrypt=False)

View File

@ -14,40 +14,42 @@ PBA_WINDOWS_FILENAME_PATH = ["monkey", "post_breach", "PBA_windows_filename"]
PBA_LINUX_FILENAME_PATH = ["monkey", "post_breach", "PBA_linux_filename"]
def remove_PBA_files():
if monkey_island.cc.services.config.ConfigService.get_config():
windows_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_WINDOWS_FILENAME_PATH
)
linux_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_LINUX_FILENAME_PATH
)
if linux_filename:
remove_file(linux_filename)
if windows_filename:
remove_file(windows_filename)
class PostBreachFilesService:
@staticmethod
def remove_PBA_files():
if monkey_island.cc.services.config.ConfigService.get_config():
windows_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_WINDOWS_FILENAME_PATH
)
linux_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_LINUX_FILENAME_PATH
)
if linux_filename:
PostBreachFilesService._remove_file(linux_filename)
if windows_filename:
PostBreachFilesService._remove_file(windows_filename)
@staticmethod
def _remove_file(file_name):
file_path = os.path.join(env_singleton.env.get_config().data_dir_abs_path, file_name)
try:
if os.path.exists(file_path):
os.remove(file_path)
except OSError as e:
logger.error("Can't remove previously uploaded post breach files: %s" % e)
def remove_file(file_name):
file_path = os.path.join(env_singleton.env.get_config().data_dir_abs_path, file_name)
try:
if os.path.exists(file_path):
os.remove(file_path)
except OSError as e:
logger.error("Can't remove previously uploaded post breach files: %s" % e)
def set_config_PBA_files(config_json):
"""
Sets PBA file info in config_json to current config's PBA file info values.
:param config_json: config_json that will be modified
"""
if monkey_island.cc.services.config.ConfigService.get_config():
linux_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_LINUX_FILENAME_PATH
)
windows_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_WINDOWS_FILENAME_PATH
)
config_json["monkey"]["post_breach"]["PBA_linux_filename"] = linux_filename
config_json["monkey"]["post_breach"]["PBA_windows_filename"] = windows_filename
@staticmethod
def set_config_PBA_files(config_json):
"""
Sets PBA file info in config_json to current config's PBA file info values.
:param config_json: config_json that will be modified
"""
if monkey_island.cc.services.config.ConfigService.get_config():
linux_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_LINUX_FILENAME_PATH
)
windows_filename = monkey_island.cc.services.config.ConfigService.get_config_value(
PBA_WINDOWS_FILENAME_PATH
)
config_json["monkey"]["post_breach"]["PBA_linux_filename"] = linux_filename
config_json["monkey"]["post_breach"]["PBA_windows_filename"] = windows_filename