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

Merge remote-tracking branch 'origin/dev' into feature/2646_Update-pre-commit-to-check-for-valid-copyright

Browse Source
This commit is contained in:
Marek Wolan
2024-06-13 12:52:09 +01:00
parent 6762fd805e ae2bbd3e00
commit fad8d936e6
54 changed files with 9740 additions and 647 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml Normal file
View File

@@ -0,0 +1,137 @@
io_settings:
save_step_metadata: false
save_pcap_logs: true
save_sys_logs: true
sys_log_level: WARNING
game:
max_episode_length: 256
ports:
- ARP
- DNS
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
agents:
- ref: client_1_red_nmap
team: RED
type: ProbabilisticAgent
observation_space: null
action_space:
options:
nodes:
- node_name: client_1
applications:
- application_name: NMAP
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
action_list:
- type: NODE_NMAP_NETWORK_SERVICE_RECON
action_map:
0:
action: NODE_NMAP_NETWORK_SERVICE_RECON
options:
source_node: client_1
target_ip_address: 192.168.10.0/24
target_port: 80
target_protocol: tcp
reward_function:
reward_components:
- type: DUMMY
agent_settings:
action_probabilities:
0: 1.0
simulation:
network:
nodes:
- hostname: switch_1
num_ports: 8
type: switch
- hostname: switch_2
num_ports: 8
type: switch
- hostname: router_1
type: router
ports:
1:
ip_address: 192.168.1.1
subnet_mask: 255.255.255.0
2:
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
acl:
1:
action: PERMIT
- hostname: client_1
type: computer
ip_address: 192.168.10.21
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
- hostname: client_2
type: computer
ip_address: 192.168.10.22
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
- hostname: server_1
type: server
ip_address: 192.168.1.10
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
- hostname: server_2
type: server
ip_address: 192.168.1.14
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
links:
- endpoint_a_hostname: router_1
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 8
- endpoint_a_hostname: router_1
endpoint_a_port: 2
endpoint_b_hostname: switch_2
endpoint_b_port: 8
- endpoint_a_hostname: client_1
endpoint_a_port: 1
endpoint_b_hostname: switch_2
endpoint_b_port: 1
- endpoint_a_hostname: client_2
endpoint_a_port: 1
endpoint_b_hostname: switch_2
endpoint_b_port: 2
- endpoint_a_hostname: server_1
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 1
- endpoint_a_hostname: server_2
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 2

135
tests/assets/configs/nmap_ping_scan_red_agent_config.yaml Normal file
View File

@@ -0,0 +1,135 @@
io_settings:
save_step_metadata: false
save_pcap_logs: true
save_sys_logs: true
sys_log_level: WARNING
game:
max_episode_length: 256
ports:
- ARP
- DNS
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
agents:
- ref: client_1_red_nmap
team: RED
type: ProbabilisticAgent
observation_space: null
action_space:
options:
nodes:
- node_name: client_1
applications:
- application_name: NMAP
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
action_list:
- type: NODE_NMAP_PING_SCAN
action_map:
0:
action: NODE_NMAP_PING_SCAN
options:
source_node: client_1
target_ip_address: 192.168.1.0/24
reward_function:
reward_components:
- type: DUMMY
agent_settings:
action_probabilities:
0: 1.0
simulation:
network:
nodes:
- hostname: switch_1
num_ports: 8
type: switch
- hostname: switch_2
num_ports: 8
type: switch
- hostname: router_1
type: router
ports:
1:
ip_address: 192.168.1.1
subnet_mask: 255.255.255.0
2:
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
acl:
1:
action: PERMIT
- hostname: client_1
type: computer
ip_address: 192.168.10.21
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
- hostname: client_2
type: computer
ip_address: 192.168.10.22
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
- hostname: server_1
type: server
ip_address: 192.168.1.10
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
- hostname: server_2
type: server
ip_address: 192.168.1.14
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
links:
- endpoint_a_hostname: router_1
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 8
- endpoint_a_hostname: router_1
endpoint_a_port: 2
endpoint_b_hostname: switch_2
endpoint_b_port: 8
- endpoint_a_hostname: client_1
endpoint_a_port: 1
endpoint_b_hostname: switch_2
endpoint_b_port: 1
- endpoint_a_hostname: client_2
endpoint_a_port: 1
endpoint_b_hostname: switch_2
endpoint_b_port: 2
- endpoint_a_hostname: server_1
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 1
- endpoint_a_hostname: server_2
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 2

135
tests/assets/configs/nmap_port_scan_red_agent_config.yaml Normal file
View File

@@ -0,0 +1,135 @@
io_settings:
save_step_metadata: false
save_pcap_logs: true
save_sys_logs: true
sys_log_level: WARNING
game:
max_episode_length: 256
ports:
- ARP
- DNS
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
agents:
- ref: client_1_red_nmap
team: RED
type: ProbabilisticAgent
observation_space: null
action_space:
options:
nodes:
- node_name: client_1
applications:
- application_name: NMAP
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
action_list:
- type: NODE_NMAP_PORT_SCAN
action_map:
0:
action: NODE_NMAP_PORT_SCAN
options:
source_node: client_1
target_ip_address: 192.168.10.0/24
reward_function:
reward_components:
- type: DUMMY
agent_settings:
action_probabilities:
0: 1.0
simulation:
network:
nodes:
- hostname: switch_1
num_ports: 8
type: switch
- hostname: switch_2
num_ports: 8
type: switch
- hostname: router_1
type: router
ports:
1:
ip_address: 192.168.1.1
subnet_mask: 255.255.255.0
2:
ip_address: 192.168.10.1
subnet_mask: 255.255.255.0
acl:
1:
action: PERMIT
- hostname: client_1
type: computer
ip_address: 192.168.10.21
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
- hostname: client_2
type: computer
ip_address: 192.168.10.22
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
- hostname: server_1
type: server
ip_address: 192.168.1.10
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
- hostname: server_2
type: server
ip_address: 192.168.1.14
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
links:
- endpoint_a_hostname: router_1
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 8
- endpoint_a_hostname: router_1
endpoint_a_port: 2
endpoint_b_hostname: switch_2
endpoint_b_port: 8
- endpoint_a_hostname: client_1
endpoint_a_port: 1
endpoint_b_hostname: switch_2
endpoint_b_port: 1
- endpoint_a_hostname: client_2
endpoint_a_port: 1
endpoint_b_hostname: switch_2
endpoint_b_port: 2
- endpoint_a_hostname: server_1
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 1
- endpoint_a_hostname: server_2
endpoint_a_port: 1
endpoint_b_hostname: switch_1
endpoint_b_port: 2

52
tests/integration_tests/system/red_applications/test_ransomware_script.py
View File

@@ -10,12 +10,8 @@ from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.simulator.system.applications.red_applications.ransomware_script import (
RansomwareAttackStage,
RansomwareScript,
)
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.software import SoftwareHealthState
@@ -86,54 +82,24 @@ def ransomware_script_db_server_green_client(example_network) -> Network:
return network
def test_repeating_ransomware_script_attack(ransomware_script_and_db_server):
def test_ransomware_script_attack(ransomware_script_and_db_server):
"""Test a repeating data manipulation attack."""
RansomwareScript, computer, db_server_service, server = ransomware_script_and_db_server
computer.apply_timestep(timestep=0)
server.apply_timestep(timestep=0)
assert db_server_service.health_state_actual is SoftwareHealthState.GOOD
assert computer.file_system.num_file_creations == 0
assert server.file_system.num_file_creations == 1
RansomwareScript.target_scan_p_of_success = 1
RansomwareScript.c2_beacon_p_of_success = 1
RansomwareScript.ransomware_encrypt_p_of_success = 1
RansomwareScript.repeat = True
RansomwareScript.attack()
assert RansomwareScript.attack_stage == RansomwareAttackStage.NOT_STARTED
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED
assert computer.file_system.num_file_creations == 1
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT
assert server.file_system.num_file_creations == 2
computer.apply_timestep(timestep=1)
server.apply_timestep(timestep=1)
assert RansomwareScript.attack_stage == RansomwareAttackStage.NOT_STARTED
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED
def test_repeating_ransomware_script_attack(ransomware_script_and_db_server):
"""Test a repeating ransowmare script attack."""
RansomwareScript, computer, db_server_service, server = ransomware_script_and_db_server
assert db_server_service.health_state_actual is SoftwareHealthState.GOOD
RansomwareScript.target_scan_p_of_success = 1
RansomwareScript.c2_beacon_p_of_success = 1
RansomwareScript.ransomware_encrypt_p_of_success = 1
RansomwareScript.repeat = False
RansomwareScript.attack()
assert RansomwareScript.attack_stage == RansomwareAttackStage.SUCCEEDED
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT
assert computer.file_system.num_file_creations == 1
computer.apply_timestep(timestep=1)
computer.pre_timestep(timestep=1)
server.apply_timestep(timestep=1)
server.pre_timestep(timestep=1)
assert RansomwareScript.attack_stage == RansomwareAttackStage.SUCCEEDED
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT
assert computer.file_system.num_file_creations == 0
def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_green_client):
@@ -154,10 +120,6 @@ def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_
assert green_db_client_connection.query("SELECT")
assert green_db_client.last_query_response.get("status_code") == 200
ransomware_script_application.target_scan_p_of_success = 1
ransomware_script_application.ransomware_encrypt_p_of_success = 1
ransomware_script_application.c2_beacon_p_of_success = 1
ransomware_script_application.repeat = False
ransomware_script_application.attack()
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT

186
tests/integration_tests/system/test_nmap.py Normal file
View File

@@ -0,0 +1,186 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
import yaml
from primaite.game.game import PrimaiteGame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.nmap import NMAP
from tests import TEST_ASSETS_ROOT
def test_ping_scan_all_on(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
expected_result = [IPv4Address("192.168.1.10"), IPv4Address("192.168.1.14")]
actual_result = client_1_nmap.ping_scan(target_ip_address=["192.168.1.10", "192.168.1.14"])
assert sorted(actual_result) == sorted(expected_result)
def test_ping_scan_all_on_full_network(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
expected_result = [IPv4Address("192.168.1.1"), IPv4Address("192.168.1.10"), IPv4Address("192.168.1.14")]
actual_result = client_1_nmap.ping_scan(target_ip_address=IPv4Network("192.168.1.0/24"))
assert sorted(actual_result) == sorted(expected_result)
def test_ping_scan_some_on(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
network.get_node_by_hostname("server_2").power_off()
expected_result = [IPv4Address("192.168.1.1"), IPv4Address("192.168.1.10")]
actual_result = client_1_nmap.ping_scan(target_ip_address=IPv4Network("192.168.1.0/24"))
assert sorted(actual_result) == sorted(expected_result)
def test_ping_scan_all_off(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
network.get_node_by_hostname("server_1").power_off()
network.get_node_by_hostname("server_2").power_off()
expected_result = []
actual_result = client_1_nmap.ping_scan(target_ip_address=["192.168.1.10", "192.168.1.14"])
assert sorted(actual_result) == sorted(expected_result)
def test_port_scan_one_node_one_port(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
client_2 = network.get_node_by_hostname("client_2")
actual_result = client_1_nmap.port_scan(
target_ip_address=client_2.network_interface[1].ip_address, target_port=Port.DNS, target_protocol=IPProtocol.TCP
)
expected_result = {IPv4Address("192.168.10.22"): {IPProtocol.TCP: [Port.DNS]}}
assert actual_result == expected_result
def sort_dict(d):
"""Recursively sorts a dictionary."""
if isinstance(d, dict):
return {k: sort_dict(v) for k, v in sorted(d.items(), key=lambda item: str(item[0]))}
elif isinstance(d, list):
return sorted(d, key=lambda item: str(item) if isinstance(item, Enum) else item)
elif isinstance(d, Enum):
return str(d)
else:
return d
def test_port_scan_full_subnet_all_ports_and_protocols(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
actual_result = client_1_nmap.port_scan(
target_ip_address=IPv4Network("192.168.10.0/24"),
)
expected_result = {
IPv4Address("192.168.10.1"): {IPProtocol.UDP: [Port.ARP]},
IPv4Address("192.168.10.22"): {
IPProtocol.TCP: [Port.HTTP, Port.FTP, Port.DNS],
IPProtocol.UDP: [Port.ARP, Port.NTP],
},
}
assert sort_dict(actual_result) == sort_dict(expected_result)
def test_network_service_recon_all_ports_and_protocols(example_network):
network = example_network
client_1 = network.get_node_by_hostname("client_1")
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
actual_result = client_1_nmap.network_service_recon(
target_ip_address=IPv4Network("192.168.10.0/24"), target_port=Port.HTTP, target_protocol=IPProtocol.TCP
)
expected_result = {IPv4Address("192.168.10.22"): {IPProtocol.TCP: [Port.HTTP]}}
assert sort_dict(actual_result) == sort_dict(expected_result)
def test_ping_scan_red_agent():
with open(TEST_ASSETS_ROOT / "configs/nmap_ping_scan_red_agent_config.yaml", "r") as file:
cfg = yaml.safe_load(file)
game = PrimaiteGame.from_config(cfg)
game.step()
expected_result = ["192.168.1.1", "192.168.1.10", "192.168.1.14"]
action_history = game.agents["client_1_red_nmap"].history
assert len(action_history) == 1
actual_result = action_history[0].response.data["live_hosts"]
assert sorted(actual_result) == sorted(expected_result)
def test_port_scan_red_agent():
with open(TEST_ASSETS_ROOT / "configs/nmap_port_scan_red_agent_config.yaml", "r") as file:
cfg = yaml.safe_load(file)
game = PrimaiteGame.from_config(cfg)
game.step()
expected_result = {
"192.168.10.1": {"udp": [219]},
"192.168.10.22": {
"tcp": [80, 21, 53],
"udp": [219, 123],
},
}
action_history = game.agents["client_1_red_nmap"].history
assert len(action_history) == 1
actual_result = action_history[0].response.data
assert sorted(actual_result) == sorted(expected_result)
def test_network_service_recon_red_agent():
with open(TEST_ASSETS_ROOT / "configs/nmap_network_service_recon_red_agent_config.yaml", "r") as file:
cfg = yaml.safe_load(file)
game = PrimaiteGame.from_config(cfg)
game.step()
expected_result = {"192.168.10.22": {"tcp": [80]}}
action_history = game.agents["client_1_red_nmap"].history
assert len(action_history) == 1
actual_result = action_history[0].response.data
assert sorted(actual_result) == sorted(expected_result)

1
tests/unit_tests/_primaite/_utils/__init__.py Normal file
View File

@@ -0,0 +1 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK

84
tests/unit_tests/_primaite/_utils/test_dict_enum_keys_conversion.py Normal file
View File

@@ -0,0 +1,84 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.utils.converters import convert_dict_enum_keys_to_enum_values
def test_simple_conversion():
"""
Test conversion of a simple dictionary with enum keys to enum values.
The original dictionary contains one level of nested dictionary with enums as keys.
The expected output should have string values of enums as keys.
"""
original_dict = {IPProtocol.UDP: {Port.ARP: {"inbound": 0, "outbound": 1016.0}}}
expected_dict = {"udp": {219: {"inbound": 0, "outbound": 1016.0}}}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
def test_no_enums():
"""
Test conversion of a dictionary with no enum keys.
The original dictionary contains only string keys.
The expected output should be identical to the original dictionary.
"""
original_dict = {"protocol": {"port": {"inbound": 0, "outbound": 1016.0}}}
expected_dict = {"protocol": {"port": {"inbound": 0, "outbound": 1016.0}}}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
def test_mixed_keys():
"""
Test conversion of a dictionary with a mix of enum and string keys.
The original dictionary contains both enums and strings as keys.
The expected output should have string values of enums and original string keys.
"""
original_dict = {
IPProtocol.TCP: {"port": {"inbound": 0, "outbound": 1016.0}},
"protocol": {Port.HTTP: {"inbound": 10, "outbound": 2020.0}},
}
expected_dict = {
"tcp": {"port": {"inbound": 0, "outbound": 1016.0}},
"protocol": {80: {"inbound": 10, "outbound": 2020.0}},
}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
def test_empty_dict():
"""
Test conversion of an empty dictionary.
The original dictionary is empty.
The expected output should also be an empty dictionary.
"""
original_dict = {}
expected_dict = {}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
def test_nested_dicts():
"""
Test conversion of a nested dictionary with multiple levels of nested dictionaries and enums as keys.
The original dictionary contains nested dictionaries with enums as keys at different levels.
The expected output should have string values of enums as keys at all levels.
"""
original_dict = {
IPProtocol.UDP: {Port.ARP: {"inbound": 0, "outbound": 1016.0, "details": {IPProtocol.TCP: {"latency": "low"}}}}
}
expected_dict = {"udp": {219: {"inbound": 0, "outbound": 1016.0, "details": {"tcp": {"latency": "low"}}}}}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
def test_non_dict_values():
"""
Test conversion of a dictionary where some values are not dictionaries.
The original dictionary contains lists and tuples as values.
The expected output should preserve these non-dictionary values while converting enum keys to string values.
"""
original_dict = {IPProtocol.UDP: [Port.ARP, Port.HTTP], "protocols": (IPProtocol.TCP, IPProtocol.UDP)}
expected_dict = {"udp": [Port.ARP, Port.HTTP], "protocols": (IPProtocol.TCP, IPProtocol.UDP)}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict