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

#1816 - Updated the DataManipulationBot to subclass DatabaseClient. Extended logging. Dropped the Link loading logging as it was clogging up the terminal output.

Browse Source
This commit is contained in:
Chris McCarthy
2023-09-11 16:15:03 +01:00
parent f19dc9892b
commit 695b3ceab4
22 changed files with 268 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/index.rst
View File

@@ -98,6 +98,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE!
source/getting_started
source/about
source/config
source/simulation
source/primaite_session
source/custom_agent
PrimAITE API <source/_autosummary/primaite>

3
docs/source/simulation.rst
View File

@@ -21,4 +21,5 @@ Contents
simulation_components/network/router
simulation_components/network/switch
simulation_components/network/network
simulation_components/internal_frame_processing
simulation_components/system/internal_frame_processing
simulation_components/system/software

2
docs/source/simulation_components/network/network.rst
View File

@@ -2,7 +2,7 @@
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _about:
.. _network:
Network
=======

2
docs/source/simulation_components/network/router.rst
View File

@@ -2,7 +2,7 @@
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _about:
.. _router:
Router Module
=============

58
docs/source/simulation_components/system/data_manipulation_bot.rst Normal file
View File

@@ -0,0 +1,58 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
DataManipulationBot
===================
The ``DataManipulationBot`` class provides functionality to connect to a ``DatabaseService`` and execute malicious SQL statements.
Overview
--------
The bot is intended to simulate a malicious actor carrying out attacks like:
- Dropping tables
- Deleting records
- Modifying data
On a database server by abusing an application's trusted database connectivity.
Usage
-----
- Create an instance and call ``configure`` to set:
- Target database server IP
- Database password (if needed)
- SQL statement payload
- Call ``run`` to connect and execute the statement.
The bot handles connecting, executing the statement, and disconnecting.
Example
-------
.. code-block:: python
client_1 = Computer(
hostname="client_1", ip_address="192.168.10.21", subnet_mask="255.255.255.0", default_gateway="192.168.10.1"
)
client_1.power_on()
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
client_1.software_manager.install(DataManipulationBot)
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
data_manipulation_bot.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DROP TABLE IF EXISTS user;")
data_manipulation_bot.run()
This would connect to the database service at 192.168.1.14, authenticate, and execute the SQL statement to drop the 'users' table.
Implementation
--------------
The bot extends ``DatabaseClient`` and leverages its connectivity.
- Uses the Application base class for lifecycle management.
- Credentials and target IP set via ``configure``.
- ``run`` handles connecting, executing statement, and disconnecting.
- SQL payload executed via ``query`` method.
- Results in malicious SQL being executed on remote database server.

0
docs/source/simulation_components/network/database_client_server.rst → docs/source/simulation_components/system/database_client_server.rst
View File

2
docs/source/simulation_components/network/internal_frame_processing.rst → docs/source/simulation_components/system/internal_frame_processing.rst
View File

@@ -2,7 +2,7 @@
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
.. _about:
.. _internal_frame_processing:
Internal Frame Processing
=========================

1
docs/source/simulation_components/network/software.rst → docs/source/simulation_components/system/software.rst
View File

@@ -16,3 +16,4 @@ Contents
:maxdepth: 8
database_client_server
data_manipulation_bot

4
src/primaite/simulator/network/container.py
View File

@@ -184,7 +184,7 @@ class Network(SimComponent):
self._node_id_map[len(self.nodes)] = node
node.parent = self
self._nx_graph.add_node(node.hostname)
_LOGGER.info(f"Added node {node.uuid} to Network {self.uuid}")
_LOGGER.debug(f"Added node {node.uuid} to Network {self.uuid}")
def get_node_by_hostname(self, hostname: str) -> Optional[Node]:
"""
@@ -245,7 +245,7 @@ class Network(SimComponent):
self._link_id_map[len(self.links)] = link
self._nx_graph.add_edge(endpoint_a.parent.hostname, endpoint_b.parent.hostname)
link.parent = self
_LOGGER.info(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}")
_LOGGER.debug(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}")
def remove_link(self, link: Link) -> None:
"""Disconnect a link from the network.

20
src/primaite/simulator/network/hardware/base.py
View File

@@ -186,7 +186,7 @@ class NIC(SimComponent):
if self.connected_node:
self.connected_node.sys_log.info(f"NIC {self} disabled")
else:
_LOGGER.info(f"NIC {self} disabled")
_LOGGER.debug(f"NIC {self} disabled")
if self.connected_link:
self.connected_link.endpoint_down()
@@ -208,7 +208,7 @@ class NIC(SimComponent):
# TODO: Inform the Node that a link has been connected
self.connected_link = link
self.enable()
_LOGGER.info(f"NIC {self} connected to Link {link}")
_LOGGER.debug(f"NIC {self} connected to Link {link}")
def disconnect_link(self):
"""Disconnect the NIC from the connected Link."""
@@ -351,7 +351,7 @@ class SwitchPort(SimComponent):
if self.connected_node:
self.connected_node.sys_log.info(f"SwitchPort {self} disabled")
else:
_LOGGER.info(f"SwitchPort {self} disabled")
_LOGGER.debug(f"SwitchPort {self} disabled")
if self.connected_link:
self.connected_link.endpoint_down()
@@ -371,7 +371,7 @@ class SwitchPort(SimComponent):
# TODO: Inform the Switch that a link has been connected
self.connected_link = link
_LOGGER.info(f"SwitchPort {self} connected to Link {link}")
_LOGGER.debug(f"SwitchPort {self} connected to Link {link}")
self.enable()
def disconnect_link(self):
@@ -477,13 +477,13 @@ class Link(SimComponent):
def endpoint_up(self):
"""Let the Link know and endpoint has been brought up."""
if self.is_up:
_LOGGER.info(f"Link {self} up")
_LOGGER.debug(f"Link {self} up")
def endpoint_down(self):
"""Let the Link know and endpoint has been brought down."""
if not self.is_up:
self.current_load = 0.0
_LOGGER.info(f"Link {self} down")
_LOGGER.debug(f"Link {self} down")
@property
def is_up(self) -> bool:
@@ -510,7 +510,7 @@ class Link(SimComponent):
"""
can_transmit = self._can_transmit(frame)
if not can_transmit:
_LOGGER.info(f"Cannot transmit frame as {self} is at capacity")
_LOGGER.debug(f"Cannot transmit frame as {self} is at capacity")
return False
receiver = self.endpoint_a
@@ -522,7 +522,7 @@ class Link(SimComponent):
# Frame transmitted successfully
# Load the frame size on the link
self.current_load += frame_size
_LOGGER.info(
_LOGGER.debug(
f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits "
f"({self.current_load_percent})"
)
@@ -1148,7 +1148,7 @@ class Node(SimComponent):
service.parent = self
service.install() # Perform any additional setup, such as creating files for this service on the node.
self.sys_log.info(f"Installed service {service.name}")
_LOGGER.info(f"Added service {service.uuid} to node {self.uuid}")
_LOGGER.debug(f"Added service {service.uuid} to node {self.uuid}")
def uninstall_service(self, service: Service) -> None:
"""Uninstall and completely remove service from this node.
@@ -1163,7 +1163,7 @@ class Node(SimComponent):
self.services.pop(service.uuid)
service.parent = None
self.sys_log.info(f"Uninstalled service {service.name}")
_LOGGER.info(f"Removed service {service.uuid} from node {self.uuid}")
_LOGGER.debug(f"Removed service {service.uuid} from node {self.uuid}")
def __contains__(self, item: Any) -> bool:
if isinstance(item, Service):

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

@@ -1,3 +1,5 @@
from ipaddress import IPv4Address
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.base import NIC
from primaite.simulator.network.hardware.nodes.computer import Computer
@@ -8,6 +10,7 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.database_service import DatabaseService
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
def client_server_routed() -> Network:
@@ -127,6 +130,9 @@ def arcd_uc2_network() -> Network:
)
client_1.power_on()
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
db_manipulation_bot.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DROP TABLE IF EXISTS user;")
# Client 2
client_2 = Computer(
@@ -145,16 +151,6 @@ def arcd_uc2_network() -> Network:
domain_controller.power_on()
network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1])
# Web Server
web_server = Server(
hostname="web_server", ip_address="192.168.1.12", subnet_mask="255.255.255.0", default_gateway="192.168.1.1"
)
web_server.power_on()
web_server.software_manager.install(DatabaseClient)
database_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
database_client.run()
network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])
# Database Server
database_server = Server(
hostname="database_server",
@@ -194,9 +190,21 @@ def arcd_uc2_network() -> Network:
database_server.software_manager.install(DatabaseService)
database_service: DatabaseService = database_server.software_manager.software["DatabaseService"] # noqa
database_service.start()
database_service._process_sql(ddl) # noqa
database_service._process_sql(ddl, None) # noqa
for insert_statement in user_insert_statements:
database_service._process_sql(insert_statement) # noqa
database_service._process_sql(insert_statement, None) # noqa
# Web Server
web_server = Server(
hostname="web_server", ip_address="192.168.1.12", subnet_mask="255.255.255.0", default_gateway="192.168.1.1"
)
web_server.power_on()
web_server.software_manager.install(DatabaseClient)
database_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
database_client.configure(server_ip_address=IPv4Address("192.168.1.14"))
network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])
database_client.run()
database_client.connect()
# Backup Server
backup_server = Server(

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

@@ -1,11 +1,12 @@
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from uuid import uuid4
from prettytable import PrettyTable
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.applications.application import Application, ApplicationOperatingState
from primaite.simulator.system.core.software_manager import SoftwareManager
@@ -20,7 +21,9 @@ class DatabaseClient(Application):
"""
server_ip_address: Optional[IPv4Address] = None
server_password: Optional[str] = None
connected: bool = False
_query_success_tracker: Dict[str, bool] = {}
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
@@ -37,15 +40,22 @@ class DatabaseClient(Application):
pass
return super().describe_state()
def connect(self, server_ip_address: IPv4Address, password: Optional[str] = None) -> bool:
def configure(self, server_ip_address: IPv4Address, server_password: Optional[str] = None):
"""
Connect to a Database Service.
Configure the DatabaseClient to communicate with a DatabaseService.
:param server_ip_address: The IPv4 Address of the Node the Database Service is running on.
:param password: The Database Service password. Is optional and has a default value of None.
:param server_ip_address: The IP address of the Node the DatabaseService is on.
:param server_password: The password on the DatabaseService.
"""
self.server_ip_address = server_ip_address
self.server_password = server_password
self.sys_log.info(f"Configured the {self.name} with {server_ip_address=}, {server_password=}.")
def connect(self) -> bool:
"""Connect to a Database Service."""
if not self.connected and self.operating_state.RUNNING:
return self._connect(server_ip_address, password)
return self._connect(self.server_ip_address, self.server_password)
return False
def _connect(
self, server_ip_address: IPv4Address, password: Optional[str] = None, is_reattempt: bool = False
@@ -75,18 +85,42 @@ class DatabaseClient(Application):
self.sys_log.info(f"DatabaseClient disconnected from {self.server_ip_address}")
self.server_ip_address = None
self.connected = False
def query(self, sql: str):
def _query(self, sql: str, query_id: str, is_reattempt: bool = False) -> bool:
if is_reattempt:
success = self._query_success_tracker.get(query_id)
if success:
return True
return False
else:
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "sql", "sql": sql, "uuid": query_id},
dest_ip_address=self.server_ip_address,
dest_port=self.port,
)
return self._query(sql=sql, query_id=query_id, is_reattempt=True)
def run(self) -> None:
"""Run the DatabaseClient."""
super().run()
self.operating_state = ApplicationOperatingState.RUNNING
self.connect()
def query(self, sql: str) -> bool:
"""
Send a query to the Database Service.
:param sql: The SQL query.
:return: True if the query was successful, otherwise False.
"""
if self.connected and self.operating_state.RUNNING:
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "sql", "sql": sql}, dest_ip_address=self.server_ip_address, dest_port=self.port
)
query_id = str(uuid4())
# Initialise the tracker of this ID to False
self._query_success_tracker[query_id] = False
return self._query(sql=sql, query_id=query_id)
def _print_data(self, data: Dict):
"""
@@ -94,13 +128,14 @@ class DatabaseClient(Application):
:param markdown: Whether to display the table in Markdown format or not. Default is `False`.
"""
table = PrettyTable(list(data.values())[0])
if data:
table = PrettyTable(list(data.values())[0])
table.align = "l"
table.title = f"{self.sys_log.hostname} Database Client"
for row in data.values():
table.add_row(row.values())
print(table)
table.align = "l"
table.title = f"{self.sys_log.hostname} Database Client"
for row in data.values():
table.add_row(row.values())
print(table)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
@@ -114,5 +149,9 @@ class DatabaseClient(Application):
if payload["type"] == "connect_response":
self.connected = payload["response"] == True
elif payload["type"] == "sql":
self._print_data(payload["data"])
query_id = payload.get("uuid")
status_code = payload.get("status_code")
self._query_success_tracker[query_id] = status_code == 200
if self._query_success_tracker[query_id]:
self._print_data(payload["data"])
return True

11
src/primaite/simulator/system/services/database_service.py
View File

@@ -81,20 +81,21 @@ class DatabaseService(Service):
status_code = 404 # service not found
return {"status_code": status_code, "type": "connect_response", "response": status_code == 200}
def _process_sql(self, query: str) -> Dict[str, Union[int, List[Any]]]:
def _process_sql(self, query: str, query_id: str) -> Dict[str, Union[int, List[Any]]]:
"""
Executes the given SQL query and returns the result.
:param query: The SQL query to be executed.
:return: Dictionary containing status code and data fetched.
"""
self.sys_log.info(f"{self.name}: Running {query}")
try:
self._cursor.execute(query)
self._conn.commit()
except OperationalError:
# Handle the case where the table does not exist.
return {"status_code": 404, "data": []}
self.sys_log.error(f"{self.name}: Error, query failed")
return {"status_code": 404, "data": {}}
data = []
description = self._cursor.description
if description:
@@ -104,7 +105,7 @@ class DatabaseService(Service):
data = self._cursor.fetchall()
if data and headers:
data = {row[0]: {header: value for header, value in zip(headers, row)} for row in data}
return {"status_code": 200, "type": "sql", "data": data}
return {"status_code": 200, "type": "sql", "data": data, "uuid": query_id}
def describe_state(self) -> Dict:
"""
@@ -134,7 +135,7 @@ class DatabaseService(Service):
self.connections.pop(session_id)
elif payload["type"] == "sql":
if session_id in self.connections:
result = self._process_sql(payload.get("sql"))
result = self._process_sql(query=payload["sql"], query_id=payload["uuid"])
else:
result = {"status_code": 401, "type": "sql"}
self.send(payload=result, session_id=session_id)

0
src/primaite/simulator/system/services/red_services/__init__.py Normal file
View File

49
src/primaite/simulator/system/services/red_services/data_manipulation_bot.py Normal file
View File

@@ -0,0 +1,49 @@
from ipaddress import IPv4Address
from typing import Optional
from primaite.simulator.system.applications.database_client import DatabaseClient
class DataManipulationBot(DatabaseClient):
"""
Red Agent Data Integration Service.
The Service represents a bot that causes files/folders in the File System to
become corrupted.
"""
server_ip_address: Optional[IPv4Address] = None
payload: Optional[str] = None
server_password: Optional[str] = None
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DataManipulationBot"
def configure(
self, server_ip_address: IPv4Address, server_password: Optional[str] = None, payload: Optional[str] = None
):
"""
Configure the DataManipulatorBot to communicate with a DatabaseService.
:param server_ip_address: The IP address of the Node the DatabaseService is on.
:param server_password: The password on the DatabaseService.
:param payload: The data manipulation query payload.
"""
self.server_ip_address = server_ip_address
self.payload = payload
self.server_password = server_password
self.sys_log.info(f"Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}.")
def run(self):
"""Run the DataManipulationBot."""
if self.server_ip_address and self.payload:
self.sys_log.info(f"Attempting to start the {self.name}")
super().run()
if not self.connected:
self.connect()
if self.connected:
self.query(self.payload)
self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
else:
self.sys_log.error(f"Failed to start the {self.name} as it requires both a target_io_address and payload.")

34
src/primaite/simulator/system/services/red_services/data_manipulator_service.py
View File

@@ -1,34 +0,0 @@
from ipaddress import IPv4Address
from typing import Any, Optional
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service
class DataManipulatorService(Service):
"""
Red Agent Data Integration Service.
The Service represents a bot that causes files/folders in the File System to
become corrupted.
"""
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulatorBot"
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def start(self, target_ip_address: IPv4Address, payload: Optional[Any] = "DELETE TABLE users", **kwargs):
"""
Run the DataManipulatorService actions.
:param: target_ip_address: The IP address of the target machine to attack
:param: payload: The payload to send to the target machine
"""
super().start()
self.software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=target_ip_address, dest_port=self.port
)

0
tests/e2e_integration_tests/__init__.py Normal file
View File

25
tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py Normal file
View File

@@ -0,0 +1,25 @@
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.services.database_service import DatabaseService
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
def test_data_manipulation(uc2_network):
client_1: Computer = uc2_network.get_node_by_hostname("client_1")
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
database_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = database_server.software_manager.software["DatabaseService"]
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
# First check that the DB client on the web_server can successfully query the users table on the database
assert db_client.query("SELECT * FROM user;")
# Now we run the DataManipulationBot
db_manipulation_bot.run()
# Now check that the DB client on the web_server cannot query the users table on the database
assert not db_client.query("SELECT * FROM user;")

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

@@ -12,9 +12,6 @@ def test_database_client_server_connection(uc2_network):
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
assert len(db_service.connections) == 0
assert db_client.connect(server_ip_address=IPv4Address("192.168.1.14"))
assert len(db_service.connections) == 1
db_client.disconnect()
@@ -27,11 +24,14 @@ def test_database_client_server_correct_password(uc2_network):
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
db_client.disconnect()
db_client.configure(server_ip_address=IPv4Address("192.168.1.14"), server_password="12345")
db_service.password = "12345"
assert len(db_service.connections) == 0
assert db_client.connect()
assert db_client.connect(server_ip_address=IPv4Address("192.168.1.14"), password="12345")
assert len(db_service.connections) == 1
@@ -41,11 +41,12 @@ def test_database_client_server_incorrect_password(uc2_network):
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
db_client.disconnect()
db_client.configure(server_ip_address=IPv4Address("192.168.1.14"), server_password="54321")
db_service.password = "12345"
assert len(db_service.connections) == 0
assert not db_client.connect(server_ip_address=IPv4Address("192.168.1.14"), password="54321")
assert not db_client.connect()
assert len(db_service.connections) == 0
@@ -53,14 +54,6 @@ 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["DatabaseClient"]
db_client.connect()
db_client.connect(server_ip_address=IPv4Address("192.168.1.14"))
db_client.query("SELECT * FROM user;")
web_server_nic = web_server.ethernet_port[1]
web_server_last_payload = web_server_nic.pcap.read()[-1]["payload"]
assert web_server_last_payload["status_code"] == 200
assert web_server_last_payload["data"]
assert db_client.query("SELECT * FROM user;")

3
tests/unit_tests/_primaite/_simulator/_network/test_container.py
View File

@@ -1,5 +1,7 @@
import json
import pytest
from primaite.simulator.network.container import Network
@@ -10,6 +12,7 @@ def test_creating_container():
assert net.links == {}
@pytest.mark.skip(reason="Skipping until we tackle serialisation")
def test_describe_state():
"""Check that we can describe network state without raising errors, and that the result is JSON serialisable."""
net = Network()

20
tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py Normal file
View File

@@ -0,0 +1,20 @@
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.base import Node
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.services.red_services.data_manipulation_bot import DataManipulationBot
def test_creation():
network = arcd_uc2_network()
client_1: Node = network.get_node_by_hostname("client_1")
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
assert data_manipulation_bot.name == "DataManipulationBot"
assert data_manipulation_bot.port == Port.POSTGRES_SERVER
assert data_manipulation_bot.protocol == IPProtocol.TCP
assert data_manipulation_bot.payload == "DROP TABLE IF EXISTS user;"

32
tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulator_service.py
View File

@@ -1,32 +0,0 @@
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.base import Node
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.services.red_services.data_manipulator_service import DataManipulatorService
def test_creation():
network = arcd_uc2_network()
client_1: Node = network.get_node_by_hostname("client_1")
client_1.software_manager.install(service_class=DataManipulatorService)
data_manipulator_service: DataManipulatorService = client_1.software_manager.software["DataManipulatorBot"]
assert data_manipulator_service.name == "DataManipulatorBot"
assert data_manipulator_service.port == Port.POSTGRES_SERVER
assert data_manipulator_service.protocol == IPProtocol.TCP
# should have no session yet
assert len(client_1.session_manager.sessions_by_uuid) == 0
try:
data_manipulator_service.start(target_ip_address=IPv4Address("192.168.1.14"))
except Exception as e:
assert False, f"Test was not supposed to throw exception: {e}"
# there should be a session after the service is started
assert len(client_1.session_manager.sessions_by_uuid) == 1