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

Merge remote-tracking branch 'origin/dev' into dev-v3.0.0b6

Browse Source
This commit is contained in:
Marek Wolan
2024-01-29 10:26:28 +00:00
parent 2ba05e7348 62197fe72b
commit 555802baaa
57 changed files with 2329 additions and 461 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -93,25 +93,25 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: domain_controller_dns_server
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
- service_name: web_server_database_client
- node_hostname: database_server
services:
- service_ref: database_service
- service_name: database_service
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
- node_hostname: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -126,7 +126,7 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
nic_num: 1
@@ -514,7 +514,7 @@ agents:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
@@ -522,8 +522,8 @@ agents:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: web_server
service_name: web_server_web_service
agent_settings:

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

@@ -31,13 +31,6 @@ agents:
action_space:
action_list:
- type: DONOTHING
# <not yet implemented>
# - type: NODE_LOGON
# - type: NODE_LOGOFF
# - type: NODE_APPLICATION_EXECUTE
# options:
# execution_definition:
# target_address: arcd.com
options:
nodes:
@@ -104,25 +97,25 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: domain_controller_dns_server
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
- service_name: web_server_database_client
- node_hostname: database_server
services:
- service_ref: database_service
- service_name: database_service
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
- node_hostname: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -137,7 +130,7 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
nic_num: 1
@@ -525,7 +518,7 @@ agents:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
@@ -533,8 +526,8 @@ agents:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: web_server
service_name: web_server_web_service
agent_settings:

63
tests/assets/configs/multi_agent_session.yaml
View File

@@ -37,13 +37,6 @@ agents:
action_space:
action_list:
- type: DONOTHING
# <not yet implemented>
# - type: NODE_LOGON
# - type: NODE_LOGOFF
# - type: NODE_APPLICATION_EXECUTE
# options:
# execution_definition:
# target_address: arcd.com
options:
nodes:
@@ -111,25 +104,25 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: domain_controller_dns_server
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
- service_name: web_server_database_client
- node_hostname: database_server
services:
- service_ref: database_service
- service_name: database_service
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
- node_hostname: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -144,7 +137,7 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
nic_num: 1
@@ -532,7 +525,7 @@ agents:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
@@ -540,8 +533,8 @@ agents:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: web_server
service_name: web_server_web_service
agent_settings:
@@ -559,25 +552,25 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: domain_controller_dns_server
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
- service_name: web_server_database_client
- node_hostname: database_server
services:
- service_ref: database_service
- service_name: database_service
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
- node_hostname: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -592,7 +585,7 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
nic_num: 1
@@ -980,7 +973,7 @@ agents:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
@@ -988,8 +981,8 @@ agents:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: web_server
service_name: web_server_web_service
agent_settings:

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

@@ -35,13 +35,6 @@ agents:
action_space:
action_list:
- type: DONOTHING
# <not yet implemented>
# - type: NODE_LOGON
# - type: NODE_LOGOFF
# - type: NODE_APPLICATION_EXECUTE
# options:
# execution_definition:
# target_address: arcd.com
options:
nodes:
@@ -109,25 +102,25 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: domain_controller_dns_server
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
- service_name: web_server_database_client
- node_hostname: database_server
services:
- service_ref: database_service
- service_name: database_service
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
- node_hostname: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
# - service_name: backup_service
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -142,7 +135,7 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
nic_num: 1
@@ -530,7 +523,7 @@ agents:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
@@ -538,8 +531,8 @@ agents:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: web_server
service_name: web_server_web_service
agent_settings:

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

@@ -105,25 +105,23 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: domain_controller_dns_server
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
- service_name: web_server_database_client
- node_hostname: database_server
services:
- service_ref: database_service
- service_name: database_service
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: backup_server
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -138,7 +136,7 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
nic_num: 1
@@ -526,7 +524,7 @@ agents:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
@@ -534,8 +532,8 @@ agents:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: web_server
service_name: web_service
agent_settings:

2
tests/conftest.py
View File

@@ -170,7 +170,7 @@ def example_network() -> Network:
-------------- --------------
| client_1 |----- ----| server_1 |
-------------- | -------------- -------------- -------------- | --------------
------| switch_1 |------| router_1 |------| switch_2 |------
------| switch_2 |------| router_1 |------| switch_1 |------
-------------- | -------------- -------------- -------------- | --------------
| client_2 |---- ----| server_2 |
-------------- --------------

2
tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py
View File

@@ -1,8 +1,8 @@
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
def test_data_manipulation(uc2_network):

2
tests/integration_tests/game_layer/test_observations.py
View File

@@ -14,7 +14,7 @@ def test_file_observation():
state = sim.describe_state()
dog_file_obs = FileObservation(
where=["network", "nodes", pc.uuid, "file_system", "folders", "root", "files", "dog.png"]
where=["network", "nodes", pc.hostname, "file_system", "folders", "root", "files", "dog.png"]
)
assert dog_file_obs.observe(state) == {"health_status": 1}
assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)})

180
tests/integration_tests/network/test_broadcast.py Normal file
View File

@@ -0,0 +1,180 @@
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Tuple
import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
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.services.service import Service
class BroadcastService(Service):
"""A service for sending broadcast and unicast messages over a network."""
def __init__(self, **kwargs):
# Set default service properties for broadcasting
kwargs["name"] = "BroadcastService"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def describe_state(self) -> Dict:
# Implement state description for the service
pass
def unicast(self, ip_address: IPv4Address):
# Send a unicast payload to a specific IP address
super().send(
payload="unicast",
dest_ip_address=ip_address,
dest_port=Port.HTTP,
)
def broadcast(self, ip_network: IPv4Network):
# Send a broadcast payload to an entire IP network
super().send(
payload="broadcast",
dest_ip_address=ip_network,
dest_port=Port.HTTP,
)
class BroadcastClient(Application):
"""A client application to receive broadcast and unicast messages."""
payloads_received: List = []
def __init__(self, **kwargs):
# Set default client properties
kwargs["name"] = "BroadcastClient"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def describe_state(self) -> Dict:
# Implement state description for the application
pass
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
# Append received payloads to the list and print a message
self.payloads_received.append(payload)
print(f"Payload: {payload} received on node {self.sys_log.hostname}")
@pytest.fixture(scope="function")
def broadcast_network() -> Network:
network = Network()
client_1 = Computer(
hostname="client_1",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
client_1.power_on()
client_1.software_manager.install(BroadcastClient)
application_1 = client_1.software_manager.software["BroadcastClient"]
application_1.run()
client_2 = Computer(
hostname="client_2",
ip_address="192.168.1.3",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
client_2.power_on()
client_2.software_manager.install(BroadcastClient)
application_2 = client_2.software_manager.software["BroadcastClient"]
application_2.run()
server_1 = Server(
hostname="server_1",
ip_address="192.168.1.1",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server_1.power_on()
server_1.software_manager.install(BroadcastService)
service: BroadcastService = server_1.software_manager.software["BroadcastService"]
service.start()
switch_1 = Switch(hostname="switch_1", num_ports=6, start_up_duration=0)
switch_1.power_on()
network.connect(endpoint_a=client_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[1])
network.connect(endpoint_a=client_2.ethernet_port[1], endpoint_b=switch_1.switch_ports[2])
network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[3])
return network
@pytest.fixture(scope="function")
def broadcast_service_and_clients(broadcast_network) -> Tuple[BroadcastService, BroadcastClient, BroadcastClient]:
client_1: BroadcastClient = broadcast_network.get_node_by_hostname("client_1").software_manager.software[
"BroadcastClient"
]
client_2: BroadcastClient = broadcast_network.get_node_by_hostname("client_2").software_manager.software[
"BroadcastClient"
]
service: BroadcastService = broadcast_network.get_node_by_hostname("server_1").software_manager.software[
"BroadcastService"
]
return service, client_1, client_2
def test_broadcast_correct_subnet(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.broadcast(IPv4Network("192.168.1.0/24"))
assert client_1.payloads_received == ["broadcast"]
assert client_2.payloads_received == ["broadcast"]
def test_broadcast_incorrect_subnet(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.broadcast(IPv4Network("192.168.2.0/24"))
assert not client_1.payloads_received
assert not client_2.payloads_received
def test_unicast_correct_address(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.unicast(IPv4Address("192.168.1.2"))
assert client_1.payloads_received == ["unicast"]
assert not client_2.payloads_received
def test_unicast_incorrect_address(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.unicast(IPv4Address("192.168.2.2"))
assert not client_1.payloads_received
assert not client_2.payloads_received

133
tests/integration_tests/network/test_routing.py
View File

@@ -1,11 +1,16 @@
from ipaddress import IPv4Address
from typing import Tuple
import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.base import Link, NIC, Node, NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import ACLAction, Router
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
@pytest.fixture(scope="function")
@@ -34,6 +39,69 @@ def pc_a_pc_b_router_1() -> Tuple[Node, Node, Router]:
return pc_a, pc_b, router_1
@pytest.fixture(scope="function")
def multi_hop_network() -> Network:
network = Network()
# Configure PC A
pc_a = Computer(
hostname="pc_a",
ip_address="192.168.0.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.0.1",
start_up_duration=0,
)
pc_a.power_on()
network.add_node(pc_a)
# Configure Router 1
router_1 = Router(hostname="router_1", start_up_duration=0)
router_1.power_on()
network.add_node(router_1)
# Configure the connection between PC A and Router 1 port 2
router_1.configure_port(2, "192.168.0.1", "255.255.255.0")
network.connect(pc_a.ethernet_port[1], router_1.ethernet_ports[2])
router_1.enable_port(2)
# Configure Router 1 ACLs
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
# Configure PC B
pc_b = Computer(
hostname="pc_b",
ip_address="192.168.2.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.2.1",
start_up_duration=0,
)
pc_b.power_on()
network.add_node(pc_b)
# Configure Router 2
router_2 = Router(hostname="router_2", start_up_duration=0)
router_2.power_on()
network.add_node(router_2)
# Configure the connection between PC B and Router 2 port 2
router_2.configure_port(2, "192.168.2.1", "255.255.255.0")
network.connect(pc_b.ethernet_port[1], router_2.ethernet_ports[2])
router_2.enable_port(2)
# Configure Router 2 ACLs
router_2.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
router_2.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
# Configure the connection between Router 1 port 1 and Router 2 port 1
router_2.configure_port(1, "192.168.1.2", "255.255.255.252")
router_1.configure_port(1, "192.168.1.1", "255.255.255.252")
network.connect(router_1.ethernet_ports[1], router_2.ethernet_ports[1])
router_1.enable_port(1)
router_2.enable_port(1)
return network
def test_ping_default_gateway(pc_a_pc_b_router_1):
pc_a, pc_b, router_1 = pc_a_pc_b_router_1
@@ -50,3 +118,68 @@ def test_host_on_other_subnet(pc_a_pc_b_router_1):
pc_a, pc_b, router_1 = pc_a_pc_b_router_1
assert pc_a.ping("192.168.1.10")
def test_no_route_no_ping(multi_hop_network):
pc_a = multi_hop_network.get_node_by_hostname("pc_a")
pc_b = multi_hop_network.get_node_by_hostname("pc_b")
assert not pc_a.ping(pc_b.ethernet_port[1].ip_address)
def test_with_routes_can_ping(multi_hop_network):
pc_a = multi_hop_network.get_node_by_hostname("pc_a")
pc_b = multi_hop_network.get_node_by_hostname("pc_b")
router_1: Router = multi_hop_network.get_node_by_hostname("router_1") # noqa
router_2: Router = multi_hop_network.get_node_by_hostname("router_2") # noqa
# Configure Route from Router 1 to PC B subnet
router_1.route_table.add_route(
address="192.168.2.0", subnet_mask="255.255.255.0", next_hop_ip_address="192.168.1.2"
)
# Configure Route from Router 2 to PC A subnet
router_2.route_table.add_route(
address="192.168.0.2", subnet_mask="255.255.255.0", next_hop_ip_address="192.168.1.1"
)
assert pc_a.ping(pc_b.ethernet_port[1].ip_address)
def test_routing_services(multi_hop_network):
pc_a = multi_hop_network.get_node_by_hostname("pc_a")
pc_b = multi_hop_network.get_node_by_hostname("pc_b")
pc_a.software_manager.install(NTPClient)
ntp_client = pc_a.software_manager.software["NTPClient"]
ntp_client.start()
pc_b.software_manager.install(NTPServer)
pc_b.software_manager.software["NTPServer"].start()
ntp_client.configure(ntp_server_ip_address=pc_b.ethernet_port[1].ip_address)
router_1: Router = multi_hop_network.get_node_by_hostname("router_1") # noqa
router_2: Router = multi_hop_network.get_node_by_hostname("router_2") # noqa
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.NTP, dst_port=Port.NTP, position=21)
router_2.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.NTP, dst_port=Port.NTP, position=21)
assert ntp_client.time is None
ntp_client.request_time()
assert ntp_client.time is None
# Configure Route from Router 1 to PC B subnet
router_1.route_table.add_route(
address="192.168.2.0", subnet_mask="255.255.255.0", next_hop_ip_address="192.168.1.2"
)
# Configure Route from Router 2 to PC A subnet
router_2.route_table.add_route(
address="192.168.0.2", subnet_mask="255.255.255.0", next_hop_ip_address="192.168.1.1"
)
ntp_client.request_time()
assert ntp_client.time is not None

180
tests/integration_tests/system/red_applications/test_dos_bot_and_server.py Normal file
View File

@@ -0,0 +1,180 @@
from ipaddress import IPv4Address
from typing import Tuple
import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.server import Server
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
from primaite.simulator.system.applications.red_applications.dos_bot import DoSAttackStage, DoSBot
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.software import SoftwareHealthState
@pytest.fixture(scope="function")
def dos_bot_and_db_server(client_server) -> Tuple[DoSBot, Computer, DatabaseService, Server]:
computer, server = client_server
# Install DoSBot on computer
computer.software_manager.install(DoSBot)
dos_bot: DoSBot = computer.software_manager.software.get("DoSBot")
dos_bot.configure(
target_ip_address=IPv4Address(server.nics.get(next(iter(server.nics))).ip_address),
target_port=Port.POSTGRES_SERVER,
)
# Install DB Server service on server
server.software_manager.install(DatabaseService)
db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService")
db_server_service.start()
return dos_bot, computer, db_server_service, server
@pytest.fixture(scope="function")
def dos_bot_db_server_green_client(example_network) -> Network:
network: Network = example_network
router_1: Router = example_network.get_node_by_hostname("router_1")
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0
)
client_1: Computer = network.get_node_by_hostname("client_1")
client_2: Computer = network.get_node_by_hostname("client_2")
server: Server = network.get_node_by_hostname("server_1")
# install DoS bot on client 1
client_1.software_manager.install(DoSBot)
dos_bot: DoSBot = client_1.software_manager.software.get("DoSBot")
dos_bot.configure(
target_ip_address=IPv4Address(server.nics.get(next(iter(server.nics))).ip_address),
target_port=Port.POSTGRES_SERVER,
)
# install db server service on server
server.software_manager.install(DatabaseService)
db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService")
db_server_service.start()
# Install DB client (green) on client 2
client_2.software_manager.install(DatabaseClient)
database_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient")
database_client.configure(server_ip_address=IPv4Address("192.168.0.1"))
database_client.run()
return network
def test_repeating_dos_attack(dos_bot_and_db_server):
dos_bot, computer, db_server_service, server = dos_bot_and_db_server
assert db_server_service.health_state_actual is SoftwareHealthState.GOOD
dos_bot.port_scan_p_of_success = 1
dos_bot.repeat = True
dos_bot.run()
assert len(dos_bot.connections) == db_server_service.max_sessions
assert len(db_server_service.connections) == db_server_service.max_sessions
assert len(dos_bot.connections) == db_server_service.max_sessions
assert dos_bot.attack_stage is DoSAttackStage.NOT_STARTED
assert db_server_service.health_state_actual is SoftwareHealthState.OVERWHELMED
db_server_service.clear_connections()
db_server_service.set_health_state(SoftwareHealthState.GOOD)
assert len(db_server_service.connections) == 0
computer.apply_timestep(timestep=1)
server.apply_timestep(timestep=1)
assert len(dos_bot.connections) == db_server_service.max_sessions
assert len(db_server_service.connections) == db_server_service.max_sessions
assert len(dos_bot.connections) == db_server_service.max_sessions
assert dos_bot.attack_stage is DoSAttackStage.NOT_STARTED
assert db_server_service.health_state_actual is SoftwareHealthState.OVERWHELMED
def test_non_repeating_dos_attack(dos_bot_and_db_server):
dos_bot, computer, db_server_service, server = dos_bot_and_db_server
assert db_server_service.health_state_actual is SoftwareHealthState.GOOD
dos_bot.port_scan_p_of_success = 1
dos_bot.repeat = False
dos_bot.run()
assert len(dos_bot.connections) == db_server_service.max_sessions
assert len(db_server_service.connections) == db_server_service.max_sessions
assert len(dos_bot.connections) == db_server_service.max_sessions
assert dos_bot.attack_stage is DoSAttackStage.COMPLETED
assert db_server_service.health_state_actual is SoftwareHealthState.OVERWHELMED
db_server_service.clear_connections()
db_server_service.set_health_state(SoftwareHealthState.GOOD)
assert len(db_server_service.connections) == 0
computer.apply_timestep(timestep=1)
server.apply_timestep(timestep=1)
assert len(dos_bot.connections) == 0
assert len(db_server_service.connections) == 0
assert len(dos_bot.connections) == 0
assert dos_bot.attack_stage is DoSAttackStage.COMPLETED
assert db_server_service.health_state_actual is SoftwareHealthState.GOOD
def test_dos_bot_database_service_connection(dos_bot_and_db_server):
dos_bot, computer, db_server_service, server = dos_bot_and_db_server
dos_bot.operating_state = ApplicationOperatingState.RUNNING
dos_bot.attack_stage = DoSAttackStage.PORT_SCAN
dos_bot._perform_dos()
assert len(dos_bot.connections) == db_server_service.max_sessions
assert len(db_server_service.connections) == db_server_service.max_sessions
assert len(dos_bot.connections) == db_server_service.max_sessions
def test_dos_blocks_green_agent_connection(dos_bot_db_server_green_client):
network: Network = dos_bot_db_server_green_client
client_1: Computer = network.get_node_by_hostname("client_1")
dos_bot: DoSBot = client_1.software_manager.software.get("DoSBot")
client_2: Computer = network.get_node_by_hostname("client_2")
green_db_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient")
server: Server = network.get_node_by_hostname("server_1")
db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService")
assert db_server_service.health_state_actual is SoftwareHealthState.GOOD
dos_bot.port_scan_p_of_success = 1
dos_bot.repeat = False
dos_bot.run()
# DoS bot fills up connection of db server service
assert len(dos_bot.connections) == db_server_service.max_sessions
assert len(db_server_service.connections) == db_server_service.max_sessions
assert len(dos_bot.connections) == db_server_service.max_sessions
assert len(green_db_client.connections) == 0
assert dos_bot.attack_stage is DoSAttackStage.COMPLETED
# db server service is overwhelmed
assert db_server_service.health_state_actual is SoftwareHealthState.OVERWHELMED
# green agent tries to connect but fails because service is overwhelmed
assert green_db_client.connect() is False
assert len(green_db_client.connections) == 0

8
tests/integration_tests/system/test_application_on_node.py
View File

@@ -65,8 +65,8 @@ def test_server_turns_off_application(populated_node):
assert app.operating_state is ApplicationOperatingState.CLOSED
def test_application_cannot_be_turned_on_when_server_is_off(populated_node):
"""Check that the application cannot be started when the server is off."""
def test_application_cannot_be_turned_on_when_computer_is_off(populated_node):
"""Check that the application cannot be started when the computer is off."""
app, computer = populated_node
assert computer.operating_state is NodeOperatingState.ON
@@ -86,8 +86,8 @@ def test_application_cannot_be_turned_on_when_server_is_off(populated_node):
assert app.operating_state is ApplicationOperatingState.CLOSED
def test_server_turns_on_application(populated_node):
"""Check that turning on the server turns on application."""
def test_computer_runs_applications(populated_node):
"""Check that turning on the computer will turn on applications."""
app, computer = populated_node
assert computer.operating_state is NodeOperatingState.ON

121
tests/integration_tests/system/test_database_on_node.py
View File

@@ -1,6 +1,9 @@
from ipaddress import IPv4Address
from typing import Tuple
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
import pytest
from primaite.simulator.network.hardware.base import Link, NIC, Node, NodeOperatingState
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.database.database_service import DatabaseService
@@ -8,57 +11,109 @@ from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
def test_database_client_server_connection(uc2_network):
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
@pytest.fixture(scope="function")
def peer_to_peer() -> Tuple[Node, Node]:
node_a = Node(hostname="node_a", operating_state=NodeOperatingState.ON)
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", operating_state=NodeOperatingState.ON)
node_a.connect_nic(nic_a)
node_a.software_manager.get_open_ports()
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
node_b = Node(hostname="node_b", operating_state=NodeOperatingState.ON)
nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0")
node_b.connect_nic(nic_b)
Link(endpoint_a=nic_a, endpoint_b=nic_b)
assert node_a.ping("192.168.0.11")
node_a.software_manager.install(DatabaseClient)
node_a.software_manager.software["DatabaseClient"].configure(server_ip_address=IPv4Address("192.168.0.11"))
node_a.software_manager.software["DatabaseClient"].run()
node_b.software_manager.install(DatabaseService)
database_service: DatabaseService = node_b.software_manager.software["DatabaseService"] # noqa
database_service.start()
return node_a, node_b
@pytest.fixture(scope="function")
def peer_to_peer_secure_db() -> Tuple[Node, Node]:
node_a = Node(hostname="node_a", operating_state=NodeOperatingState.ON)
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", operating_state=NodeOperatingState.ON)
node_a.connect_nic(nic_a)
node_a.software_manager.get_open_ports()
node_b = Node(hostname="node_b", operating_state=NodeOperatingState.ON)
nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0")
node_b.connect_nic(nic_b)
Link(endpoint_a=nic_a, endpoint_b=nic_b)
assert node_a.ping("192.168.0.11")
node_a.software_manager.install(DatabaseClient)
node_a.software_manager.software["DatabaseClient"].configure(server_ip_address=IPv4Address("192.168.0.11"))
node_a.software_manager.software["DatabaseClient"].run()
node_b.software_manager.install(DatabaseService)
database_service: DatabaseService = node_b.software_manager.software["DatabaseService"] # noqa
database_service.password = "12345"
database_service.start()
return node_a, node_b
def test_database_client_server_connection(peer_to_peer):
node_a, node_b = peer_to_peer
db_client: DatabaseClient = node_a.software_manager.software["DatabaseClient"]
db_service: DatabaseService = node_b.software_manager.software["DatabaseService"]
db_client.connect()
assert len(db_client.connections) == 1
assert len(db_service.connections) == 1
db_client.disconnect()
assert len(db_client.connections) == 0
assert len(db_service.connections) == 0
def test_database_client_server_correct_password(uc2_network):
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
def test_database_client_server_correct_password(peer_to_peer_secure_db):
node_a, node_b = peer_to_peer_secure_db
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
db_client: DatabaseClient = node_a.software_manager.software["DatabaseClient"]
db_client.disconnect()
db_client.configure(server_ip_address=IPv4Address("192.168.1.14"), server_password="12345")
db_service.password = "12345"
assert db_client.connect()
db_service: DatabaseService = node_b.software_manager.software["DatabaseService"]
db_client.configure(server_ip_address=IPv4Address("192.168.0.11"), server_password="12345")
db_client.connect()
assert len(db_client.connections) == 1
assert len(db_service.connections) == 1
def test_database_client_server_incorrect_password(uc2_network):
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
def test_database_client_server_incorrect_password(peer_to_peer_secure_db):
node_a, node_b = peer_to_peer_secure_db
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
db_client: DatabaseClient = node_a.software_manager.software["DatabaseClient"]
db_client.disconnect()
db_client.configure(server_ip_address=IPv4Address("192.168.1.14"), server_password="54321")
db_service.password = "12345"
db_service: DatabaseService = node_b.software_manager.software["DatabaseService"]
assert not db_client.connect()
# should fail
db_client.connect()
assert len(db_client.connections) == 0
assert len(db_service.connections) == 0
db_client.configure(server_ip_address=IPv4Address("192.168.0.11"), server_password="wrongpass")
db_client.connect()
assert len(db_client.connections) == 0
assert len(db_service.connections) == 0
def test_database_client_query(uc2_network):
"""Tests DB query across the network returns HTTP status 200 and date."""
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
assert db_client.connected
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_client.connect()
assert db_client.query("SELECT")
@@ -66,13 +121,13 @@ def test_database_client_query(uc2_network):
def test_create_database_backup(uc2_network):
"""Run the backup_database method and check if the FTP server has the relevant file."""
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
# back up should be created
assert db_service.backup_database() is True
backup_server: Server = uc2_network.get_node_by_hostname("backup_server")
ftp_server: FTPServer = backup_server.software_manager.software.get("FTPServer")
ftp_server: FTPServer = backup_server.software_manager.software["FTPServer"]
# backup file should exist in the backup server
assert ftp_server.file_system.get_file(folder_name=db_service.uuid, file_name="database.db") is not None
@@ -81,7 +136,7 @@ def test_create_database_backup(uc2_network):
def test_restore_backup(uc2_network):
"""Run the restore_backup method and check if the backup is properly restored."""
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
# create a back up
assert db_service.backup_database() is True
@@ -107,7 +162,7 @@ def test_database_client_cannot_query_offline_database_server(uc2_network):
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
assert db_client.connected
assert len(db_client.connections)
assert db_client.query("SELECT") is True

86
tests/integration_tests/system/test_ntp_client_server.py Normal file
View File

@@ -0,0 +1,86 @@
from ipaddress import IPv4Address
from time import sleep
from typing import Tuple
import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.protocols.ntp import NTPPacket
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
# Create simple network for testing
# Define one node to be an NTP server and another node to be a NTP Client.
@pytest.fixture(scope="function")
def create_ntp_network(client_server) -> Tuple[NTPClient, Computer, NTPServer, Server]:
"""
+------------+ +------------+
| ntp | | ntp |
| client_1 +------------+ server_1 |
| | | |
+------------+ +------------+
"""
client, server = client_server
server.power_on()
server.software_manager.install(NTPServer)
ntp_server: NTPServer = server.software_manager.software.get("NTPServer")
ntp_server.start()
client.power_on()
client.software_manager.install(NTPClient)
ntp_client: NTPClient = client.software_manager.software.get("NTPClient")
ntp_client.start()
return ntp_client, client, ntp_server, server
def test_ntp_client_server(create_ntp_network):
ntp_client, client, ntp_server, server = create_ntp_network
ntp_server: NTPServer = server.software_manager.software["NTPServer"]
ntp_client: NTPClient = client.software_manager.software["NTPClient"]
assert ntp_server.operating_state == ServiceOperatingState.RUNNING
assert ntp_client.operating_state == ServiceOperatingState.RUNNING
ntp_client.configure(ntp_server_ip_address=IPv4Address("192.168.0.2"))
assert ntp_client.time is None
ntp_client.request_time()
assert ntp_client.time is not None
first_time = ntp_client.time
sleep(0.1)
ntp_client.apply_timestep(1) # Check time advances
second_time = ntp_client.time
assert first_time < second_time
# Test ntp client behaviour when ntp server is unavailable.
def test_ntp_server_failure(create_ntp_network):
ntp_client, client, ntp_server, server = create_ntp_network
ntp_server: NTPServer = server.software_manager.software["NTPServer"]
ntp_client: NTPClient = client.software_manager.software["NTPClient"]
assert ntp_client.operating_state == ServiceOperatingState.RUNNING
assert ntp_client.operating_state == ServiceOperatingState.RUNNING
ntp_client.configure(ntp_server_ip_address=IPv4Address("192.168.0.2"))
# Turn off ntp server.
ntp_server.stop()
assert ntp_server.operating_state == ServiceOperatingState.STOPPED
# And request a time update.
ntp_client.request_time()
assert ntp_client.time is None
# Restart ntp server.
ntp_server.start()
assert ntp_server.operating_state == ServiceOperatingState.RUNNING
ntp_client.request_time()
assert ntp_client.time is not None

6
tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py
View File

@@ -53,12 +53,12 @@ def test_node_os_scan(node, service, application):
# TODO implement processes
# add services to node
service.health_state_actual = SoftwareHealthState.COMPROMISED
service.set_health_state(SoftwareHealthState.COMPROMISED)
node.install_service(service=service)
assert service.health_state_visible == SoftwareHealthState.UNUSED
# add application to node
application.health_state_actual = SoftwareHealthState.COMPROMISED
application.set_health_state(SoftwareHealthState.COMPROMISED)
node.install_application(application=application)
assert application.health_state_visible == SoftwareHealthState.UNUSED
@@ -101,7 +101,7 @@ def test_node_red_scan(node, service, application):
assert service.revealed_to_red is False
# add application to node
application.health_state_actual = SoftwareHealthState.COMPROMISED
application.set_health_state(SoftwareHealthState.COMPROMISED)
node.install_application(application=application)
assert application.revealed_to_red is False

0
tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/__init__.py → tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/__init__.py
View File

4
tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py → tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_data_manipulation_bot.py
View File

@@ -5,7 +5,7 @@ from primaite.simulator.network.networks import arcd_uc2_network
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 ApplicationOperatingState
from primaite.simulator.system.services.red_services.data_manipulation_bot import (
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import (
DataManipulationAttackStage,
DataManipulationBot,
)
@@ -70,4 +70,4 @@ def test_dm_bot_perform_data_manipulation_success(dm_bot):
dm_bot._perform_data_manipulation(p_of_success=1.0)
assert dm_bot.attack_stage in (DataManipulationAttackStage.SUCCEEDED, DataManipulationAttackStage.FAILED)
assert dm_bot.connected
assert len(dm_bot.connections)

90
tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_dos_bot.py Normal file
View File

@@ -0,0 +1,90 @@
from ipaddress import IPv4Address
import pytest
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.red_applications.dos_bot import DoSAttackStage, DoSBot
@pytest.fixture(scope="function")
def dos_bot() -> DoSBot:
computer = Computer(
hostname="compromised_pc",
ip_address="192.168.0.1",
subnet_mask="255.255.255.0",
operating_state=NodeOperatingState.ON,
)
computer.software_manager.install(DoSBot)
dos_bot: DoSBot = computer.software_manager.software.get("DoSBot")
dos_bot.configure(target_ip_address=IPv4Address("192.168.0.1"))
dos_bot.set_original_state()
return dos_bot
def test_dos_bot_creation(dos_bot):
"""Test that the DoS bot is installed on a node."""
assert dos_bot is not None
def test_dos_bot_reset(dos_bot):
assert dos_bot.target_ip_address == IPv4Address("192.168.0.1")
assert dos_bot.target_port is Port.POSTGRES_SERVER
assert dos_bot.payload is None
assert dos_bot.repeat is False
dos_bot.configure(
target_ip_address=IPv4Address("192.168.1.1"), target_port=Port.HTTP, payload="payload", repeat=True
)
# should reset the relevant items
dos_bot.reset_component_for_episode(episode=0)
assert dos_bot.target_ip_address == IPv4Address("192.168.0.1")
assert dos_bot.target_port is Port.POSTGRES_SERVER
assert dos_bot.payload is None
assert dos_bot.repeat is False
dos_bot.configure(
target_ip_address=IPv4Address("192.168.1.1"), target_port=Port.HTTP, payload="payload", repeat=True
)
dos_bot.set_original_state()
dos_bot.reset_component_for_episode(episode=1)
# should reset to the configured value
assert dos_bot.target_ip_address == IPv4Address("192.168.1.1")
assert dos_bot.target_port is Port.HTTP
assert dos_bot.payload == "payload"
assert dos_bot.repeat is True
def test_dos_bot_cannot_run_when_node_offline(dos_bot):
dos_bot_node: Computer = dos_bot.parent
assert dos_bot_node.operating_state is NodeOperatingState.ON
dos_bot_node.power_off()
for i in range(dos_bot_node.shut_down_duration + 1):
dos_bot_node.apply_timestep(timestep=i)
assert dos_bot_node.operating_state is NodeOperatingState.OFF
dos_bot._application_loop()
# assert not run
assert dos_bot.attack_stage is DoSAttackStage.NOT_STARTED
def test_dos_bot_not_configured(dos_bot):
dos_bot.target_ip_address = None
dos_bot.operating_state = ApplicationOperatingState.RUNNING
dos_bot._application_loop()
def test_dos_bot_perform_port_scan(dos_bot):
dos_bot._perform_port_scan(p_of_success=1)
assert dos_bot.attack_stage is DoSAttackStage.PORT_SCAN

0
tests/unit_tests/_primaite/_simulator/_system/_applications/test_application_actions.py Normal file
View File

50
tests/unit_tests/_primaite/_simulator/_system/_applications/test_applications.py Normal file
View File

@@ -0,0 +1,50 @@
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.software import SoftwareHealthState
def test_scan(application):
assert application.operating_state == ApplicationOperatingState.CLOSED
assert application.health_state_visible == SoftwareHealthState.UNUSED
application.run()
assert application.operating_state == ApplicationOperatingState.RUNNING
assert application.health_state_visible == SoftwareHealthState.UNUSED
application.scan()
assert application.operating_state == ApplicationOperatingState.RUNNING
assert application.health_state_visible == SoftwareHealthState.GOOD
def test_run_application(application):
assert application.operating_state == ApplicationOperatingState.CLOSED
assert application.health_state_actual == SoftwareHealthState.UNUSED
application.run()
assert application.operating_state == ApplicationOperatingState.RUNNING
assert application.health_state_actual == SoftwareHealthState.GOOD
def test_close_application(application):
application.run()
assert application.operating_state == ApplicationOperatingState.RUNNING
assert application.health_state_actual == SoftwareHealthState.GOOD
application.close()
assert application.operating_state == ApplicationOperatingState.CLOSED
assert application.health_state_actual == SoftwareHealthState.GOOD
def test_application_describe_states(application):
assert application.operating_state == ApplicationOperatingState.CLOSED
assert application.health_state_actual == SoftwareHealthState.UNUSED
assert SoftwareHealthState.UNUSED.value == application.describe_state().get("health_state_actual")
application.run()
assert SoftwareHealthState.GOOD.value == application.describe_state().get("health_state_actual")
application.set_health_state(SoftwareHealthState.COMPROMISED)
assert SoftwareHealthState.COMPROMISED.value == application.describe_state().get("health_state_actual")
application.patch()
assert SoftwareHealthState.PATCHING.value == application.describe_state().get("health_state_actual")

30
tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py
View File

@@ -1,5 +1,6 @@
from ipaddress import IPv4Address
from typing import Tuple, Union
from uuid import uuid4
import pytest
@@ -62,18 +63,26 @@ def test_disconnect_when_client_is_closed(database_client_on_computer):
def test_disconnect(database_client_on_computer):
"""Database client should set connected to False and remove the database server ip address."""
"""Database client should remove the connection."""
database_client, computer = database_client_on_computer
database_client.connected = True
database_client._connections[str(uuid4())] = {"item": True}
assert len(database_client.connections) == 1
assert database_client.operating_state is ApplicationOperatingState.RUNNING
assert database_client.server_ip_address is not None
database_client.disconnect()
assert database_client.connected is False
assert database_client.server_ip_address is None
assert len(database_client.connections) == 0
uuid = str(uuid4())
database_client._connections[uuid] = {"item": True}
assert len(database_client.connections) == 1
database_client.disconnect(connection_id=uuid)
assert len(database_client.connections) == 0
def test_query_when_client_is_closed(database_client_on_computer):
@@ -86,19 +95,6 @@ def test_query_when_client_is_closed(database_client_on_computer):
assert database_client.query(sql="test") is False
def test_query_failed_reattempt(database_client_on_computer):
"""Database client query should return False if the reattempt fails."""
database_client, computer = database_client_on_computer
def return_false():
return False
database_client.connect = return_false
database_client.connected = False
assert database_client.query(sql="test", is_reattempt=True) is False
def test_query_fail_to_connect(database_client_on_computer):
"""Database client query should return False if the connect attempt fails."""
database_client, computer = database_client_on_computer

124
tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py
View File

@@ -1,3 +1,5 @@
from uuid import uuid4
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.software import SoftwareHealthState
@@ -17,52 +19,174 @@ def test_scan(service):
def test_start_service(service):
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
service.start()
assert service.operating_state == ServiceOperatingState.RUNNING
assert service.health_state_actual == SoftwareHealthState.GOOD
def test_stop_service(service):
service.start()
assert service.operating_state == ServiceOperatingState.RUNNING
assert service.health_state_actual == SoftwareHealthState.GOOD
service.stop()
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.GOOD
def test_pause_and_resume_service(service):
assert service.operating_state == ServiceOperatingState.STOPPED
service.resume()
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
service.start()
assert service.health_state_actual == SoftwareHealthState.GOOD
service.pause()
assert service.operating_state == ServiceOperatingState.PAUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
service.resume()
assert service.operating_state == ServiceOperatingState.RUNNING
assert service.health_state_actual == SoftwareHealthState.GOOD
def test_restart(service):
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
service.restart()
# Service is STOPPED. Restart will only work if the service was PAUSED or RUNNING
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
service.start()
assert service.operating_state == ServiceOperatingState.RUNNING
assert service.health_state_actual == SoftwareHealthState.GOOD
service.restart()
# Service is RUNNING. Restart should work
assert service.operating_state == ServiceOperatingState.RESTARTING
assert service.health_state_actual == SoftwareHealthState.GOOD
timestep = 0
while service.operating_state == ServiceOperatingState.RESTARTING:
service.apply_timestep(timestep)
assert service.health_state_actual == SoftwareHealthState.GOOD
timestep += 1
assert service.operating_state == ServiceOperatingState.RUNNING
assert service.health_state_actual == SoftwareHealthState.GOOD
def test_restart_compromised(service):
service.start()
assert service.health_state_actual == SoftwareHealthState.GOOD
# compromise the service
service.set_health_state(SoftwareHealthState.COMPROMISED)
service.restart()
assert service.operating_state == ServiceOperatingState.RESTARTING
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
"""
Service should be compromised even after reset.
Only way to remove compromised status is via patching.
"""
timestep = 0
while service.operating_state == ServiceOperatingState.RESTARTING:
service.apply_timestep(timestep)
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
timestep += 1
assert service.operating_state == ServiceOperatingState.RUNNING
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
def test_compromised_service_remains_compromised(service):
"""
Tests that a compromised service stays compromised.
The only way that the service can be uncompromised is by running patch.
"""
service.start()
assert service.health_state_actual == SoftwareHealthState.GOOD
service.set_health_state(SoftwareHealthState.COMPROMISED)
service.stop()
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
service.start()
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
service.disable()
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
service.enable()
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
service.pause()
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
service.resume()
assert service.health_state_actual == SoftwareHealthState.COMPROMISED
def test_service_patching(service):
service.start()
assert service.health_state_actual == SoftwareHealthState.GOOD
service.set_health_state(SoftwareHealthState.COMPROMISED)
service.patch()
assert service.health_state_actual == SoftwareHealthState.PATCHING
for i in range(service.patching_duration + 1):
service.apply_timestep(i)
assert service.health_state_actual == SoftwareHealthState.GOOD
def test_enable_disable(service):
service.disable()
assert service.operating_state == ServiceOperatingState.DISABLED
assert service.health_state_actual == SoftwareHealthState.UNUSED
service.enable()
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
def test_overwhelm_service(service):
service.max_sessions = 2
service.start()
uuid = str(uuid4())
assert service.add_connection(connection_id=uuid) # should be true
assert service.health_state_actual == SoftwareHealthState.GOOD
assert not service.add_connection(connection_id=uuid) # fails because connection already exists
assert service.health_state_actual == SoftwareHealthState.GOOD
assert service.add_connection(connection_id=str(uuid4())) # succeed
assert service.health_state_actual == SoftwareHealthState.GOOD
assert not service.add_connection(connection_id=str(uuid4())) # fail because at capacity
assert service.health_state_actual is SoftwareHealthState.OVERWHELMED
def test_create_and_remove_connections(service):
service.start()
uuid = str(uuid4())
assert service.add_connection(connection_id=uuid) # should be true
assert len(service.connections) == 1
assert service.health_state_actual is SoftwareHealthState.GOOD
assert service.remove_connection(connection_id=uuid) # should be true
assert len(service.connections) == 0
assert service.health_state_actual is SoftwareHealthState.GOOD

29
tests/unit_tests/_primaite/_simulator/_system/test_software.py Normal file
View File

@@ -0,0 +1,29 @@
from typing import Dict
import pytest
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.software import Software, SoftwareHealthState
class TestSoftware(Software):
def describe_state(self) -> Dict:
pass
@pytest.fixture(scope="function")
def software(file_system):
return TestSoftware(
name="TestSoftware", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="test_service")
)
def test_software_creation(software):
assert software is not None
def test_software_set_health_state(software):
assert software.health_state_actual == SoftwareHealthState.UNUSED
software.set_health_state(SoftwareHealthState.GOOD)
assert software.health_state_actual == SoftwareHealthState.GOOD