From c096d06bcdf49f10769b320840ca1ee94ddda318 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Tue, 26 Sep 2023 15:14:24 +0100 Subject: [PATCH 1/7] #1796: pre installing system software --- .../system/ftp_client_server.rst | 4 ++-- src/primaite/simulator/network/hardware/base.py | 12 ++++++++++++ src/primaite/simulator/network/networks.py | 8 -------- .../_simulator/_system/_services/test_dns.py | 2 -- .../_simulator/_system/_services/test_ftp.py | 2 -- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/source/simulation_components/system/ftp_client_server.rst b/docs/source/simulation_components/system/ftp_client_server.rst index 0e4aeea3..94aef925 100644 --- a/docs/source/simulation_components/system/ftp_client_server.rst +++ b/docs/source/simulation_components/system/ftp_client_server.rst @@ -94,11 +94,11 @@ Example peer to peer network Install the FTP Server ^^^^^^^^^^^^^^^^^^^^^^ +FTP Client should be pre installed on nodes + .. code-block:: python srv.software_manager.install(FTPServer) - pc1.software_manager.install(FTPClient) - client: FTPClient = pc1.software_manager.software['FTPClient'] ftpserv: FTPServer = srv.software_manager.software['FTPServer'] Setting up the FTP Server diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index dd2130d2..00b1f097 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -25,6 +25,8 @@ from primaite.simulator.system.core.session_manager import SessionManager from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.processes.process import Process +from primaite.simulator.system.services.dns.dns_client import DNSClient +from primaite.simulator.system.services.ftp.ftp_client import FTPClient from primaite.simulator.system.services.service import Service _LOGGER = getLogger(__name__) @@ -937,6 +939,16 @@ class Node(SimComponent): self.arp.nics = self.nics self.session_manager.software_manager = self.software_manager + self._install_system_software() + + def _install_system_software(self): + """Install System Software - software that is usually provided with the OS.""" + # DNS Client + self.software_manager.install(DNSClient) + + # FTP + self.software_manager.install(FTPClient) + def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 63cb05e0..6122146b 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -10,9 +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.database_service import DatabaseService -from primaite.simulator.system.services.dns.dns_client import DNSClient from primaite.simulator.system.services.dns.dns_server import DNSServer -from primaite.simulator.system.services.ftp.ftp_client import FTPClient from primaite.simulator.system.services.ftp.ftp_server import FTPServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot @@ -137,14 +135,11 @@ def arcd_uc2_network() -> Network: dns_server=IPv4Address("192.168.1.10"), ) client_1.power_on() - client_1.software_manager.install(DNSClient) 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_1.software_manager.install(FTPClient) - # Client 2 client_2 = Computer( hostname="client_2", @@ -154,11 +149,8 @@ def arcd_uc2_network() -> Network: dns_server=IPv4Address("192.168.1.10"), ) client_2.power_on() - client_2.software_manager.install(DNSClient) network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2]) - client_2.software_manager.install(FTPClient) - # Domain Controller domain_controller = Server( hostname="domain_controller", diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py index f501d14a..d86791cd 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py @@ -21,8 +21,6 @@ def dns_server() -> Node: @pytest.fixture(scope="function") def dns_client() -> Node: node = Node(hostname="dns_client") - node.software_manager.install(software_class=DNSClient) - node.software_manager.software["DNSClient"].start() return node diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py index ea563a88..fce4a487 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py @@ -21,8 +21,6 @@ def ftp_server() -> Node: @pytest.fixture(scope="function") def ftp_client() -> Node: node = Node(hostname="ftp_client") - node.software_manager.install(software_class=FTPClient) - node.software_manager.software["FTPClient"].start() return node From 6202d320a6948b6831bd8801ca79c4b3f3447eaf Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Thu, 28 Sep 2023 12:23:49 +0100 Subject: [PATCH 2/7] #1796: Add ability to create and restore database backups + add more sys log messages + remove the link size checks temporarily --- .../simulator/network/hardware/base.py | 4 +- src/primaite/simulator/network/networks.py | 1 + .../services/database/database_service.py | 101 ++++++++++++++++++ .../system/services/ftp/ftp_client.py | 12 +-- .../system/services/ftp/ftp_server.py | 3 + .../system/services/ftp/ftp_service.py | 9 +- .../system/test_database_on_node.py | 28 +++++ .../system/test_ftp_client_server.py | 4 +- 8 files changed, 146 insertions(+), 16 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 00b1f097..7c08f9fc 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -498,7 +498,9 @@ class Link(SimComponent): def _can_transmit(self, frame: Frame) -> bool: if self.is_up: frame_size_Mbits = frame.size_Mbits # noqa - Leaving it as Mbits as this is how they're expressed - return self.current_load + frame_size_Mbits <= self.bandwidth + # return self.current_load + frame_size_Mbits <= self.bandwidth + # TODO: re add this check once packet size limiting and MTU checks are implemented + return True return False def transmit_frame(self, sender_nic: Union[NIC, SwitchPort], frame: Frame) -> bool: diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 6122146b..f54e1172 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -231,6 +231,7 @@ 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.configure_backup(backup_server=IPv4Address("192.168.1.16")) database_service._process_sql(ddl, None) # noqa for insert_statement in user_insert_statements: database_service._process_sql(insert_statement, None) # noqa diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 62120fc7..268bd54f 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -1,5 +1,6 @@ import sqlite3 from datetime import datetime +from ipaddress import IPv4Address from sqlite3 import OperationalError from typing import Any, Dict, List, Optional, Union @@ -9,6 +10,7 @@ from primaite.simulator.file_system.file_system import File from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.core.software_manager import SoftwareManager +from primaite.simulator.system.services.ftp.ftp_client import FTPClient from primaite.simulator.system.services.service import Service, ServiceOperatingState from primaite.simulator.system.software import SoftwareHealthState @@ -23,6 +25,15 @@ class DatabaseService(Service): password: Optional[str] = None connections: Dict[str, datetime] = {} + backup_server: IPv4Address = None + """IP address of the backup server.""" + + latest_backup_directory: str = None + """Directory of latest backup.""" + + latest_backup_file_name: str = None + """File name of latest backup.""" + def __init__(self, **kwargs): kwargs["name"] = "DatabaseService" kwargs["port"] = Port.POSTGRES_SERVER @@ -58,6 +69,96 @@ class DatabaseService(Service): table.add_row([row]) print(table) + def configure_backup(self, backup_server: IPv4Address): + """ + Set up the database backup. + + :param: backup_server_ip: The IP address of the backup server + """ + self.backup_server = backup_server + + def backup_database( + self, backup_directory: Optional[str] = "db_backup", backup_file_name: Optional[str] = None + ) -> bool: + """ + Create a backup of the database to the configured backup server. + + :param: backup_directory: Name of directory where backup will be stored. Optional. + :type: backup_directory: Optional[str] + + :param: backup_file_name: Name of file where backup will be stored. Optional. + :type: backup_file_name: Optional[str] + """ + # check if the backup server was configured + if self.backup_server is None: + self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.") + return False + + if backup_file_name is None: + backup_file_name = f"{datetime.now().strftime('%d-%m-%Y_%H-%M')}.db" + + software_manager: SoftwareManager = self.software_manager + ftp_client_service: FTPClient = software_manager.software["FTPClient"] + + # send backup copy of database file to FTP server + response = ftp_client_service.send_file( + dest_ip_address=self.backup_server, + src_file_name=self._db_file.name, + src_folder_name=self._db_file.folder.name, + dest_folder_name=backup_directory, + dest_file_name=backup_file_name, + ) + + if response: + self.latest_backup_directory = backup_directory + self.latest_backup_file_name = backup_file_name + return True + + self.sys_log.error("Unable to create database backup.") + return False + + def restore_backup(self, backup_directory: Optional[str] = None, backup_file_name: Optional[str] = None) -> bool: + """ + Restore a backup from backup server. + + :param: backup_directory: Name of directory where backup will be stored. Optional. + :type: backup_directory: Optional[str] + + :param: backup_file_name: Name of file where backup will be stored. Optional. + :type: backup_file_name: Optional[str] + """ + if backup_directory is None: + backup_directory = self.latest_backup_directory + + if backup_file_name is None: + backup_file_name = self.latest_backup_file_name + + software_manager: SoftwareManager = self.software_manager + ftp_client_service: FTPClient = software_manager.software["FTPClient"] + + # retrieve backup file from backup server + response = ftp_client_service.request_file( + src_folder_name=backup_directory, + src_file_name=backup_file_name, + dest_folder_name="downloads", + dest_file_name="database.db", + dest_ip_address=self.backup_server, + ) + + if response: + # replace db file + self.file_system.delete_file(folder_name=self.folder.name, file_name="downloads.db") + + self.file_system.move_file( + src_folder_name="downloads", src_file_name="database.db", dst_folder_name=self.folder.name + ) + self._db_file = self.file_system.get_file(folder_name=self.folder.name, file_name="database.db") + + return self._db_file is not None + + self.sys_log.error("Unable to restore database backup.") + return False + def _create_db_file(self): """Creates the Simulation File and sqlite file in the file system.""" self._db_file: File = self.file_system.create_file(folder_name="database", file_name="database.db", real=True) diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 33fe32be..8359e8a0 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -107,7 +107,7 @@ class FTPClient(FTPServiceABC): payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port ) if payload.status_code == FTPStatusCode.OK: - self.connected = False + self.connected = None return True return False @@ -159,8 +159,9 @@ class FTPClient(FTPServiceABC): if not self.connected: return False else: + self.sys_log.info(f"Sending file {src_folder_name}/{src_file_name} to {str(dest_ip_address)}") # send STOR request - self._send_data( + return self._send_data( file=file_to_transfer, dest_folder_name=dest_folder_name, dest_file_name=dest_file_name, @@ -168,9 +169,6 @@ class FTPClient(FTPServiceABC): dest_port=dest_port, ) - # send disconnect - return self._disconnect_from_server(dest_ip_address=dest_ip_address, dest_port=dest_port) - def request_file( self, dest_ip_address: IPv4Address, @@ -222,14 +220,12 @@ class FTPClient(FTPServiceABC): "dest_folder_name": dest_folder_name, }, ) + self.sys_log.info(f"Requesting file {src_folder_name}/{src_file_name} from {str(dest_ip_address)}") software_manager: SoftwareManager = self.software_manager software_manager.send_payload_to_session_manager( payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port ) - # send disconnect - self._disconnect_from_server(dest_ip_address=dest_ip_address, dest_port=dest_port) - # the payload should have ok status code if payload.status_code == FTPStatusCode.OK: self.sys_log.info(f"File {src_folder_name}/{src_file_name} found in FTP server.") diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index 83c883f1..93f8b45b 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/src/primaite/simulator/system/services/ftp/ftp_server.py @@ -55,6 +55,9 @@ class FTPServer(FTPServiceABC): if session_id: session_details = self._get_session_details(session_id) + if payload.ftp_command is not None: + self.sys_log.info(f"Received FTP {payload.ftp_command.name} command.") + # process server specific commands, otherwise call super if payload.ftp_command == FTPCommand.PORT: # check that the port is valid diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index f47b8f64..c6d63751 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -25,6 +25,9 @@ class FTPServiceABC(Service, ABC): :param: session_id: session ID linked to the FTP Packet. Optional. :type: session_id: Optional[str] """ + if payload.ftp_command is not None: + self.sys_log.info(f"Received FTP {payload.ftp_command.name} command.") + # handle STOR request if payload.ftp_command == FTPCommand.STOR: # check that the file is created in the computed hosting the FTP server @@ -48,11 +51,7 @@ class FTPServiceABC(Service, ABC): file_name = payload.ftp_command_args["dest_file_name"] folder_name = payload.ftp_command_args["dest_folder_name"] file_size = payload.ftp_command_args["file_size"] - self.file_system.create_file( - file_name=file_name, - folder_name=folder_name, - size=file_size, - ) + self.file_system.create_file(file_name=file_name, folder_name=folder_name, size=file_size, real=True) self.sys_log.info( f"Created item in {self.sys_log.hostname}: {payload.ftp_command_args['dest_folder_name']}/" f"{payload.ftp_command_args['dest_file_name']}" diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index c69f131c..38196041 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/tests/integration_tests/system/test_database_on_node.py @@ -3,6 +3,7 @@ from ipaddress import IPv4Address 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 +from primaite.simulator.system.services.ftp.ftp_server import FTPServer def test_database_client_server_connection(uc2_network): @@ -57,3 +58,30 @@ def test_database_client_query(uc2_network): db_client.connect() assert db_client.query("SELECT * FROM user;") + + +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["DatabaseService"] + + # back up should be created + assert db_service.backup_database(backup_file_name="test_file.db") is True + + backup_server: Server = uc2_network.get_node_by_hostname("backup_server") + 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_backup", file_name="test_file.db") is not None + + +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["DatabaseService"] + + # create a back up + assert db_service.backup_database(backup_file_name="test_file.db") is True + + # back up should be restored + assert db_service.restore_backup(backup_file_name="test_file.db") is True diff --git a/tests/integration_tests/system/test_ftp_client_server.py b/tests/integration_tests/system/test_ftp_client_server.py index fbbe6011..48dc2960 100644 --- a/tests/integration_tests/system/test_ftp_client_server.py +++ b/tests/integration_tests/system/test_ftp_client_server.py @@ -23,7 +23,7 @@ def test_ftp_client_store_file_in_server(uc2_network): # create file on ftp client ftp_client.file_system.create_file(file_name="test_file.txt") - ftp_client.send_file( + assert ftp_client.send_file( src_folder_name="root", src_file_name="test_file.txt", dest_folder_name="client_1_backup", @@ -50,7 +50,7 @@ def test_ftp_client_retrieve_file_from_server(uc2_network): # create file on ftp server ftp_server.file_system.create_file(file_name="test_file.txt", folder_name="file_share") - ftp_client.request_file( + assert ftp_client.request_file( src_folder_name="file_share", src_file_name="test_file.txt", dest_folder_name="downloads", From bca3e6344eab3d7fe310fcb316420afd47e981a7 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Thu, 28 Sep 2023 14:09:32 +0100 Subject: [PATCH 3/7] #1796: documentation --- CHANGELOG.md | 3 +++ .../simulation_components/system/database_client_server.rst | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7147f82b..cc9c26d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,9 @@ SessionManager. - File System - ability to emulate a node's file system during a simulation - Example notebooks - There is currently 1 jupyter notebook which walks through using PrimAITE 1. Creating a simulation - this notebook explains how to build up a simulation using the Python package. (WIP) +- Database: + - `DatabaseClient` and `DatabaseService` created to allow emulation of database actions + - Ability to `backup_database` and `restore_backup` for a `DatabaseService` - Red Agent Services: - Data Manipulator Bot - A red agent service which sends a payload to a target machine. (By default this payload is a SQL query that breaks a database) - DNS Services: `DNSClient` and `DNSServer` diff --git a/docs/source/simulation_components/system/database_client_server.rst b/docs/source/simulation_components/system/database_client_server.rst index 99bbe25e..ef911e0e 100644 --- a/docs/source/simulation_components/system/database_client_server.rst +++ b/docs/source/simulation_components/system/database_client_server.rst @@ -60,6 +60,12 @@ Usage - Retrieve results in a dictionary. - Disconnect when finished. +To create database backups: + +- Configure the backup server the ``DatabaseService`` by providing the Backup server ``IPv4Address`` with ``configure_backup`` +- Create a backup using ``backup_database``. This fails if the backup server is not configured. +- Restore a backup using ``restore_backup``. By default, this uses the database created via ``backup_database``. + Implementation ^^^^^^^^^^^^^^ From 3dc8a0f222636b8f0ee6c9e2703e7077a9765643 Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 29 Sep 2023 20:14:42 +0100 Subject: [PATCH 4/7] #1796 - Made the FTP copy real files. Hardcoded the DatabaseService backup folder and filename. Added db restore and final query check to the data manipulation e2e test. --- .../services/database/database_service.py | 39 +++++++++---------- .../system/services/ftp/ftp_client.py | 1 + .../system/services/ftp/ftp_service.py | 10 ++++- .../test_uc2_data_manipulation_scenario.py | 8 ++++ 4 files changed, 36 insertions(+), 22 deletions(-) diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 268bd54f..f874b89b 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -41,6 +41,9 @@ class DatabaseService(Service): super().__init__(**kwargs) self._db_file: File self._create_db_file() + self._connect() + + def _connect(self): self._conn = sqlite3.connect(self._db_file.sim_path) self._cursor = self._conn.cursor() @@ -51,8 +54,10 @@ class DatabaseService(Service): :return: List of table names. """ sql = "SELECT name FROM sqlite_master WHERE type='table' AND name != 'sqlite_sequence';" - results = self._process_sql(sql) - return [row[0] for row in results["data"]] + results = self._process_sql(sql, None) + if isinstance(results["data"], dict): + return list(results["data"].keys()) + return [] def show(self, markdown: bool = False): """ @@ -77,9 +82,7 @@ class DatabaseService(Service): """ self.backup_server = backup_server - def backup_database( - self, backup_directory: Optional[str] = "db_backup", backup_file_name: Optional[str] = None - ) -> bool: + def backup_database(self) -> bool: """ Create a backup of the database to the configured backup server. @@ -94,8 +97,7 @@ class DatabaseService(Service): self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.") return False - if backup_file_name is None: - backup_file_name = f"{datetime.now().strftime('%d-%m-%Y_%H-%M')}.db" + self._conn.close() software_manager: SoftwareManager = self.software_manager ftp_client_service: FTPClient = software_manager.software["FTPClient"] @@ -105,19 +107,19 @@ class DatabaseService(Service): dest_ip_address=self.backup_server, src_file_name=self._db_file.name, src_folder_name=self._db_file.folder.name, - dest_folder_name=backup_directory, - dest_file_name=backup_file_name, + dest_folder_name=str(self.uuid), + dest_file_name="database.db", + real_file_path=self._db_file.sim_path, ) + self._connect() if response: - self.latest_backup_directory = backup_directory - self.latest_backup_file_name = backup_file_name return True self.sys_log.error("Unable to create database backup.") return False - def restore_backup(self, backup_directory: Optional[str] = None, backup_file_name: Optional[str] = None) -> bool: + def restore_backup(self) -> bool: """ Restore a backup from backup server. @@ -127,32 +129,27 @@ class DatabaseService(Service): :param: backup_file_name: Name of file where backup will be stored. Optional. :type: backup_file_name: Optional[str] """ - if backup_directory is None: - backup_directory = self.latest_backup_directory - - if backup_file_name is None: - backup_file_name = self.latest_backup_file_name - software_manager: SoftwareManager = self.software_manager ftp_client_service: FTPClient = software_manager.software["FTPClient"] # retrieve backup file from backup server response = ftp_client_service.request_file( - src_folder_name=backup_directory, - src_file_name=backup_file_name, + src_folder_name=str(self.uuid), + src_file_name="database.db", dest_folder_name="downloads", dest_file_name="database.db", dest_ip_address=self.backup_server, ) if response: + self._conn.close() # replace db file self.file_system.delete_file(folder_name=self.folder.name, file_name="downloads.db") - self.file_system.move_file( src_folder_name="downloads", src_file_name="database.db", dst_folder_name=self.folder.name ) self._db_file = self.file_system.get_file(folder_name=self.folder.name, file_name="database.db") + self._connect() return self._db_file is not None diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 8359e8a0..c22f704b 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -119,6 +119,7 @@ class FTPClient(FTPServiceABC): dest_folder_name: str, dest_file_name: str, dest_port: Optional[Port] = Port.FTP, + real_file_path: Optional[str] = None, ) -> bool: """ Send a file to a target IP address. diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index c6d63751..5314b6a3 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -1,3 +1,4 @@ +import shutil from abc import ABC from ipaddress import IPv4Address from typing import Optional @@ -51,11 +52,17 @@ class FTPServiceABC(Service, ABC): file_name = payload.ftp_command_args["dest_file_name"] folder_name = payload.ftp_command_args["dest_folder_name"] file_size = payload.ftp_command_args["file_size"] - self.file_system.create_file(file_name=file_name, folder_name=folder_name, size=file_size, real=True) + real_file_path = payload.ftp_command_args.get("real_file_path") + is_real = real_file_path is not None + file = self.file_system.create_file( + file_name=file_name, folder_name=folder_name, size=file_size, real=is_real + ) self.sys_log.info( f"Created item in {self.sys_log.hostname}: {payload.ftp_command_args['dest_folder_name']}/" f"{payload.ftp_command_args['dest_file_name']}" ) + if is_real: + shutil.copy(real_file_path, file.sim_path) # file should exist return self.file_system.get_file(file_name=file_name, folder_name=folder_name) is not None except Exception as e: @@ -99,6 +106,7 @@ class FTPServiceABC(Service, ABC): "dest_folder_name": dest_folder_name, "dest_file_name": dest_file_name, "file_size": file.sim_size, + "real_file_path": file.sim_path if file.real else None, }, packet_payload_size=file.sim_size, ) diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index 50998f09..955fa20e 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -15,6 +15,8 @@ def test_data_manipulation(uc2_network): web_server: Server = uc2_network.get_node_by_hostname("web_server") db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] + db_service.backup_database() + # 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;") @@ -23,3 +25,9 @@ def test_data_manipulation(uc2_network): # 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;") + + # Now restore the database + db_service.restore_backup() + + # Now 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;") From 84405d7ed3578f4f99e924149bac52ebc441a55b Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 29 Sep 2023 20:19:26 +0100 Subject: [PATCH 5/7] #1796 - Added docstring to the test_uc2_data_manipulation_scenario.py --- .../e2e_integration_tests/test_uc2_data_manipulation_scenario.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index 955fa20e..13f4d1f3 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -6,6 +6,7 @@ from primaite.simulator.system.services.red_services.data_manipulation_bot impor def test_data_manipulation(uc2_network): + """Tests the UC2 data manipulation scenario end-to-end. Is a work in progress.""" client_1: Computer = uc2_network.get_node_by_hostname("client_1") db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"] From fdebfce406b8e28d1e3efb841d8e2d7ec6124187 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Mon, 2 Oct 2023 12:14:59 +0100 Subject: [PATCH 6/7] #1796: Fix test + making the restore test better --- .../system/test_database_on_node.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index 38196041..92056981 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/tests/integration_tests/system/test_database_on_node.py @@ -66,13 +66,13 @@ def test_create_database_backup(uc2_network): db_service: DatabaseService = db_server.software_manager.software["DatabaseService"] # back up should be created - assert db_service.backup_database(backup_file_name="test_file.db") is True + 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["FTPServer"] # backup file should exist in the backup server - assert ftp_server.file_system.get_file(folder_name="db_backup", file_name="test_file.db") is not None + assert ftp_server.file_system.get_file(folder_name=db_service.uuid, file_name="database.db") is not None def test_restore_backup(uc2_network): @@ -81,7 +81,14 @@ def test_restore_backup(uc2_network): db_service: DatabaseService = db_server.software_manager.software["DatabaseService"] # create a back up - assert db_service.backup_database(backup_file_name="test_file.db") is True + assert db_service.backup_database() is True + + # delete database locally + db_service.file_system.delete_file(folder_name="database", file_name="database.db") + + assert db_service.file_system.get_file(folder_name="database", file_name="database.db") is None # back up should be restored - assert db_service.restore_backup(backup_file_name="test_file.db") is True + assert db_service.restore_backup() is True + + assert db_service.file_system.get_file(folder_name="database", file_name="database.db") is not None From 97f0267539c67e847f416d5de46629edd23cac21 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Wed, 4 Oct 2023 11:33:18 +0100 Subject: [PATCH 7/7] #1796: apply PR suggestions + fix tests --- CHANGELOG.md | 2 +- .../system/database_client_server.rst | 2 +- .../simulator/network/hardware/base.py | 8 +------- .../network/hardware/nodes/computer.py | 13 ++++++++++++ .../services/database/database_service.py | 20 ++----------------- .../system/services/ftp/ftp_client.py | 9 ++++++--- .../_simulator/_system/_services/test_dns.py | 10 ++++++++-- .../_simulator/_system/_services/test_ftp.py | 10 ++++++++-- 8 files changed, 40 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cc9c26d1..a5bc08f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,7 +29,7 @@ SessionManager. 1. Creating a simulation - this notebook explains how to build up a simulation using the Python package. (WIP) - Database: - `DatabaseClient` and `DatabaseService` created to allow emulation of database actions - - Ability to `backup_database` and `restore_backup` for a `DatabaseService` + - Ability for `DatabaseService` to backup its data to another server via FTP and restore data from backup - Red Agent Services: - Data Manipulator Bot - A red agent service which sends a payload to a target machine. (By default this payload is a SQL query that breaks a database) - DNS Services: `DNSClient` and `DNSServer` diff --git a/docs/source/simulation_components/system/database_client_server.rst b/docs/source/simulation_components/system/database_client_server.rst index ef911e0e..32568477 100644 --- a/docs/source/simulation_components/system/database_client_server.rst +++ b/docs/source/simulation_components/system/database_client_server.rst @@ -62,7 +62,7 @@ Usage To create database backups: -- Configure the backup server the ``DatabaseService`` by providing the Backup server ``IPv4Address`` with ``configure_backup`` +- Configure the backup server on the ``DatabaseService`` by providing the Backup server ``IPv4Address`` with ``configure_backup`` - Create a backup using ``backup_database``. This fails if the backup server is not configured. - Restore a backup using ``restore_backup``. By default, this uses the database created via ``backup_database``. diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 7c08f9fc..2725ab1a 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -25,8 +25,6 @@ from primaite.simulator.system.core.session_manager import SessionManager from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.processes.process import Process -from primaite.simulator.system.services.dns.dns_client import DNSClient -from primaite.simulator.system.services.ftp.ftp_client import FTPClient from primaite.simulator.system.services.service import Service _LOGGER = getLogger(__name__) @@ -945,11 +943,7 @@ class Node(SimComponent): def _install_system_software(self): """Install System Software - software that is usually provided with the OS.""" - # DNS Client - self.software_manager.install(DNSClient) - - # FTP - self.software_manager.install(FTPClient) + pass def describe_state(self) -> Dict: """ diff --git a/src/primaite/simulator/network/hardware/nodes/computer.py b/src/primaite/simulator/network/hardware/nodes/computer.py index 5452666b..3ceb5291 100644 --- a/src/primaite/simulator/network/hardware/nodes/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/computer.py @@ -1,4 +1,6 @@ from primaite.simulator.network.hardware.base import NIC, Node +from primaite.simulator.system.services.dns.dns_client import DNSClient +from primaite.simulator.system.services.ftp.ftp_client import FTPClient class Computer(Node): @@ -36,3 +38,14 @@ class Computer(Node): def __init__(self, **kwargs): super().__init__(**kwargs) self.connect_nic(NIC(ip_address=kwargs["ip_address"], subnet_mask=kwargs["subnet_mask"])) + self._install_system_software() + + def _install_system_software(self): + """Install System Software - software that is usually provided with the OS.""" + # DNS Client + self.software_manager.install(DNSClient) + + # FTP + self.software_manager.install(FTPClient) + + super()._install_system_software() diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index f874b89b..0a6de8c3 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -83,15 +83,7 @@ class DatabaseService(Service): self.backup_server = backup_server def backup_database(self) -> bool: - """ - Create a backup of the database to the configured backup server. - - :param: backup_directory: Name of directory where backup will be stored. Optional. - :type: backup_directory: Optional[str] - - :param: backup_file_name: Name of file where backup will be stored. Optional. - :type: backup_file_name: Optional[str] - """ + """Create a backup of the database to the configured backup server.""" # check if the backup server was configured if self.backup_server is None: self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.") @@ -120,15 +112,7 @@ class DatabaseService(Service): return False def restore_backup(self) -> bool: - """ - Restore a backup from backup server. - - :param: backup_directory: Name of directory where backup will be stored. Optional. - :type: backup_directory: Optional[str] - - :param: backup_file_name: Name of file where backup will be stored. Optional. - :type: backup_file_name: Optional[str] - """ + """Restore a backup from backup server.""" software_manager: SoftwareManager = self.software_manager ftp_client_service: FTPClient = software_manager.software["FTPClient"] diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index c22f704b..648b2494 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -107,7 +107,7 @@ class FTPClient(FTPServiceABC): payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port ) if payload.status_code == FTPStatusCode.OK: - self.connected = None + self.connected = False return True return False @@ -162,13 +162,16 @@ class FTPClient(FTPServiceABC): else: self.sys_log.info(f"Sending file {src_folder_name}/{src_file_name} to {str(dest_ip_address)}") # send STOR request - return self._send_data( + if self._send_data( file=file_to_transfer, dest_folder_name=dest_folder_name, dest_file_name=dest_file_name, dest_ip_address=dest_ip_address, dest_port=dest_port, - ) + ): + return self._disconnect_from_server(dest_ip_address=dest_ip_address, dest_port=dest_port) + + return False def request_file( self, diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py index d86791cd..31718387 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py @@ -3,6 +3,8 @@ from ipaddress import IPv4Address import pytest from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.server import Server from primaite.simulator.network.protocols.dns import DNSPacket, DNSReply, DNSRequest from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port @@ -12,7 +14,9 @@ from primaite.simulator.system.services.dns.dns_server import DNSServer @pytest.fixture(scope="function") def dns_server() -> Node: - node = Node(hostname="dns_server") + node = Server( + hostname="dns_server", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" + ) node.software_manager.install(software_class=DNSServer) node.software_manager.software["DNSServer"].start() return node @@ -20,7 +24,9 @@ def dns_server() -> Node: @pytest.fixture(scope="function") def dns_client() -> Node: - node = Node(hostname="dns_client") + node = Computer( + hostname="dns_client", ip_address="192.168.1.11", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" + ) return node diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py index fce4a487..3ccb0c99 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py @@ -3,6 +3,8 @@ from ipaddress import IPv4Address import pytest from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.server import Server from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port @@ -12,7 +14,9 @@ from primaite.simulator.system.services.ftp.ftp_server import FTPServer @pytest.fixture(scope="function") def ftp_server() -> Node: - node = Node(hostname="ftp_server") + node = Server( + hostname="ftp_server", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" + ) node.software_manager.install(software_class=FTPServer) node.software_manager.software["FTPServer"].start() return node @@ -20,7 +24,9 @@ def ftp_server() -> Node: @pytest.fixture(scope="function") def ftp_client() -> Node: - node = Node(hostname="ftp_client") + node = Computer( + hostname="ftp_client", ip_address="192.168.1.11", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" + ) return node