Remove real database in favour of simulated

This commit is contained in:
Marek Wolan
2023-11-18 03:40:08 +00:00
parent 815bdfe603
commit 77f3806ba7
9 changed files with 53 additions and 102 deletions

View File

@@ -41,7 +41,7 @@ Example
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
client_1.software_manager.install(DataManipulationBot) client_1.software_manager.install(DataManipulationBot)
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software["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.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DELETE")
data_manipulation_bot.run() 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. This would connect to the database service at 192.168.1.14, authenticate, and execute the SQL statement to drop the 'users' table.

View File

@@ -81,7 +81,7 @@ game_config:
# options: # options:
# execution_definition: # execution_definition:
# server_ip: 192.168.1.14 # server_ip: 192.168.1.14
# payload: "DROP TABLE IF EXISTS user;" # payload: "DELETE"
# success_rate: 80% # success_rate: 80%
- type: NODE_FILE_DELETE - type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT - type: NODE_FILE_CORRUPT

View File

@@ -140,7 +140,7 @@ def arcd_uc2_network() -> Network:
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
client_1.software_manager.install(DataManipulationBot) client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["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;") db_manipulation_bot.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DELETE")
# Client 2 # Client 2
client_2 = Computer( client_2 = Computer(

View File

@@ -2,13 +2,14 @@ from ipaddress import IPv4Address
from typing import Any, Dict, Optional from typing import Any, Dict, Optional
from uuid import uuid4 from uuid import uuid4
from prettytable import PrettyTable from primaite import getLogger
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, ApplicationOperatingState from primaite.simulator.system.applications.application import Application, ApplicationOperatingState
from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.core.software_manager import SoftwareManager
_LOGGER = getLogger(__name__)
class DatabaseClient(Application): class DatabaseClient(Application):
""" """
@@ -148,21 +149,6 @@ class DatabaseClient(Application):
self._query_success_tracker[query_id] = False self._query_success_tracker[query_id] = False
return self._query(sql=sql, query_id=query_id) return self._query(sql=sql, query_id=query_id)
def _print_data(self, data: Dict):
"""
Display the contents of the Folder in tabular format.
:param markdown: Whether to display the table in Markdown format or not. Default is `False`.
"""
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)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool: def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
""" """
Receive a payload from the Software Manager. Receive a payload from the Software Manager.
@@ -179,5 +165,5 @@ class DatabaseClient(Application):
status_code = payload.get("status_code") status_code = payload.get("status_code")
self._query_success_tracker[query_id] = status_code == 200 self._query_success_tracker[query_id] = status_code == 200
if self._query_success_tracker[query_id]: if self._query_success_tracker[query_id]:
self._print_data(payload["data"]) _LOGGER.debug(f"Received payload {payload}")
return True return True

View File

@@ -1,10 +1,6 @@
import sqlite3
from datetime import datetime from datetime import datetime
from ipaddress import IPv4Address from ipaddress import IPv4Address
from sqlite3 import OperationalError from typing import Any, Dict, List, Literal, Optional, Union
from typing import Any, Dict, List, Optional, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.file_system.file_system import File from primaite.simulator.file_system.file_system import File
from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.network_layer import IPProtocol
@@ -19,7 +15,7 @@ class DatabaseService(Service):
""" """
A class for simulating a generic SQL Server service. A class for simulating a generic SQL Server service.
This class inherits from the `Service` class and provides methods to manage and query a SQLite database. This class inherits from the `Service` class and provides methods to simulate a SQL database.
""" """
password: Optional[str] = None password: Optional[str] = None
@@ -41,38 +37,6 @@ class DatabaseService(Service):
super().__init__(**kwargs) super().__init__(**kwargs)
self._db_file: File self._db_file: File
self._create_db_file() self._create_db_file()
self._connect()
def _connect(self):
self._conn = sqlite3.connect(self._db_file.sim_path)
self._cursor = self._conn.cursor()
def tables(self) -> List[str]:
"""
Get a list of table names present in the database.
:return: List of table names.
"""
sql = "SELECT name FROM sqlite_master WHERE type='table' AND name != 'sqlite_sequence';"
results = self._process_sql(sql, None)
if isinstance(results["data"], dict):
return list(results["data"].keys())
return []
def show(self, markdown: bool = False):
"""
Prints a list of table names in the database using PrettyTable.
:param markdown: Whether to output the table in Markdown format.
"""
table = PrettyTable(["Table"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.file_system.sys_log.hostname} Database"
for row in self.tables():
table.add_row([row])
print(table)
def configure_backup(self, backup_server: IPv4Address): def configure_backup(self, backup_server: IPv4Address):
""" """
@@ -89,8 +53,6 @@ class DatabaseService(Service):
self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.") self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.")
return False return False
self._conn.close()
software_manager: SoftwareManager = self.software_manager software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software["FTPClient"] ftp_client_service: FTPClient = software_manager.software["FTPClient"]
@@ -98,12 +60,10 @@ class DatabaseService(Service):
response = ftp_client_service.send_file( response = ftp_client_service.send_file(
dest_ip_address=self.backup_server, dest_ip_address=self.backup_server,
src_file_name=self._db_file.name, src_file_name=self._db_file.name,
src_folder_name=self._db_file.folder.name, src_folder_name=self.folder.name,
dest_folder_name=str(self.uuid), dest_folder_name=str(self.uuid),
dest_file_name="database.db", dest_file_name="database.db",
real_file_path=self._db_file.sim_path,
) )
self._connect()
if response: if response:
return True return True
@@ -125,25 +85,29 @@ class DatabaseService(Service):
dest_ip_address=self.backup_server, dest_ip_address=self.backup_server,
) )
if response: if not response:
self._conn.close() self.sys_log.error("Unable to restore database backup.")
# replace db file return False
self.file_system.delete_file(folder_name=self.folder.name, file_name="downloads.db")
self.file_system.copy_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 # replace db file
self.file_system.delete_file(folder_name=self.folder.name, file_name="downloads.db")
self.file_system.copy_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.sys_log.error("Unable to restore database backup.") if self._db_file is None:
return False self.sys_log.error("Copying database backup failed.")
return False
self.set_health_state(SoftwareHealthState.GOOD)
return True
def _create_db_file(self): def _create_db_file(self):
"""Creates the Simulation File and sqlite file in the file system.""" """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) self._db_file: File = self.file_system.create_file(folder_name="database", file_name="database.db")
self.folder = self._db_file.folder self.folder = self.file_system.get_folder_by_id(self._db_file.folder_id)
def _process_connect( def _process_connect(
self, session_id: str, password: Optional[str] = None self, session_id: str, password: Optional[str] = None
@@ -163,31 +127,32 @@ class DatabaseService(Service):
status_code = 404 # service not found status_code = 404 # service not found
return {"status_code": status_code, "type": "connect_response", "response": status_code == 200} return {"status_code": status_code, "type": "connect_response", "response": status_code == 200}
def _process_sql(self, query: str, query_id: str) -> Dict[str, Union[int, List[Any]]]: def _process_sql(self, query: Literal["SELECT", "DELETE"], query_id: str) -> Dict[str, Union[int, List[Any]]]:
""" """
Executes the given SQL query and returns the result. Executes the given SQL query and returns the result.
Possible queries:
- SELECT : returns the data
- DELETE : deletes the data
:param query: The SQL query to be executed. :param query: The SQL query to be executed.
:return: Dictionary containing status code and data fetched. :return: Dictionary containing status code and data fetched.
""" """
self.sys_log.info(f"{self.name}: Running {query}") self.sys_log.info(f"{self.name}: Running {query}")
try: if query == "SELECT":
self._cursor.execute(query) if self.health_state_actual == SoftwareHealthState.GOOD:
self._conn.commit() return {"status_code": 200, "type": "sql", "data": True, "uuid": query_id}
except OperationalError: else:
# Handle the case where the table does not exist. return {"status_code": 404, "data": False}
self.sys_log.error(f"{self.name}: Error, query failed") elif query == "DELETE":
return {"status_code": 404, "data": {}} if self.health_state_actual == SoftwareHealthState.GOOD:
data = [] self.health_state_actual = SoftwareHealthState.COMPROMISED
description = self._cursor.description return {"status_code": 200, "type": "sql", "data": False, "uuid": query_id}
if description: else:
headers = [] return {"status_code": 404, "data": False}
for header in description: else:
headers.append(header[0]) # Invalid query
data = self._cursor.fetchall() return {"status_code": 500, "data": False}
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, "uuid": query_id}
def describe_state(self) -> Dict: def describe_state(self) -> Dict:
""" """

View File

@@ -106,7 +106,7 @@ class WebServer(Service):
# get data from DatabaseServer # get data from DatabaseServer
db_client: DatabaseClient = self.software_manager.software["DatabaseClient"] db_client: DatabaseClient = self.software_manager.software["DatabaseClient"]
# get all users # get all users
if db_client.query("SELECT * FROM user;"): if db_client.query("SELECT"):
# query succeeded # query succeeded
response.status_code = HttpStatusCode.OK response.status_code = HttpStatusCode.OK

View File

@@ -19,16 +19,16 @@ def test_data_manipulation(uc2_network):
db_service.backup_database() db_service.backup_database()
# First check that the DB client on the web_server can successfully query the users table on the 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;") assert db_client.query("SELECT")
# Now we run the DataManipulationBot # Now we run the DataManipulationBot
db_manipulation_bot.run() db_manipulation_bot.run()
# Now check that the DB client on the web_server cannot query the users table on the database # 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;") assert not db_client.query("SELECT")
# Now restore the database # Now restore the database
db_service.restore_backup() db_service.restore_backup()
# Now check that the DB client on the web_server can successfully query the users table on the database # 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;") assert db_client.query("SELECT")

View File

@@ -57,7 +57,7 @@ def test_database_client_query(uc2_network):
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_client.connect() db_client.connect()
assert db_client.query("SELECT * FROM user;") assert db_client.query("SELECT")
def test_create_database_backup(uc2_network): def test_create_database_backup(uc2_network):

View File

@@ -17,4 +17,4 @@ def test_creation():
assert data_manipulation_bot.name == "DataManipulationBot" assert data_manipulation_bot.name == "DataManipulationBot"
assert data_manipulation_bot.port == Port.POSTGRES_SERVER assert data_manipulation_bot.port == Port.POSTGRES_SERVER
assert data_manipulation_bot.protocol == IPProtocol.TCP assert data_manipulation_bot.protocol == IPProtocol.TCP
assert data_manipulation_bot.payload == "DROP TABLE IF EXISTS user;" assert data_manipulation_bot.payload == "DELETE"