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

#2258: setting up test that verifies game config parsing

Browse Source
This commit is contained in:
Czar Echavez
2024-02-07 18:09:54 +00:00
parent cfc254753f
commit f21ee857a7
3 changed files with 213 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/primaite/game/game.py
View File

@@ -31,6 +31,23 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
_LOGGER = getLogger(__name__)
APPLICATION_TYPES_MAPPING = {
"WebBrowser": WebBrowser,
"DataManipulationBot": DataManipulationBot,
}
SERVICE_TYPES_MAPPING = {
"DNSClient": DNSClient,
"DNSServer": DNSServer,
"DatabaseClient": DatabaseClient,
"DatabaseService": DatabaseService,
"WebServer": WebServer,
"FTPClient": FTPClient,
"FTPServer": FTPServer,
"NTPClient": NTPClient,
"NTPServer": NTPServer,
}
class PrimaiteGameOptions(BaseModel):
"""
@@ -238,20 +255,9 @@ class PrimaiteGame:
new_service = None
service_ref = service_cfg["ref"]
service_type = service_cfg["type"]
service_types_mapping = {
"DNSClient": DNSClient, # key is equal to the 'name' attr of the service class itself.
"DNSServer": DNSServer,
"DatabaseClient": DatabaseClient,
"DatabaseService": DatabaseService,
"WebServer": WebServer,
"FTPClient": FTPClient,
"FTPServer": FTPServer,
"NTPClient": NTPClient,
"NTPServer": NTPServer,
}
if service_type in service_types_mapping:
if service_type in SERVICE_TYPES_MAPPING:
_LOGGER.debug(f"installing {service_type} on node {new_node.hostname}")
new_node.software_manager.install(service_types_mapping[service_type])
new_node.software_manager.install(SERVICE_TYPES_MAPPING[service_type])
new_service = new_node.software_manager.software[service_type]
game.ref_map_services[service_ref] = new_service.uuid
else:
@@ -280,12 +286,9 @@ class PrimaiteGame:
new_application = None
application_ref = application_cfg["ref"]
application_type = application_cfg["type"]
application_types_mapping = {
"WebBrowser": WebBrowser,
"DataManipulationBot": DataManipulationBot,
}
if application_type in application_types_mapping:
new_node.software_manager.install(application_types_mapping[application_type])
if application_type in APPLICATION_TYPES_MAPPING:
new_node.software_manager.install(APPLICATION_TYPES_MAPPING[application_type])
new_application = new_node.software_manager.software[application_type]
game.ref_map_applications[application_ref] = new_application.uuid
else:

114
tests/assets/configs/basic_switched_network.yaml Normal file
View File

@@ -0,0 +1,114 @@
training_config:
rl_framework: SB3
rl_algorithm: PPO
seed: 333
n_learn_episodes: 1
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
save_step_metadata: false
save_pcap_logs: true
save_sys_logs: true
game:
max_episode_length: 256
ports:
- ARP
- DNS
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
agents:
- ref: client_2_green_user
team: GREEN
type: GreenWebBrowsingAgent
observation_space:
type: UC2GreenObservation
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_2
applications:
- application_name: WebBrowser
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
reward_function:
reward_components:
- type: DUMMY
agent_settings:
start_settings:
start_step: 5
frequency: 4
variance: 3
simulation:
network:
nodes:
- ref: switch_1
type: switch
hostname: switch_1
num_ports: 8
- ref: client_1
type: computer
hostname: client_1
ip_address: 192.168.10.21
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- ref: client_1_web_browser
type: WebBrowser
options:
target_url: http://arcd.com/users/
- ref: data_manipulation_bot
type: DataManipulationBot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
services:
- ref: client_1_dns_client
type: DNSClient
- ref: client_2
type: computer
hostname: client_2
ip_address: 192.168.10.22
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
# pre installed services and applications
links:
- ref: switch_1___client_1
endpoint_a_ref: switch_1
endpoint_a_port: 1
endpoint_b_ref: client_1
endpoint_b_port: 1
- ref: switch_1___client_2
endpoint_a_ref: switch_1
endpoint_a_port: 2
endpoint_b_ref: client_2
endpoint_b_port: 1

77
tests/integration_tests/game_configuration.py Normal file
View File

@@ -0,0 +1,77 @@
from pathlib import Path
from typing import Union
import yaml
from primaite.config.load import example_config_path
from primaite.game.agent.data_manipulation_bot import DataManipulationAgent
from primaite.game.agent.interface import ProxyAgent, RandomAgent
from primaite.game.game import APPLICATION_TYPES_MAPPING, PrimaiteGame, SERVICE_TYPES_MAPPING
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from tests import TEST_ASSETS_ROOT
BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml"
def load_config(config_path: Union[str, Path]) -> PrimaiteGame:
"""Returns a PrimaiteGame object which loads the contents of a given yaml path."""
with open(config_path, "r") as f:
cfg = yaml.safe_load(f)
return PrimaiteGame.from_config(cfg)
def test_example_config():
"""Test that the example config can be parsed properly."""
game = load_config(example_config_path())
assert len(game.agents) == 3 # red, blue and green agent
# green agent
assert game.agents[0].agent_name == "client_2_green_user"
assert isinstance(game.agents[0], RandomAgent)
# red agent
assert game.agents[1].agent_name == "client_1_data_manipulation_red_bot"
assert isinstance(game.agents[1], DataManipulationAgent)
# blue agent
assert game.agents[2].agent_name == "defender"
assert isinstance(game.agents[2], ProxyAgent)
network: Network = game.simulation.network
assert len(network.nodes) == 10 # 10 nodes in example network
assert len(network.routers) == 1 # 1 router in network
assert len(network.switches) == 2 # 2 switches in network
assert len(network.servers) == 5 # 5 servers in network
def test_node_software_install():
"""Test that software can be installed on a node."""
game = load_config(BASIC_CONFIG)
client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
client_2: Computer = game.simulation.network.get_node_by_hostname("client_2")
system_software = {DNSClient, FTPClient, WebBrowser}
# check that system software is installed on client 1
for software in system_software:
assert client_1.software_manager.software.get(software.__name__) is not None
# check that system software is installed on client 2
for software in system_software:
assert client_2.software_manager.software.get(software.__name__) is not None
# check that applications have been installed on client 1
for applications in APPLICATION_TYPES_MAPPING:
assert client_1.software_manager.software.get(applications) is not None
# check that services have been installed on client 1
# for service in SERVICE_TYPES_MAPPING:
# assert client_1.software_manager.software.get(service) is not None