#2700 Add DatabaseConfigure action

This commit is contained in:
Marek Wolan
2024-07-01 11:31:27 +01:00
parent e7fc01781e
commit c34cb6d7ce
4 changed files with 123 additions and 4 deletions

View File

@@ -17,6 +17,7 @@ from gymnasium import spaces
from pydantic import BaseModel, Field, field_validator, ValidationInfo from pydantic import BaseModel, Field, field_validator, ValidationInfo
from primaite import getLogger from primaite import getLogger
from primaite.interface.request import RequestFormat
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
@@ -245,6 +246,27 @@ class NodeApplicationInstallAction(AbstractAction):
] ]
class ConfigureDatabaseClientAction(AbstractAction):
"""Action which sets config parameters for a database client on a node."""
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, options: Dict) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
if node_name is None:
return ["do_nothing"]
ConfigureDatabaseClientAction._Opts.model_validate(options) # check that options adhere to schema
return ["network", "node", node_name, "application", "DatabaseClient", "configure", options]
class NodeApplicationRemoveAction(AbstractAction): class NodeApplicationRemoveAction(AbstractAction):
"""Action which removes/uninstalls an application.""" """Action which removes/uninstalls an application."""
@@ -1045,6 +1067,7 @@ class ActionManager:
"NODE_NMAP_PING_SCAN": NodeNMAPPingScanAction, "NODE_NMAP_PING_SCAN": NodeNMAPPingScanAction,
"NODE_NMAP_PORT_SCAN": NodeNMAPPortScanAction, "NODE_NMAP_PORT_SCAN": NodeNMAPPortScanAction,
"NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction, "NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction,
"CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction,
} }
"""Dictionary which maps action type strings to the corresponding action class.""" """Dictionary which maps action type strings to the corresponding action class."""

View File

@@ -8,13 +8,14 @@ from uuid import uuid4
from prettytable import MARKDOWN, PrettyTable from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel from pydantic import BaseModel
from primaite.interface.request import RequestResponse from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.utils.validators import IPV4Address
class DatabaseClientConnection(BaseModel): class DatabaseClientConnection(BaseModel):
@@ -96,6 +97,14 @@ class DatabaseClient(Application):
""" """
rm = super()._init_request_manager() rm = super()._init_request_manager()
rm.add_request("execute", RequestType(func=lambda request, context: RequestResponse.from_bool(self.execute()))) rm.add_request("execute", RequestType(func=lambda request, context: RequestResponse.from_bool(self.execute())))
def _configure(request: RequestFormat, context: Dict) -> RequestResponse:
ip, pw = request[-1].get("server_ip_address"), request[-1].get("server_password")
ip = None if ip is None else IPV4Address(ip)
success = self.configure(server_ip_address=ip, server_password=pw)
return RequestResponse.from_bool(success)
rm.add_request("configure", RequestType(func=lambda request, context: _configure(request, context)))
return rm return rm
def execute(self) -> bool: def execute(self) -> bool:
@@ -141,16 +150,17 @@ class DatabaseClient(Application):
table.add_row([connection_id, connection.is_active]) table.add_row([connection_id, connection.is_active])
print(table.get_string(sortby="Connection ID")) print(table.get_string(sortby="Connection ID"))
def configure(self, server_ip_address: IPv4Address, server_password: Optional[str] = None): def configure(self, server_ip_address: Optional[IPv4Address] = None, server_password: Optional[str] = None) -> bool:
""" """
Configure the DatabaseClient to communicate with a DatabaseService. Configure the DatabaseClient to communicate with a DatabaseService.
:param server_ip_address: The IP address of the Node the DatabaseService is on. :param server_ip_address: The IP address of the Node the DatabaseService is on.
:param server_password: The password on the DatabaseService. :param server_password: The password on the DatabaseService.
""" """
self.server_ip_address = server_ip_address self.server_ip_address = server_ip_address or self.server_ip_address
self.server_password = server_password self.server_password = server_password or self.server_password
self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.") self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.")
return True
def connect(self) -> bool: def connect(self) -> bool:
"""Connect the native client connection.""" """Connect the native client connection."""

View File

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

View File

@@ -0,0 +1,85 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from primaite.game.agent.actions import ConfigureDatabaseClientAction
from primaite.simulator.system.applications.database_client import DatabaseClient
from tests.conftest import ControlledAgent
class TestConfigureDatabaseAction:
def test_configure_ip_password(self, game_and_agent):
game, agent = game_and_agent
agent: ControlledAgent
agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager)
# make sure there is a database client on this node
client_1 = game.simulation.network.get_node_by_hostname("client_1")
client_1.software_manager.install(DatabaseClient)
db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"]
action = (
"CONFIGURE_DATABASE_CLIENT",
{
"node_id": 0,
"options": {
"server_ip_address": "192.168.1.99",
"server_password": "admin123",
},
},
)
agent.store_action(action)
game.step()
assert db_client.server_ip_address == IPv4Address("192.168.1.99")
assert db_client.server_password == "admin123"
def test_configure_ip(self, game_and_agent):
game, agent = game_and_agent
agent: ControlledAgent
agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager)
# make sure there is a database client on this node
client_1 = game.simulation.network.get_node_by_hostname("client_1")
client_1.software_manager.install(DatabaseClient)
db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"]
action = (
"CONFIGURE_DATABASE_CLIENT",
{
"node_id": 0,
"options": {
"server_ip_address": "192.168.1.99",
},
},
)
agent.store_action(action)
game.step()
assert db_client.server_ip_address == IPv4Address("192.168.1.99")
assert db_client.server_password is None
def test_configure_password(self, game_and_agent):
game, agent = game_and_agent
agent: ControlledAgent
agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager)
# make sure there is a database client on this node
client_1 = game.simulation.network.get_node_by_hostname("client_1")
client_1.software_manager.install(DatabaseClient)
db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"]
old_ip = db_client.server_ip_address
action = (
"CONFIGURE_DATABASE_CLIENT",
{
"node_id": 0,
"options": {
"server_password": "admin123",
},
},
)
agent.store_action(action)
game.step()
assert db_client.server_ip_address == old_ip
assert db_client.server_password is "admin123"