Make db manipulation bot work with db client
This commit is contained in:
@@ -409,9 +409,6 @@ class PrimaiteGame:
|
|||||||
# CREATE REWARD FUNCTION
|
# CREATE REWARD FUNCTION
|
||||||
reward_function = RewardFunction.from_config(reward_function_cfg)
|
reward_function = RewardFunction.from_config(reward_function_cfg)
|
||||||
|
|
||||||
# OTHER AGENT SETTINGS
|
|
||||||
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
|
|
||||||
|
|
||||||
# CREATE AGENT
|
# CREATE AGENT
|
||||||
if agent_type == "probabilistic_agent":
|
if agent_type == "probabilistic_agent":
|
||||||
# TODO: implement non-random agents and fix this parsing
|
# TODO: implement non-random agents and fix this parsing
|
||||||
@@ -424,6 +421,7 @@ class PrimaiteGame:
|
|||||||
settings=settings,
|
settings=settings,
|
||||||
)
|
)
|
||||||
elif agent_type == "ProxyAgent":
|
elif agent_type == "ProxyAgent":
|
||||||
|
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
|
||||||
new_agent = ProxyAgent(
|
new_agent = ProxyAgent(
|
||||||
agent_name=agent_cfg["ref"],
|
agent_name=agent_cfg["ref"],
|
||||||
action_space=action_space,
|
action_space=action_space,
|
||||||
@@ -433,6 +431,8 @@ class PrimaiteGame:
|
|||||||
)
|
)
|
||||||
game.rl_agents[agent_cfg["ref"]] = new_agent
|
game.rl_agents[agent_cfg["ref"]] = new_agent
|
||||||
elif agent_type == "RedDatabaseCorruptingAgent":
|
elif agent_type == "RedDatabaseCorruptingAgent":
|
||||||
|
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
|
||||||
|
|
||||||
new_agent = DataManipulationAgent(
|
new_agent = DataManipulationAgent(
|
||||||
agent_name=agent_cfg["ref"],
|
agent_name=agent_cfg["ref"],
|
||||||
action_space=action_space,
|
action_space=action_space,
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ class IPProtocol(Enum):
|
|||||||
.. _List of IPProtocols:
|
.. _List of IPProtocols:
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
NONE = "none"
|
||||||
|
"""Placeholder for a non-port."""
|
||||||
TCP = "tcp"
|
TCP = "tcp"
|
||||||
"""Transmission Control Protocol."""
|
"""Transmission Control Protocol."""
|
||||||
UDP = "udp"
|
UDP = "udp"
|
||||||
|
|||||||
@@ -85,6 +85,13 @@ class DatabaseClient(Application):
|
|||||||
return self.connected
|
return self.connected
|
||||||
|
|
||||||
def check_connection(self, connection_id: str) -> bool:
|
def check_connection(self, connection_id: str) -> bool:
|
||||||
|
"""Check whether the connection can be successfully re-established.
|
||||||
|
|
||||||
|
:param connection_id: connection ID to check
|
||||||
|
:type connection_id: str
|
||||||
|
:return: Whether the connection was successfully re-established.
|
||||||
|
:rtype: bool
|
||||||
|
"""
|
||||||
if not self._can_perform_action():
|
if not self._can_perform_action():
|
||||||
return False
|
return False
|
||||||
print(self.query("SELECT * FROM pg_stat_activity", connection_id=connection_id))
|
print(self.query("SELECT * FROM pg_stat_activity", connection_id=connection_id))
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
from enum import IntEnum
|
from enum import IntEnum
|
||||||
from ipaddress import IPv4Address
|
from ipaddress import IPv4Address
|
||||||
from typing import Optional
|
from typing import Dict, Optional
|
||||||
|
|
||||||
from primaite import getLogger
|
from primaite import getLogger
|
||||||
from primaite.game.science import simulate_trial
|
from primaite.game.science import simulate_trial
|
||||||
from primaite.simulator.core import RequestManager, RequestType
|
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
|
||||||
|
from primaite.simulator.system.applications.application import Application
|
||||||
from primaite.simulator.system.applications.database_client import DatabaseClient
|
from primaite.simulator.system.applications.database_client import DatabaseClient
|
||||||
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
@@ -32,12 +35,12 @@ class DataManipulationAttackStage(IntEnum):
|
|||||||
"Signifies that the attack has failed."
|
"Signifies that the attack has failed."
|
||||||
|
|
||||||
|
|
||||||
class DataManipulationBot(DatabaseClient):
|
class DataManipulationBot(Application):
|
||||||
"""A bot that simulates a script which performs a SQL injection attack."""
|
"""A bot that simulates a script which performs a SQL injection attack."""
|
||||||
|
|
||||||
server_ip_address: Optional[IPv4Address] = None
|
# server_ip_address: Optional[IPv4Address] = None
|
||||||
payload: Optional[str] = None
|
payload: Optional[str] = None
|
||||||
server_password: Optional[str] = None
|
# server_password: Optional[str] = None
|
||||||
port_scan_p_of_success: float = 0.1
|
port_scan_p_of_success: float = 0.1
|
||||||
data_manipulation_p_of_success: float = 0.1
|
data_manipulation_p_of_success: float = 0.1
|
||||||
|
|
||||||
@@ -46,8 +49,31 @@ class DataManipulationBot(DatabaseClient):
|
|||||||
"Whether to repeat attacking once finished."
|
"Whether to repeat attacking once finished."
|
||||||
|
|
||||||
def __init__(self, **kwargs):
|
def __init__(self, **kwargs):
|
||||||
|
kwargs["name"] = "DataManipulationBot"
|
||||||
|
kwargs["port"] = Port.NONE
|
||||||
|
kwargs["protocol"] = IPProtocol.NONE
|
||||||
|
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
self.name = "DataManipulationBot"
|
|
||||||
|
def describe_state(self) -> Dict:
|
||||||
|
"""
|
||||||
|
Produce a dictionary describing the current state of this object.
|
||||||
|
|
||||||
|
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||||
|
|
||||||
|
:return: Current state of this object and child objects.
|
||||||
|
:rtype: Dict
|
||||||
|
"""
|
||||||
|
state = super().describe_state()
|
||||||
|
return state
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _host_db_client(self) -> DatabaseClient:
|
||||||
|
"""Return the database client that is installed on the same machine as the DataManipulationBot."""
|
||||||
|
db_client = self.software_manager.software.get("DatabaseClient")
|
||||||
|
if db_client is None:
|
||||||
|
_LOGGER.info(f"{self.__class__.__name__} cannot find a database client on its host.")
|
||||||
|
return db_client
|
||||||
|
|
||||||
def _init_request_manager(self) -> RequestManager:
|
def _init_request_manager(self) -> RequestManager:
|
||||||
rm = super()._init_request_manager()
|
rm = super()._init_request_manager()
|
||||||
@@ -76,8 +102,8 @@ class DataManipulationBot(DatabaseClient):
|
|||||||
:param repeat: Whether to repeat attacking once finished.
|
:param repeat: Whether to repeat attacking once finished.
|
||||||
"""
|
"""
|
||||||
self.server_ip_address = server_ip_address
|
self.server_ip_address = server_ip_address
|
||||||
self.payload = payload
|
|
||||||
self.server_password = server_password
|
self.server_password = server_password
|
||||||
|
self.payload = payload
|
||||||
self.port_scan_p_of_success = port_scan_p_of_success
|
self.port_scan_p_of_success = port_scan_p_of_success
|
||||||
self.data_manipulation_p_of_success = data_manipulation_p_of_success
|
self.data_manipulation_p_of_success = data_manipulation_p_of_success
|
||||||
self.repeat = repeat
|
self.repeat = repeat
|
||||||
@@ -123,15 +149,17 @@ class DataManipulationBot(DatabaseClient):
|
|||||||
|
|
||||||
:param p_of_success: Probability of successfully performing data manipulation, by default 0.1.
|
:param p_of_success: Probability of successfully performing data manipulation, by default 0.1.
|
||||||
"""
|
"""
|
||||||
|
self._host_db_client.server_ip_address = self.server_ip_address
|
||||||
|
self._host_db_client.server_password = self.server_password
|
||||||
if self.attack_stage == DataManipulationAttackStage.PORT_SCAN:
|
if self.attack_stage == DataManipulationAttackStage.PORT_SCAN:
|
||||||
# perform the actual data manipulation attack
|
# perform the actual data manipulation attack
|
||||||
if simulate_trial(p_of_success):
|
if simulate_trial(p_of_success):
|
||||||
self.sys_log.info(f"{self.name}: Performing data manipulation")
|
self.sys_log.info(f"{self.name}: Performing data manipulation")
|
||||||
# perform the attack
|
# perform the attack
|
||||||
if not len(self.connections):
|
if not len(self._host_db_client.connections):
|
||||||
self.connect()
|
self._host_db_client.connect()
|
||||||
if len(self.connections):
|
if len(self._host_db_client.connections):
|
||||||
self.query(self.payload)
|
self._host_db_client.query(self.payload)
|
||||||
self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
|
self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
|
||||||
attack_successful = True
|
attack_successful = True
|
||||||
if attack_successful:
|
if attack_successful:
|
||||||
|
|||||||
Reference in New Issue
Block a user