#2059: moved connection handling from Service to IOSoftware + changes that now utilise connections from IOSoftware + dos bot attacking now works + tests

This commit is contained in:
Czar Echavez
2023-12-12 17:01:03 +00:00
parent 1b124035a1
commit 4f79d2ad36
12 changed files with 510 additions and 88 deletions

View File

@@ -5,7 +5,7 @@ from uuid import uuid4
from primaite import getLogger
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application, ApplicationOperatingState
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.core.software_manager import SoftwareManager
_LOGGER = getLogger(__name__)
@@ -23,7 +23,6 @@ class DatabaseClient(Application):
server_ip_address: Optional[IPv4Address] = None
server_password: Optional[str] = None
connections: Dict[str, Dict] = {}
_query_success_tracker: Dict[str, bool] = {}
def __init__(self, **kwargs):
@@ -143,7 +142,7 @@ class DatabaseClient(Application):
dest_ip_address=self.server_ip_address,
dest_port=self.port,
)
self.connections.pop(connection_id)
self.remove_connection(connection_id=connection_id)
self.sys_log.info(
f"{self.name}: DatabaseClient disconnected connection {connection_id} from {self.server_ip_address}"
@@ -181,8 +180,6 @@ class DatabaseClient(Application):
def run(self) -> None:
"""Run the DatabaseClient."""
super().run()
if self.operating_state == ApplicationOperatingState.RUNNING:
self.connect()
def query(self, sql: str, connection_id: Optional[str] = None) -> bool:
"""
@@ -221,7 +218,8 @@ class DatabaseClient(Application):
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "connect_response":
if payload["response"] is True:
self.connections[payload.get("connection_id")] = payload
# add connection
self.add_connection(connection_id=payload.get("connection_id"), session_id=session_id)
elif payload["type"] == "sql":
query_id = payload.get("uuid")
status_code = payload.get("status_code")

View File

@@ -5,7 +5,6 @@ from typing import Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
_LOGGER = getLogger(__name__)
@@ -177,9 +176,9 @@ class DataManipulationBot(DatabaseClient):
This is the core loop where the bot sequentially goes through the stages of the attack.
"""
if self.operating_state != ApplicationOperatingState.RUNNING:
if not self._can_perform_action():
return
if self.server_ip_address and self.payload and self.operating_state:
if self.server_ip_address and self.payload:
self.sys_log.info(f"{self.name}: Running")
self._logon()
self._perform_port_scan(p_of_success=self.port_scan_p_of_success)

View File

@@ -0,0 +1,184 @@
from enum import IntEnum
from ipaddress import IPv4Address
from typing import Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient
_LOGGER = getLogger(__name__)
class DoSAttackStage(IntEnum):
"""Enum representing the different stages of a Denial of Service attack."""
NOT_STARTED = 0
"Attack not yet started."
PORT_SCAN = 1
"Attack is in discovery stage - checking if provided ip and port are open."
ATTACKING = 2
"Denial of Service attack is in progress."
COMPLETED = 3
"Attack is completed."
class DoSBot(DatabaseClient, Application):
"""A bot that simulates a Denial of Service attack."""
target_ip_address: Optional[IPv4Address] = None
"""IP address of the target service."""
target_port: Optional[Port] = None
"""Port of the target service."""
payload: Optional[str] = None
"""Payload to deliver to the target service as part of the denial of service attack."""
repeat: bool = False
"""If true, the Denial of Service bot will keep performing the attack."""
attack_stage: DoSAttackStage = DoSAttackStage.NOT_STARTED
"""Current stage of the DoS kill chain."""
port_scan_p_of_success: float = 0.1
"""Probability of port scanning being sucessful."""
dos_intensity: float = 0.25
"""How much of the max sessions will be used by the DoS when attacking."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DoSBot"
self.max_sessions = 1000 # override normal max sessions
def set_original_state(self):
"""Set the original state of the Denial of Service Bot."""
_LOGGER.debug(f"Setting {self.name} original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {
"target_ip_address",
"target_port",
"payload",
"repeat",
"attack_stage",
"max_sessions",
"port_scan_p_of_success",
"dos_intensity",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting {self.name} state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
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()))
return rm
def configure(
self,
target_ip_address: IPv4Address,
target_port: Optional[Port] = Port.POSTGRES_SERVER,
payload: Optional[str] = None,
repeat: bool = False,
max_sessions: int = 1000,
):
"""
Configure the Denial of Service bot.
:param: target_ip_address: The IP address of the Node containing the target service.
:param: target_port: The port of the target service. Optional - Default is `Port.HTTP`
:param: payload: The payload the DoS Bot will throw at the target service. Optional - Default is `None`
:param: repeat: If True, the bot will maintain the attack. Optional - Default is `True`
:param: max_sessions: The maximum number of sessions the DoS bot will attack with. Optional - Default is 1000
"""
self.target_ip_address = target_ip_address
self.target_port = target_port
self.payload = payload
self.repeat = repeat
self.max_sessions = max_sessions
self.sys_log.info(
f"{self.name}: Configured the {self.name} with {target_ip_address=}, {target_port=}, {payload=}, {repeat=}."
)
def run(self):
"""Run the Denial of Service Bot."""
super().run()
self._application_loop()
def _application_loop(self):
"""
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
# 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
self.clear_connections()
self._perform_port_scan(p_of_success=self.port_scan_p_of_success)
self._perform_dos()
if self.repeat and self.attack_stage is DoSAttackStage.ATTACKING:
self.attack_stage = DoSAttackStage.NOT_STARTED
else:
self.attack_stage = DoSAttackStage.COMPLETED
def _perform_port_scan(self, p_of_success: Optional[float] = 0.1):
"""
Perform a simulated port scan to check for open SQL ports.
Advances the attack stage to `PORT_SCAN` if successful.
:param p_of_success: Probability of successful port scan, by default 0.1.
"""
if self.attack_stage == DoSAttackStage.NOT_STARTED:
# perform a port scan to identify that the SQL port is open on the server
if simulate_trial(p_of_success):
self.sys_log.info(f"{self.name}: Performing port scan")
# perform the port scan
port_is_open = True # Temporary; later we can implement NMAP port scan.
if port_is_open:
self.sys_log.info(f"{self.name}: ")
self.attack_stage = DoSAttackStage.PORT_SCAN
def _perform_dos(self):
"""
Perform the Denial of Service attack.
DoSBot does this by clogging up the available connections to a service.
"""
if not self.attack_stage == DoSAttackStage.PORT_SCAN:
return
self.attack_stage = DoSAttackStage.ATTACKING
self.server_ip_address = self.target_ip_address
self.port = self.target_port
dos_sessions = int(float(self.max_sessions) * self.dos_intensity)
for i in range(dos_sessions):
self.connect()
def apply_timestep(self, timestep: int) -> None:
"""
Apply a timestep to the bot, iterate through the application loop.
:param timestep: The timestep value to update the bot's state.
"""
self._application_loop()