Merge pull request #1642 from guardicore/1538-reduce-network-chatter

1538 reduce network chatter
This commit is contained in:
Mike Salvatore 2021-12-08 08:10:04 -05:00 committed by GitHub
commit abf274a8d3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
30 changed files with 418 additions and 218 deletions

View File

@ -1,7 +1,6 @@
import json
import logging
import platform
from datetime import datetime
from pprint import pformat
from socket import gethostname
from urllib.parse import urljoin
@ -12,12 +11,12 @@ from requests.exceptions import ConnectionError
import infection_monkey.monkeyfs as monkeyfs
import infection_monkey.tunnel as tunnel
from common.common_consts.api_url_consts import T1216_PBA_FILE_DOWNLOAD_PATH
from common.common_consts.time_formats import DEFAULT_TIME_FORMAT
from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT
from infection_monkey.config import GUID, WormConfiguration
from infection_monkey.network.info import local_ips
from infection_monkey.transport.http import HTTPConnectProxy
from infection_monkey.transport.tcp import TcpProxy
from infection_monkey.utils import agent_process
from infection_monkey.utils.environment import is_windows_os
requests.packages.urllib3.disable_warnings()
@ -52,7 +51,7 @@ class ControlClient(object):
"description": " ".join(platform.uname()),
"config": WormConfiguration.as_dict(),
"parent": parent,
"launch_time": str(datetime.now().strftime(DEFAULT_TIME_FORMAT)),
"launch_time": agent_process.get_start_time(),
}
if ControlClient.proxies:

View File

@ -19,9 +19,12 @@ class ControlChannel(IControlChannel):
self._control_channel_server = server
def should_agent_stop(self) -> bool:
if not self._control_channel_server:
logger.error("Agent should stop because it can't connect to the C&C server.")
return True
try:
response = requests.get( # noqa: DUO123
f"{self._control_channel_server}/api/monkey_control/{self._agent_id}",
f"https://{self._control_channel_server}/api/monkey_control/needs-to-stop/{self._agent_id}",
verify=False,
proxies=ControlClient.proxies,
timeout=SHORT_REQUEST_TIMEOUT,

View File

@ -8,8 +8,9 @@ import time
import infection_monkey.tunnel as tunnel
from common.utils.attack_utils import ScanStatus, UsageEnum
from common.version import get_version
from infection_monkey.config import WormConfiguration
from infection_monkey.config import GUID, WormConfiguration
from infection_monkey.control import ControlClient
from infection_monkey.master.control_channel import ControlChannel
from infection_monkey.master.mock_master import MockMaster
from infection_monkey.model import DELAY_DELETE_CMD
from infection_monkey.network.firewall import app as firewall
@ -73,8 +74,9 @@ class InfectionMonkey:
if is_windows_os():
T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
if InfectionMonkey._is_monkey_alive_by_config():
logger.info("Monkey marked 'not alive' from configuration.")
should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
if should_stop:
logger.info("The Monkey Island has instructed this agent to stop")
return
if InfectionMonkey._is_upgrade_to_64_needed():
@ -126,10 +128,6 @@ class InfectionMonkey:
logger.debug("default server set to: %s" % self._opts.server)
return True
@staticmethod
def _is_monkey_alive_by_config():
return not WormConfiguration.alive
@staticmethod
def _is_upgrade_to_64_needed():
return WindowsUpgrader.should_upgrade()

View File

@ -0,0 +1,8 @@
import os
import psutil
def get_start_time() -> float:
agent_process = psutil.Process(os.getpid())
return agent_process.create_time()

View File

@ -8,6 +8,11 @@ from werkzeug.exceptions import NotFound
from common.common_consts.api_url_consts import T1216_PBA_FILE_DOWNLOAD_PATH
from monkey_island.cc.database import database, mongo
from monkey_island.cc.resources.agent_controls import (
StartedOnIsland,
StopAgentCheck,
StopAllAgents,
)
from monkey_island.cc.resources.attack.attack_report import AttackReport
from monkey_island.cc.resources.auth.auth import Authenticate, init_jwt
from monkey_island.cc.resources.auth.registration import Registration
@ -30,8 +35,6 @@ from monkey_island.cc.resources.island_mode import IslandMode
from monkey_island.cc.resources.local_run import LocalRun
from monkey_island.cc.resources.log import Log
from monkey_island.cc.resources.monkey import Monkey
from monkey_island.cc.resources.monkey_control.started_on_island import StartedOnIsland
from monkey_island.cc.resources.monkey_control.stop_agent_check import StopAgentCheck
from monkey_island.cc.resources.monkey_download import MonkeyDownload
from monkey_island.cc.resources.netmap import NetMap
from monkey_island.cc.resources.node import Node
@ -97,6 +100,7 @@ def init_app_config(app, mongo_url):
# See https://flask.palletsprojects.com/en/1.1.x/config/#JSON_SORT_KEYS.
app.config["JSON_SORT_KEYS"] = False
app.url_map.strict_slashes = False
app.json_encoder = CustomJSONEncoder
@ -124,30 +128,26 @@ def init_api_resources(api):
api.add_resource(
Monkey,
"/api/monkey",
"/api/monkey/",
"/api/monkey/<string:guid>",
"/api/monkey/<string:guid>/<string:config_format>",
)
api.add_resource(Bootloader, "/api/bootloader/<string:os>")
api.add_resource(LocalRun, "/api/local-monkey", "/api/local-monkey/")
api.add_resource(ClientRun, "/api/client-monkey", "/api/client-monkey/")
api.add_resource(
Telemetry, "/api/telemetry", "/api/telemetry/", "/api/telemetry/<string:monkey_guid>"
)
api.add_resource(LocalRun, "/api/local-monkey")
api.add_resource(ClientRun, "/api/client-monkey")
api.add_resource(Telemetry, "/api/telemetry", "/api/telemetry/<string:monkey_guid>")
api.add_resource(IslandMode, "/api/island-mode")
api.add_resource(IslandConfiguration, "/api/configuration/island", "/api/configuration/island/")
api.add_resource(IslandConfiguration, "/api/configuration/island")
api.add_resource(ConfigurationExport, "/api/configuration/export")
api.add_resource(ConfigurationImport, "/api/configuration/import")
api.add_resource(
MonkeyDownload,
"/api/monkey/download",
"/api/monkey/download/",
"/api/monkey/download/<string:path>",
)
api.add_resource(NetMap, "/api/netmap", "/api/netmap/")
api.add_resource(Edge, "/api/netmap/edge", "/api/netmap/edge/")
api.add_resource(Node, "/api/netmap/node", "/api/netmap/node/")
api.add_resource(NetMap, "/api/netmap")
api.add_resource(Edge, "/api/netmap/edge")
api.add_resource(Node, "/api/netmap/node")
api.add_resource(NodeStates, "/api/netmap/nodeStates")
api.add_resource(SecurityReport, "/api/report/security")
@ -158,9 +158,9 @@ def init_api_resources(api):
api.add_resource(MonkeyExploitation, "/api/exploitations/monkey")
api.add_resource(ZeroTrustFindingEvent, "/api/zero-trust/finding-event/<string:finding_id>")
api.add_resource(TelemetryFeed, "/api/telemetry-feed", "/api/telemetry-feed/")
api.add_resource(Log, "/api/log", "/api/log/")
api.add_resource(IslandLog, "/api/log/island/download", "/api/log/island/download/")
api.add_resource(TelemetryFeed, "/api/telemetry-feed")
api.add_resource(Log, "/api/log")
api.add_resource(IslandLog, "/api/log/island/download")
api.add_resource(PBAFileDownload, "/api/pba/download/<string:filename>")
api.add_resource(T1216PBAFileDownload, T1216_PBA_FILE_DOWNLOAD_PATH)
api.add_resource(
@ -170,10 +170,11 @@ def init_api_resources(api):
"/api/fileUpload/<string:file_type>?restore=<string:filename>",
)
api.add_resource(PropagationCredentials, "/api/propagation-credentials/<string:guid>")
api.add_resource(RemoteRun, "/api/remote-monkey", "/api/remote-monkey/")
api.add_resource(VersionUpdate, "/api/version-update", "/api/version-update/")
api.add_resource(RemoteRun, "/api/remote-monkey")
api.add_resource(VersionUpdate, "/api/version-update")
api.add_resource(StartedOnIsland, "/api/monkey_control/started_on_island")
api.add_resource(StopAgentCheck, "/api/monkey_control/<int:monkey_guid>")
api.add_resource(StopAgentCheck, "/api/monkey_control/needs-to-stop/<int:monkey_guid>")
api.add_resource(StopAllAgents, "/api/monkey_control/stop-all-agents")
api.add_resource(ScoutSuiteAuth, "/api/scoutsuite_auth/<string:provider>")
api.add_resource(AWSKeys, "/api/aws_keys")

View File

@ -0,0 +1 @@
from .agent_controls import AgentControls

View File

@ -0,0 +1,7 @@
from mongoengine import Document, FloatField
class AgentControls(Document):
# Timestamp of the last "kill all agents" command
last_stop_all = FloatField(default=None)

View File

@ -1,4 +1,4 @@
from mongoengine import EmbeddedDocument
from mongoengine import EmbeddedDocument, BooleanField
class Config(EmbeddedDocument):
@ -8,5 +8,6 @@ class Config(EmbeddedDocument):
See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist
"""
alive = BooleanField()
meta = {"strict": False}
pass

View File

@ -9,6 +9,7 @@ from mongoengine import (
DoesNotExist,
DynamicField,
EmbeddedDocumentField,
FloatField,
ListField,
ReferenceField,
StringField,
@ -20,6 +21,10 @@ from monkey_island.cc.server_utils.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURAT
from monkey_island.cc.services.utils.network_utils import local_ip_addresses
class ParentNotFoundError(Exception):
"""Raise when trying to get a parent of monkey that doesn't have one"""
class Monkey(Document):
"""
This class has 2 main section:
@ -38,7 +43,7 @@ class Monkey(Document):
description = StringField()
hostname = StringField()
ip_addresses = ListField(StringField())
launch_time = StringField()
launch_time = FloatField()
keepalive = DateTimeField()
modifytime = DateTimeField()
# TODO make "parent" an embedded document, so this can be removed and the schema explained (
@ -95,6 +100,18 @@ class Monkey(Document):
monkey_is_dead = True
return monkey_is_dead
def has_parent(self):
for p in self.parent:
if p[0] != self.guid:
return True
return False
def get_parent(self):
if self.has_parent():
return Monkey.objects(guid=self.parent[0][0]).first()
else:
raise ParentNotFoundError(f"No parent was found for agent with GUID {self.guid}")
def get_os(self):
os = "unknown"
if self.description.lower().find("linux") != -1:

View File

@ -0,0 +1,3 @@
from .stop_all_agents import StopAllAgents
from .started_on_island import StartedOnIsland
from .stop_agent_check import StopAgentCheck

View File

@ -0,0 +1,8 @@
import flask_restful
from monkey_island.cc.services.infection_lifecycle import should_agent_die
class StopAgentCheck(flask_restful.Resource):
def get(self, monkey_guid: int):
return {"stop_agent": should_agent_die(monkey_guid)}

View File

@ -0,0 +1,23 @@
import json
import flask_restful
from flask import make_response, request
from monkey_island.cc.resources.auth.auth import jwt_required
from monkey_island.cc.resources.utils.semaphores import agent_killing_mutex
from monkey_island.cc.services.infection_lifecycle import set_stop_all, should_agent_die
class StopAllAgents(flask_restful.Resource):
@jwt_required
def post(self):
with agent_killing_mutex:
data = json.loads(request.data)
if data["kill_time"]:
set_stop_all(data["kill_time"])
return make_response({}, 200)
else:
return make_response({}, 400)
def get(self, monkey_guid):
return {"stop_agent": should_agent_die(monkey_guid)}

View File

@ -8,6 +8,7 @@ from flask import request
from monkey_island.cc.database import mongo
from monkey_island.cc.models.monkey_ttl import create_monkey_ttl_document
from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore
from monkey_island.cc.resources.utils.semaphores import agent_killing_mutex
from monkey_island.cc.server_utils.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
from monkey_island.cc.services.config import ConfigService
from monkey_island.cc.services.edge.edge import EdgeService
@ -66,6 +67,7 @@ class Monkey(flask_restful.Resource):
# Called on monkey wakeup to initialize local configuration
@TestTelemStore.store_exported_telem
def post(self, **kw):
with agent_killing_mutex:
monkey_json = json.loads(request.data)
monkey_json["creds"] = []
monkey_json["dead"] = False
@ -108,7 +110,9 @@ class Monkey(flask_restful.Resource):
)
else:
parent_to_add = (parent, None)
elif (not parent or parent == monkey_json.get("guid")) and "ip_addresses" in monkey_json:
elif (
not parent or parent == monkey_json.get("guid")
) and "ip_addresses" in monkey_json:
exploit_telem = [
x
for x in mongo.db.telemetry.find(
@ -139,7 +143,9 @@ class Monkey(flask_restful.Resource):
ttl = create_monkey_ttl_document(DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS)
monkey_json["ttl_ref"] = ttl.id
mongo.db.monkey.update({"guid": monkey_json["guid"]}, {"$set": monkey_json}, upsert=True)
mongo.db.monkey.update(
{"guid": monkey_json["guid"]}, {"$set": monkey_json}, upsert=True
)
# Merge existing scanned node with new monkey
@ -154,7 +160,9 @@ class Monkey(flask_restful.Resource):
if existing_node:
node_id = existing_node["_id"]
EdgeService.update_all_dst_nodes(old_dst_node_id=node_id, new_dst_node_id=new_monkey_id)
EdgeService.update_all_dst_nodes(
old_dst_node_id=node_id, new_dst_node_id=new_monkey_id
)
for creds in existing_node["creds"]:
NodeService.add_credentials_to_monkey(new_monkey_id, creds)
mongo.db.node.remove({"_id": node_id})

View File

@ -1,9 +0,0 @@
import flask_restful
class StopAgentCheck(flask_restful.Resource):
def get(self, monkey_guid: int):
if monkey_guid % 2:
return {"stop_agent": True}
else:
return {"stop_agent": False}

View File

@ -6,7 +6,7 @@ from flask import jsonify, make_response, request
from monkey_island.cc.database import mongo
from monkey_island.cc.resources.auth.auth import jwt_required
from monkey_island.cc.services.database import Database
from monkey_island.cc.services.infection_lifecycle import InfectionLifecycle
from monkey_island.cc.services.infection_lifecycle import get_completed_steps
from monkey_island.cc.services.utils.network_utils import local_ip_addresses
logger = logging.getLogger(__name__)
@ -21,8 +21,6 @@ class Root(flask_restful.Resource):
return self.get_server_info()
elif action == "reset":
return jwt_required(Database.reset_db)()
elif action == "killall":
return jwt_required(InfectionLifecycle.kill_all)()
elif action == "is-up":
return {"is-up": True}
else:
@ -33,5 +31,5 @@ class Root(flask_restful.Resource):
return jsonify(
ip_addresses=local_ip_addresses(),
mongo=str(mongo.db),
completed_steps=InfectionLifecycle.get_completed_steps(),
completed_steps=get_completed_steps(),
)

View File

@ -0,0 +1,5 @@
from gevent.lock import BoundedSemaphore
# Mutex avoids race condition between monkeys
# being marked dead and monkey waking up as alive
agent_killing_mutex = BoundedSemaphore()

View File

@ -3,6 +3,7 @@ import logging
from flask import jsonify
from monkey_island.cc.database import mongo
from monkey_island.cc.models.agent_controls import AgentControls
from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations
from monkey_island.cc.services.config import ConfigService
@ -23,6 +24,7 @@ class Database(object):
if not x.startswith("system.") and not x == AttackMitigations.COLLECTION_NAME
]
ConfigService.init_config()
Database.init_agent_controls()
logger.info("DB was reset")
return jsonify(status="OK")
@ -31,6 +33,10 @@ class Database(object):
mongo.db[collection_name].drop()
logger.info("Dropped collection {}".format(collection_name))
@staticmethod
def init_agent_controls():
AgentControls().save()
@staticmethod
def is_mitigations_missing() -> bool:
return bool(AttackMitigations.COLLECTION_NAME not in mongo.db.list_collection_names())

View File

@ -1,9 +1,7 @@
import logging
from datetime import datetime
from flask import jsonify
from monkey_island.cc.database import mongo
from monkey_island.cc.models import Monkey
from monkey_island.cc.models.agent_controls import AgentControls
from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore
from monkey_island.cc.services.config import ConfigService
from monkey_island.cc.services.node import NodeService
@ -16,25 +14,39 @@ from monkey_island.cc.services.reporting.report_generation_synchronisation impor
logger = logging.getLogger(__name__)
class InfectionLifecycle:
@staticmethod
def kill_all():
mongo.db.monkey.update(
{"dead": False},
{"$set": {"config.alive": False, "modifytime": datetime.now()}},
upsert=False,
multi=True,
)
logger.info("Kill all monkeys was called")
return jsonify(status="OK")
def set_stop_all(time: float):
for monkey in Monkey.objects():
monkey.config.alive = False
monkey.save()
agent_controls = AgentControls.objects.first()
agent_controls.last_stop_all = time
agent_controls.save()
def should_agent_die(guid: int) -> bool:
monkey = Monkey.objects(guid=str(guid)).first()
return _is_monkey_marked_dead(monkey) or _is_monkey_killed_manually(monkey)
def _is_monkey_marked_dead(monkey: Monkey) -> bool:
return not monkey.config.alive
def _is_monkey_killed_manually(monkey: Monkey) -> bool:
if monkey.has_parent():
launch_timestamp = monkey.get_parent().launch_time
else:
launch_timestamp = monkey.launch_time
kill_timestamp = AgentControls.objects.first().last_stop_all
return int(kill_timestamp) >= int(launch_timestamp)
@staticmethod
def get_completed_steps():
is_any_exists = NodeService.is_any_monkey_exists()
infection_done = NodeService.is_monkey_finished_running()
if infection_done:
InfectionLifecycle._on_finished_infection()
_on_finished_infection()
report_done = ReportService.is_report_generated()
else: # Infection is not done
report_done = False
@ -46,7 +58,7 @@ class InfectionLifecycle:
report_done=report_done,
)
@staticmethod
def _on_finished_infection():
# Checking is_report_being_generated here, because we don't want to wait to generate a
# report; rather,

View File

@ -3,6 +3,7 @@ from typing import List
from monkey_island.cc.database import mongo
from monkey_island.cc.services.node import NodeService
from monkey_island.cc.services.utils.formatting import timestamp_to_date
@dataclass
@ -27,5 +28,5 @@ def monkey_to_manual_exploitation(monkey: dict) -> ManualExploitation:
return ManualExploitation(
hostname=monkey["hostname"],
ip_addresses=monkey["ip_addresses"],
start_time=monkey["launch_time"],
start_time=timestamp_to_date(monkey["launch_time"]),
)

View File

@ -0,0 +1,7 @@
from datetime import datetime
from common.common_consts.time_formats import DEFAULT_TIME_FORMAT
def timestamp_to_date(timestamp: int) -> str:
return datetime.fromtimestamp(timestamp).strftime(DEFAULT_TIME_FORMAT)

View File

@ -1,6 +1,6 @@
{
"name": "infection-monkey",
"version": "1.11.0",
"version": "1.12.0",
"lockfileVersion": 1,
"requires": true,
"dependencies": {
@ -11270,14 +11270,6 @@
"prop-types": "^15.7.2"
}
},
"react-toggle": {
"version": "4.1.2",
"resolved": "https://registry.npmjs.org/react-toggle/-/react-toggle-4.1.2.tgz",
"integrity": "sha512-4Ohw31TuYQdhWfA6qlKafeXx3IOH7t4ZHhmRdwsm1fQREwOBGxJT+I22sgHqR/w8JRdk+AeMCJXPImEFSrNXow==",
"requires": {
"classnames": "^2.2.5"
}
},
"react-tooltip-lite": {
"version": "1.12.0",
"resolved": "https://registry.npmjs.org/react-tooltip-lite/-/react-tooltip-lite-1.12.0.tgz",

View File

@ -111,7 +111,6 @@
"react-router-dom": "^5.3.0",
"react-spinners": "^0.9.0",
"react-table": "^6.10.3",
"react-toggle": "^4.1.2",
"react-tooltip-lite": "^1.12.0",
"redux": "^4.1.1",
"sha3": "^2.1.4",

View File

@ -20,7 +20,6 @@ import GettingStartedPage from './pages/GettingStartedPage';
import 'normalize.css/normalize.css';
import 'styles/App.css';
import 'react-toggle/style.css';
import 'react-table/react-table.css';
import LoadingScreen from './ui-components/LoadingScreen';
import SidebarLayoutComponent from "./layouts/SidebarLayoutComponent";

View File

@ -2,7 +2,6 @@ import React from 'react';
import {FontAwesomeIcon} from '@fortawesome/react-fontawesome'
import {faHandPointLeft} from '@fortawesome/free-solid-svg-icons/faHandPointLeft'
import {faQuestionCircle} from '@fortawesome/free-solid-svg-icons/faQuestionCircle'
import Toggle from 'react-toggle';
import {OverlayTrigger, Tooltip} from 'react-bootstrap';
import download from 'downloadjs'
import AuthComponent from '../../AuthComponent';
@ -67,32 +66,6 @@ class PreviewPaneComponent extends AuthComponent {
);
}
forceKill(event, asset) {
let newConfig = asset.config;
newConfig['alive'] = !event.target.checked;
this.authFetch('/api/monkey/' + asset.guid,
{
method: 'PATCH',
headers: {'Content-Type': 'application/json'},
body: JSON.stringify({config: newConfig})
});
}
forceKillRow(asset) {
return (
<tr>
<th>
Force Kill&nbsp;
{this.generateToolTip('If this is on, monkey will die next time it communicates')}
</th>
<td>
<Toggle id={asset.id} checked={!asset.config.alive} icons={false} disabled={asset.dead}
onChange={(e) => this.forceKill(e, asset)}/>
</td>
</tr>
);
}
unescapeLog(st) {
return st.substr(1, st.length - 2) // remove quotation marks on beginning and end of string.
@ -193,7 +166,6 @@ class PreviewPaneComponent extends AuthComponent {
{this.ipsRow(asset)}
{this.servicesRow(asset)}
{this.accessibleRow(asset)}
{this.forceKillRow(asset)}
{this.downloadLogRow(asset)}
</tbody>
</table>

View File

@ -84,9 +84,14 @@ class MapPageComponent extends AuthComponent {
}
killAllMonkeys = () => {
this.authFetch('/api?action=killall')
this.authFetch('/api/monkey_control/stop-all-agents',
{
method: 'POST',
headers: {'Content-Type': 'application/json'},
body: JSON.stringify({kill_time: Date.now()})
})
.then(res => res.json())
.then(res => this.setState({killPressed: (res.status === 'OK')}));
.then(res => {this.setState({killPressed: true}); console.log(res)});
};
renderKillDialogModal = () => {

View File

@ -3,7 +3,7 @@ import uuid
import pytest
from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError
from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError, ParentNotFoundError
from monkey_island.cc.models.monkey_ttl import MonkeyTtl
logger = logging.getLogger(__name__)
@ -162,3 +162,35 @@ class TestMonkey:
cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info()
assert cache_info_after_query.hits == 2
@pytest.mark.usefixtures("uses_database")
def test_has_parent(self):
monkey_1 = Monkey(guid=str(uuid.uuid4()))
monkey_2 = Monkey(guid=str(uuid.uuid4()))
monkey_1.parent = [[monkey_2.guid]]
monkey_1.save()
assert monkey_1.has_parent()
@pytest.mark.usefixtures("uses_database")
def test_has_no_parent(self):
monkey_1 = Monkey(guid=str(uuid.uuid4()))
monkey_1.parent = [[monkey_1.guid]]
monkey_1.save()
assert not monkey_1.has_parent()
@pytest.mark.usefixtures("uses_database")
def test_get_parent(self):
monkey_1 = Monkey(guid=str(uuid.uuid4()))
monkey_2 = Monkey(guid=str(uuid.uuid4()))
monkey_1.parent = [[monkey_2.guid]]
monkey_1.save()
monkey_2.save()
assert monkey_1.get_parent().guid == monkey_2.guid
@pytest.mark.usefixtures("uses_database")
def test_get_parent_no_parent(self):
monkey_1 = Monkey(guid=str(uuid.uuid4()))
monkey_1.parent = [[monkey_1.guid]]
monkey_1.save()
with pytest.raises(ParentNotFoundError):
monkey_1.get_parent()

View File

@ -0,0 +1,101 @@
import uuid
import pytest
from monkey_island.cc.models import Config, Monkey
from monkey_island.cc.models.agent_controls import AgentControls
from monkey_island.cc.services.infection_lifecycle import should_agent_die
@pytest.mark.usefixtures("uses_database")
def test_should_agent_die_by_config(monkeypatch):
monkey = Monkey(guid=str(uuid.uuid4()))
monkey.config = Config(alive=False)
monkey.save()
assert should_agent_die(monkey.guid)
monkeypatch.setattr(
"monkey_island.cc.services.infection_lifecycle._is_monkey_killed_manually", lambda _: False
)
monkey.config.alive = True
monkey.save()
assert not should_agent_die(monkey.guid)
def create_monkey(launch_time):
monkey = Monkey(guid=str(uuid.uuid4()))
monkey.config = Config(alive=True)
monkey.launch_time = launch_time
monkey.save()
return monkey
def create_kill_event(event_time):
kill_event = AgentControls(last_stop_all=event_time)
kill_event.save()
return kill_event
def create_parent(child_monkey, launch_time):
monkey_parent = Monkey(guid=str(uuid.uuid4()))
child_monkey.parent = [[monkey_parent.guid]]
monkey_parent.launch_time = launch_time
monkey_parent.save()
child_monkey.save()
@pytest.mark.usefixtures("uses_database")
def test_was_agent_killed_manually(monkeypatch):
monkey = create_monkey(launch_time=2)
create_kill_event(event_time=3)
assert should_agent_die(monkey.guid)
@pytest.mark.usefixtures("uses_database")
def test_agent_killed_on_wakeup(monkeypatch):
monkey = create_monkey(launch_time=2)
create_kill_event(event_time=2)
assert should_agent_die(monkey.guid)
@pytest.mark.usefixtures("uses_database")
def test_manual_kill_dont_affect_new_monkeys(monkeypatch):
monkey = create_monkey(launch_time=3)
create_kill_event(event_time=2)
assert not should_agent_die(monkey.guid)
@pytest.mark.usefixtures("uses_database")
def test_parent_manually_killed(monkeypatch):
monkey = create_monkey(launch_time=3)
create_parent(child_monkey=monkey, launch_time=1)
create_kill_event(event_time=2)
assert should_agent_die(monkey.guid)
@pytest.mark.usefixtures("uses_database")
def test_parent_manually_killed_on_wakeup(monkeypatch):
monkey = create_monkey(launch_time=3)
create_parent(child_monkey=monkey, launch_time=2)
create_kill_event(event_time=2)
assert should_agent_die(monkey.guid)
@pytest.mark.usefixtures("uses_database")
def test_manual_kill_dont_affect_new_monkeys_with_parent(monkeypatch):
monkey = create_monkey(launch_time=3)
create_parent(child_monkey=monkey, launch_time=2)
create_kill_event(event_time=1)
assert not should_agent_die(monkey.guid)

View File

@ -3,6 +3,7 @@ Everything in this file is what Vulture found as dead code but either isn't real
dead or is kept deliberately. Referencing these in a file like this makes sure that
Vulture doesn't mark these as dead again.
"""
from monkey_island.cc import app
from monkey_island.cc.models import Report
fake_monkey_dir_path # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
@ -100,6 +101,8 @@ EnvironmentCollector # unused class (monkey/infection_monkey/system_info/collec
ProcessListCollector # unused class (monkey/infection_monkey/system_info/collectors/process_list_collector.py:18)
_.coinit_flags # unused attribute (monkey/infection_monkey/system_info/windows_info_collector.py:11)
_.representations # unused attribute (monkey/monkey_island/cc/app.py:180)
_.representations # unused attribute (monkey/monkey_island/cc/app.py:180)
app.url_map.strict_slashes
_.log_message # unused method (monkey/infection_monkey/transport/http.py:188)
_.log_message # unused method (monkey/infection_monkey/transport/http.py:109)
_.version_string # unused method (monkey/infection_monkey/transport/http.py:148)