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

#1796: apply PR suggestions + fix tests

Browse Source
This commit is contained in:
Czar.Echavez
2023-10-04 11:33:18 +01:00
parent fdebfce406
commit 97f0267539
8 changed files with 40 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CHANGELOG.md
View File

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

2
docs/source/simulation_components/system/database_client_server.rst
View File

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

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

@@ -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:
"""

13
src/primaite/simulator/network/hardware/nodes/computer.py
View File

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

20
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -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"]

9
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -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,

10
tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py
View File

@@ -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

10
tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py
View File

@@ -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