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

Database Manipulation Bot bug fixes

Browse Source
This commit is contained in:
Jake Walker
2023-11-24 11:10:34 +00:00
parent 4e104a4d99
commit ff8b773c10
9 changed files with 124 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/config/_package_data/example_config.yaml
View File

@@ -665,7 +665,7 @@ simulation:
options: options:
port_scan_p_of_success: 0.1 port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1 data_manipulation_p_of_success: 0.1
payload: "DROP TABLE IF EXISTS user;" payload: "DELETE"
server_ip: 192.168.1.14 server_ip: 192.168.1.14
services: services:
- ref: client_1_dns_client - ref: client_1_dns_client

4
src/primaite/game/agent/interface.py
View File

@@ -58,7 +58,7 @@ class AbstractAgent(ABC):
action_space: Optional[ActionManager], action_space: Optional[ActionManager],
observation_space: Optional[ObservationManager], observation_space: Optional[ObservationManager],
reward_function: Optional[RewardFunction], reward_function: Optional[RewardFunction],
agent_settings: Optional[AgentSettings], agent_settings: Optional[AgentSettings] = None,
) -> None: ) -> None:
""" """
Initialize an agent. Initialize an agent.
@@ -217,7 +217,7 @@ class DataManipulationAgent(AbstractScriptedAgent):
:return: _description_ :return: _description_
:rtype: Tuple[str, Dict] :rtype: Tuple[str, Dict]
""" """
current_timestep = self.action_space.session.step_counter current_timestep = self.action_manager.session.step_counter
if current_timestep < self.next_execution_timestep: if current_timestep < self.next_execution_timestep:
return "DONOTHING", {"dummy": 0} return "DONOTHING", {"dummy": 0}

7
src/primaite/simulator/network/networks.py
View File

@@ -140,7 +140,12 @@ def arcd_uc2_network() -> Network:
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
client_1.software_manager.install(DataManipulationBot) client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"] db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
db_manipulation_bot.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DELETE") db_manipulation_bot.configure(
server_ip_address=IPv4Address("192.168.1.14"),
payload="DELETE",
port_scan_p_of_success=1.0,
data_manipulation_p_of_success=1.0,
)
# Client 2 # Client 2
client_2 = Computer( client_2 = Computer(

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

@@ -129,9 +129,9 @@ class DatabaseClient(Application):
) )
return self._query(sql=sql, query_id=query_id, is_reattempt=True) return self._query(sql=sql, query_id=query_id, is_reattempt=True)
def execute(self) -> None: def run(self) -> None:
"""Run the DatabaseClient.""" """Run the DatabaseClient."""
super().execute() super().run()
if self.operating_state == ApplicationOperatingState.RUNNING: if self.operating_state == ApplicationOperatingState.RUNNING:
self.connect() self.connect()

8
src/primaite/simulator/system/services/red_services/data_manipulation_bot.py
View File

@@ -50,7 +50,7 @@ class DataManipulationBot(DatabaseClient):
def _init_request_manager(self) -> RequestManager: def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager() rm = super()._init_request_manager()
rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.execute())) rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run()))
return rm return rm
@@ -139,13 +139,13 @@ class DataManipulationBot(DatabaseClient):
self.sys_log.info(f"{self.name}: Data manipulation failed") self.sys_log.info(f"{self.name}: Data manipulation failed")
self.attack_stage = DataManipulationAttackStage.FAILED self.attack_stage = DataManipulationAttackStage.FAILED
def execute(self): def run(self):
""" """
Execute the Data Manipulation Bot. Run the Data Manipulation Bot.
Calls the parent classes execute method before starting the application loop. Calls the parent classes execute method before starting the application loop.
""" """
super().execute() super().run()
self._application_loop() self._application_loop()
def _application_loop(self): def _application_loop(self):

39
tests/assets/configs/bad_primaite_session.yaml
View File

@@ -2,9 +2,17 @@ training_config:
rl_framework: SB3 rl_framework: SB3
rl_algorithm: PPO rl_algorithm: PPO
se3ed: 333 # Purposeful typo to check that error is raised with bad configuration. se3ed: 333 # Purposeful typo to check that error is raised with bad configuration.
n_learn_steps: 2560 n_learn_episodes: 25
n_eval_episodes: 5 n_eval_episodes: 5
max_steps_per_episode: 128
deterministic_eval: false
n_agents: 1
agent_references:
- defender
io_settings:
save_checkpoints: true
checkpoint_interval: 5
game_config: game_config:
@@ -49,6 +57,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: agent_settings:
start_settings:
start_step: 5 start_step: 5
frequency: 4 frequency: 4
variance: 3 variance: 3
@@ -65,8 +74,8 @@ game_config:
observations: observations:
- logon_status - logon_status
- operating_status - operating_status
services: applications:
- service_ref: data_manipulation_bot - application_ref: data_manipulation_bot
observations: observations:
operating_status operating_status
health_status health_status
@@ -76,22 +85,19 @@ game_config:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
#<not yet implemented #<not yet implemented
# - type: NODE_APPLICATION_EXECUTE - type: NODE_APPLICATION_EXECUTE
# options: # - type: NODE_FILE_DELETE
# execution_definition: # - type: NODE_FILE_CORRUPT
# server_ip: 192.168.1.14
# payload: "DROP TABLE IF EXISTS user;"
# success_rate: 80%
- type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT
# - type: NODE_FOLDER_DELETE # - type: NODE_FOLDER_DELETE
# - type: NODE_FOLDER_CORRUPT # - type: NODE_FOLDER_CORRUPT
- type: NODE_OS_SCAN # - type: NODE_OS_SCAN
# - type: NODE_LOGON # - type: NODE_LOGON
# - type: NODE_LOGOFF # - type: NODE_LOGOFF
options: options:
nodes: nodes:
- node_ref: client_1 - node_ref: client_1
applications:
- application_ref: data_manipulation_bot
max_folders_per_node: 1 max_folders_per_node: 1
max_files_per_folder: 1 max_files_per_folder: 1
max_services_per_node: 1 max_services_per_node: 1
@@ -101,6 +107,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self) agent_settings: # options specific to this particular agent type, basically args of __init__(self)
start_settings:
start_step: 25 start_step: 25
frequency: 20 frequency: 20
variance: 5 variance: 5
@@ -652,9 +659,15 @@ simulation:
subnet_mask: 255.255.255.0 subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1 default_gateway: 192.168.10.1
dns_server: 192.168.1.10 dns_server: 192.168.1.10
services: applications:
- ref: data_manipulation_bot - ref: data_manipulation_bot
type: DataManipulationBot type: DataManipulationBot
options:
port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1
payload: "DELETE"
server_ip: 192.168.1.14
services:
- ref: client_1_dns_client - ref: client_1_dns_client
type: DNSClient type: DNSClient

33
tests/assets/configs/eval_only_primaite_session.yaml
View File

@@ -10,6 +10,10 @@ training_config:
agent_references: agent_references:
- defender - defender
io_settings:
save_checkpoints: true
checkpoint_interval: 5
game_config: game_config:
ports: ports:
@@ -53,6 +57,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: agent_settings:
start_settings:
start_step: 5 start_step: 5
frequency: 4 frequency: 4
variance: 3 variance: 3
@@ -69,8 +74,8 @@ game_config:
observations: observations:
- logon_status - logon_status
- operating_status - operating_status
services: applications:
- service_ref: data_manipulation_bot - application_ref: data_manipulation_bot
observations: observations:
operating_status operating_status
health_status health_status
@@ -80,22 +85,19 @@ game_config:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
#<not yet implemented #<not yet implemented
# - type: NODE_APPLICATION_EXECUTE - type: NODE_APPLICATION_EXECUTE
# options: # - type: NODE_FILE_DELETE
# execution_definition: # - type: NODE_FILE_CORRUPT
# server_ip: 192.168.1.14
# payload: "DROP TABLE IF EXISTS user;"
# success_rate: 80%
- type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT
# - type: NODE_FOLDER_DELETE # - type: NODE_FOLDER_DELETE
# - type: NODE_FOLDER_CORRUPT # - type: NODE_FOLDER_CORRUPT
- type: NODE_OS_SCAN # - type: NODE_OS_SCAN
# - type: NODE_LOGON # - type: NODE_LOGON
# - type: NODE_LOGOFF # - type: NODE_LOGOFF
options: options:
nodes: nodes:
- node_ref: client_1 - node_ref: client_1
applications:
- application_ref: data_manipulation_bot
max_folders_per_node: 1 max_folders_per_node: 1
max_files_per_folder: 1 max_files_per_folder: 1
max_services_per_node: 1 max_services_per_node: 1
@@ -105,6 +107,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self) agent_settings: # options specific to this particular agent type, basically args of __init__(self)
start_settings:
start_step: 25 start_step: 25
frequency: 20 frequency: 20
variance: 5 variance: 5
@@ -656,9 +659,15 @@ simulation:
subnet_mask: 255.255.255.0 subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1 default_gateway: 192.168.10.1
dns_server: 192.168.1.10 dns_server: 192.168.1.10
services: applications:
- ref: data_manipulation_bot - ref: data_manipulation_bot
type: DataManipulationBot type: DataManipulationBot
options:
port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1
payload: "DELETE"
server_ip: 192.168.1.14
services:
- ref: client_1_dns_client - ref: client_1_dns_client
type: DNSClient type: DNSClient

29
tests/assets/configs/test_primaite_session.yaml
View File

@@ -57,6 +57,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: agent_settings:
start_settings:
start_step: 5 start_step: 5
frequency: 4 frequency: 4
variance: 3 variance: 3
@@ -73,8 +74,8 @@ game_config:
observations: observations:
- logon_status - logon_status
- operating_status - operating_status
services: applications:
- service_ref: data_manipulation_bot - application_ref: data_manipulation_bot
observations: observations:
operating_status operating_status
health_status health_status
@@ -84,22 +85,19 @@ game_config:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
#<not yet implemented #<not yet implemented
# - type: NODE_APPLICATION_EXECUTE - type: NODE_APPLICATION_EXECUTE
# options: # - type: NODE_FILE_DELETE
# execution_definition: # - type: NODE_FILE_CORRUPT
# server_ip: 192.168.1.14
# payload: "DROP TABLE IF EXISTS user;"
# success_rate: 80%
- type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT
# - type: NODE_FOLDER_DELETE # - type: NODE_FOLDER_DELETE
# - type: NODE_FOLDER_CORRUPT # - type: NODE_FOLDER_CORRUPT
- type: NODE_OS_SCAN # - type: NODE_OS_SCAN
# - type: NODE_LOGON # - type: NODE_LOGON
# - type: NODE_LOGOFF # - type: NODE_LOGOFF
options: options:
nodes: nodes:
- node_ref: client_1 - node_ref: client_1
applications:
- application_ref: data_manipulation_bot
max_folders_per_node: 1 max_folders_per_node: 1
max_files_per_folder: 1 max_files_per_folder: 1
max_services_per_node: 1 max_services_per_node: 1
@@ -109,6 +107,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self) agent_settings: # options specific to this particular agent type, basically args of __init__(self)
start_settings:
start_step: 25 start_step: 25
frequency: 20 frequency: 20
variance: 5 variance: 5
@@ -660,9 +659,15 @@ simulation:
subnet_mask: 255.255.255.0 subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1 default_gateway: 192.168.10.1
dns_server: 192.168.1.10 dns_server: 192.168.1.10
services: applications:
- ref: data_manipulation_bot - ref: data_manipulation_bot
type: DataManipulationBot type: DataManipulationBot
options:
port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1
payload: "DELETE"
server_ip: 192.168.1.14
services:
- ref: client_1_dns_client - ref: client_1_dns_client
type: DNSClient type: DNSClient

33
tests/assets/configs/train_only_primaite_session.yaml
View File

@@ -10,6 +10,10 @@ training_config:
agent_references: agent_references:
- defender - defender
io_settings:
save_checkpoints: true
checkpoint_interval: 5
game_config: game_config:
ports: ports:
@@ -53,6 +57,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: agent_settings:
start_settings:
start_step: 5 start_step: 5
frequency: 4 frequency: 4
variance: 3 variance: 3
@@ -69,8 +74,8 @@ game_config:
observations: observations:
- logon_status - logon_status
- operating_status - operating_status
services: applications:
- service_ref: data_manipulation_bot - application_ref: data_manipulation_bot
observations: observations:
operating_status operating_status
health_status health_status
@@ -80,22 +85,19 @@ game_config:
action_list: action_list:
- type: DONOTHING - type: DONOTHING
#<not yet implemented #<not yet implemented
# - type: NODE_APPLICATION_EXECUTE - type: NODE_APPLICATION_EXECUTE
# options: # - type: NODE_FILE_DELETE
# execution_definition: # - type: NODE_FILE_CORRUPT
# server_ip: 192.168.1.14
# payload: "DROP TABLE IF EXISTS user;"
# success_rate: 80%
- type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT
# - type: NODE_FOLDER_DELETE # - type: NODE_FOLDER_DELETE
# - type: NODE_FOLDER_CORRUPT # - type: NODE_FOLDER_CORRUPT
- type: NODE_OS_SCAN # - type: NODE_OS_SCAN
# - type: NODE_LOGON # - type: NODE_LOGON
# - type: NODE_LOGOFF # - type: NODE_LOGOFF
options: options:
nodes: nodes:
- node_ref: client_1 - node_ref: client_1
applications:
- application_ref: data_manipulation_bot
max_folders_per_node: 1 max_folders_per_node: 1
max_files_per_folder: 1 max_files_per_folder: 1
max_services_per_node: 1 max_services_per_node: 1
@@ -105,6 +107,7 @@ game_config:
- type: DUMMY - type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self) agent_settings: # options specific to this particular agent type, basically args of __init__(self)
start_settings:
start_step: 25 start_step: 25
frequency: 20 frequency: 20
variance: 5 variance: 5
@@ -656,9 +659,15 @@ simulation:
subnet_mask: 255.255.255.0 subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1 default_gateway: 192.168.10.1
dns_server: 192.168.1.10 dns_server: 192.168.1.10
services: applications:
- ref: data_manipulation_bot - ref: data_manipulation_bot
type: DataManipulationBot type: DataManipulationBot
options:
port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1
payload: "DELETE"
server_ip: 192.168.1.14
services:
- ref: client_1_dns_client - ref: client_1_dns_client
type: DNSClient type: DNSClient