oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make all requests return a RequestResponse

Browse Source
This commit is contained in:
Marek Wolan
2024-03-08 14:58:34 +00:00
parent d331224b45
commit beb51834f9
13 changed files with 275 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/primaite/simulator/system/applications/database_client.py
View File

@@ -3,6 +3,7 @@ from typing import Any, Dict, Optional
from uuid import uuid4
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
@@ -37,7 +38,7 @@ class DatabaseClient(Application):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request("execute", RequestType(func=lambda request, context: self.execute()))
rm.add_request("execute", RequestType(func=lambda request, context: RequestResponse.from_bool(self.execute())))
return rm
def execute(self) -> bool:

18
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -4,6 +4,7 @@ from typing import Dict, Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
@@ -76,7 +77,10 @@ class DataManipulationBot(Application):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.attack()))
rm.add_request(
name="execute",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.attack())),
)
return rm
@@ -179,21 +183,21 @@ class DataManipulationBot(Application):
"""
super().run()
def attack(self):
def attack(self) -> bool:
"""Perform the attack steps after opening the application."""
if not self._can_perform_action():
_LOGGER.debug("Data manipulation application attempted to execute but it cannot perform actions right now.")
self.run()
self._application_loop()
return self._application_loop()
def _application_loop(self):
def _application_loop(self) -> bool:
"""
The main application loop of the bot, handling the attack process.
This is the core loop where the bot sequentially goes through the stages of the attack.
"""
if not self._can_perform_action():
return
return False
if self.server_ip_address and self.payload:
self.sys_log.info(f"{self.name}: Running")
self._logon()
@@ -205,8 +209,12 @@ class DataManipulationBot(Application):
DataManipulationAttackStage.FAILED,
):
self.attack_stage = DataManipulationAttackStage.NOT_STARTED
return True
else:
self.sys_log.error(f"{self.name}: Failed to start as it requires both a target_ip_address and payload.")
return False
def apply_timestep(self, timestep: int) -> None:
"""

17
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -4,6 +4,7 @@ from typing import Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.database_client import DatabaseClient
@@ -59,7 +60,10 @@ class DoSBot(DatabaseClient):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run()))
rm.add_request(
name="execute",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.run())),
)
return rm
@@ -97,26 +101,26 @@ class DoSBot(DatabaseClient):
f"{repeat=}, {port_scan_p_of_success=}, {dos_intensity=}, {max_sessions=}."
)
def run(self):
def run(self) -> bool:
"""Run the Denial of Service Bot."""
super().run()
self._application_loop()
return self._application_loop()
def _application_loop(self):
def _application_loop(self) -> bool:
"""
The main application loop for the Denial of Service bot.
The loop goes through the stages of a DoS attack.
"""
if not self._can_perform_action():
return
return False
# DoS bot cannot do anything without a target
if not self.target_ip_address or not self.target_port:
self.sys_log.error(
f"{self.name} is not properly configured. {self.target_ip_address=}, {self.target_port=}"
)
return
return True
self.clear_connections()
self._perform_port_scan(p_of_success=self.port_scan_p_of_success)
@@ -126,6 +130,7 @@ class DoSBot(DatabaseClient):
self.attack_stage = DoSAttackStage.NOT_STARTED
else:
self.attack_stage = DoSAttackStage.COMPLETED
return True
def _perform_port_scan(self, p_of_success: Optional[float] = 0.1):
"""

6
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -6,6 +6,7 @@ from urllib.parse import urlparse
from pydantic import BaseModel, ConfigDict
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.http import (
HttpRequestMethod,
@@ -52,7 +53,10 @@ class WebBrowser(Application):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(
name="execute", request_type=RequestType(func=lambda request, context: self.get_webpage()) # noqa
name="execute",
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(self.get_webpage())
), # noqa
)
return rm

46
src/primaite/simulator/system/services/service.py
View File

@@ -3,6 +3,7 @@ from enum import Enum
from typing import Any, Dict, Optional
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.software import IOSoftware, SoftwareHealthState
@@ -80,14 +81,14 @@ class Service(IOSoftware):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request("scan", RequestType(func=lambda request, context: self.scan()))
rm.add_request("stop", RequestType(func=lambda request, context: self.stop()))
rm.add_request("start", RequestType(func=lambda request, context: self.start()))
rm.add_request("pause", RequestType(func=lambda request, context: self.pause()))
rm.add_request("resume", RequestType(func=lambda request, context: self.resume()))
rm.add_request("restart", RequestType(func=lambda request, context: self.restart()))
rm.add_request("disable", RequestType(func=lambda request, context: self.disable()))
rm.add_request("enable", RequestType(func=lambda request, context: self.enable()))
rm.add_request("scan", RequestType(func=lambda request, context: RequestResponse.from_bool(self.scan())))
rm.add_request("stop", RequestType(func=lambda request, context: RequestResponse.from_bool(self.stop())))
rm.add_request("start", RequestType(func=lambda request, context: RequestResponse.from_bool(self.start())))
rm.add_request("pause", RequestType(func=lambda request, context: RequestResponse.from_bool(self.pause())))
rm.add_request("resume", RequestType(func=lambda request, context: RequestResponse.from_bool(self.resume())))
rm.add_request("restart", RequestType(func=lambda request, context: RequestResponse.from_bool(self.restart())))
rm.add_request("disable", RequestType(func=lambda request, context: RequestResponse.from_bool(self.disable())))
rm.add_request("enable", RequestType(func=lambda request, context: RequestResponse.from_bool(self.enable())))
return rm
@abstractmethod
@@ -106,17 +107,19 @@ class Service(IOSoftware):
state["health_state_visible"] = self.health_state_visible.value
return state
def stop(self) -> None:
def stop(self) -> bool:
"""Stop the service."""
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.sys_log.info(f"Stopping service {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
return True
return False
def start(self, **kwargs) -> None:
def start(self, **kwargs) -> bool:
"""Start the service."""
# cant start the service if the node it is on is off
if not super()._can_perform_action():
return
return False
if self.operating_state == ServiceOperatingState.STOPPED:
self.sys_log.info(f"Starting service {self.name}")
@@ -124,36 +127,47 @@ class Service(IOSoftware):
# set software health state to GOOD if initially set to UNUSED
if self.health_state_actual == SoftwareHealthState.UNUSED:
self.set_health_state(SoftwareHealthState.GOOD)
return True
return False
def pause(self) -> None:
def pause(self) -> bool:
"""Pause the service."""
if self.operating_state == ServiceOperatingState.RUNNING:
self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.PAUSED
return True
return False
def resume(self) -> None:
def resume(self) -> bool:
"""Resume paused service."""
if self.operating_state == ServiceOperatingState.PAUSED:
self.sys_log.info(f"Resuming service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING
return True
return False
def restart(self) -> None:
def restart(self) -> bool:
"""Restart running service."""
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.RESTARTING
self.restart_countdown = self.restart_duration
return True
return False
def disable(self) -> None:
def disable(self) -> bool:
"""Disable the service."""
self.sys_log.info(f"Disabling Application {self.name}")
self.operating_state = ServiceOperatingState.DISABLED
return True
def enable(self) -> None:
def enable(self) -> bool:
"""Enable the disabled service."""
if self.operating_state == ServiceOperatingState.DISABLED:
self.sys_log.info(f"Enabling Application {self.name}")
self.operating_state = ServiceOperatingState.STOPPED
return True
return False
def apply_timestep(self, timestep: int) -> None:
"""

19
src/primaite/simulator/system/software.py
View File

@@ -5,6 +5,7 @@ from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, TYPE_CHECKING, Union
from primaite.interface.request import RequestResponse
from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
@@ -105,16 +106,18 @@ class Software(SimComponent):
rm.add_request(
"compromise",
RequestType(
func=lambda request, context: self.set_health_state(SoftwareHealthState.COMPROMISED),
func=lambda request, context: RequestResponse.from_bool(
self.set_health_state(SoftwareHealthState.COMPROMISED)
),
),
)
rm.add_request(
"patch",
RequestType(
func=lambda request, context: self.patch(),
func=lambda request, context: RequestResponse.from_bool(self.patch()),
),
)
rm.add_request("scan", RequestType(func=lambda request, context: self.scan()))
rm.add_request("scan", RequestType(func=lambda request, context: RequestResponse.from_bool(self.scan())))
return rm
def _get_session_details(self, session_id: str) -> Session:
@@ -148,7 +151,7 @@ class Software(SimComponent):
)
return state
def set_health_state(self, health_state: SoftwareHealthState) -> None:
def set_health_state(self, health_state: SoftwareHealthState) -> bool:
"""
Assign a new health state to this software.
@@ -160,6 +163,7 @@ class Software(SimComponent):
:type health_state: SoftwareHealthState
"""
self.health_state_actual = health_state
return True
def install(self) -> None:
"""
@@ -180,15 +184,18 @@ class Software(SimComponent):
"""
pass
def scan(self) -> None:
def scan(self) -> bool:
"""Update the observed health status to match the actual health status."""
self.health_state_visible = self.health_state_actual
return True
def patch(self) -> None:
def patch(self) -> bool:
"""Perform a patch on the software."""
if self.health_state_actual in (SoftwareHealthState.COMPROMISED, SoftwareHealthState.GOOD):
self._patching_countdown = self.patching_duration
self.set_health_state(SoftwareHealthState.PATCHING)
return True
return False
def _update_patch_status(self) -> None:
"""Update the patch status of the software."""