#1796: apply PR suggestions + fix tests

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`

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``.
 

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
+        pass
-        self.software_manager.install(DNSClient)
-
-        # FTP
-        self.software_manager.install(FTPClient)
 
     def describe_state(self) -> Dict:
         """

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()

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."""
-        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.")
@@ -120,15 +112,7 @@ class DatabaseService(Service):
         return False
 
     def restore_backup(self) -> bool:
-        """
+        """Restore a backup from backup server."""
-        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]
-        """
         software_manager: SoftwareManager = self.software_manager
         ftp_client_service: FTPClient = software_manager.software["FTPClient"]
 

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,

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
 
 

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