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

Merge branch 'feature/1816_Database-Service-(Network-and-User-Interaction)' into feature/1752-dns-server-and-client

Browse Source
This commit is contained in:
Czar Echavez
2023-09-11 16:22:18 +01:00
parent 4c67365d4f 695b3ceab4
commit ee730d4ab0
22 changed files with 268 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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
)