Merged PR 348: #2462 - Refactor of DatabaseClient and DatabaseServer

## Summary
Refactor of `DatabaseClient` and `DatabaseService` to update how connection IDs are generated. These are now provided by DatabaseService when establishing a connection.
Creation of `DatabaseClientConnection` class. This is used by `DatabaseClient` to hold a dictionary of active db connections.

## Test process
Tests have been updated to reflect the changes and all pass

## Checklist
- [X] PR is linked to a **work item**
- [X] **acceptance criteria** of linked ticket are met
- [X] performed **self-review** of the code
- [X] written **tests** for any new functionality added with this PR
- [X] updated the **documentation** if this PR changes or adds functionality
- [ ] written/updated **design docs** if this PR implements new functionality
- [X] updated the **change log**
- [X] ran **pre-commit** checks for code style
- [X] attended to any **TO-DOs** left in the code

Related work items: #2462
This commit is contained in:
Charlie Crane
2024-04-26 14:52:21 +00:00
parent e1ac6255ad
commit 5ee23dcb17
21 changed files with 502 additions and 156 deletions

View File

@@ -5,6 +5,8 @@ from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
@@ -298,7 +300,7 @@ class IOSoftware(Software):
"""Return the public version of connections."""
return copy.copy(self._connections)
def add_connection(self, connection_id: str, session_id: Optional[str] = None) -> bool:
def add_connection(self, connection_id: Union[str, int], session_id: Optional[str] = None) -> bool:
"""
Create a new connection to this service.
@@ -323,6 +325,7 @@ class IOSoftware(Software):
if session_id:
session_details = self._get_session_details(session_id)
self._connections[connection_id] = {
"session_id": session_id,
"ip_address": session_details.with_ip_address if session_details else None,
"time": datetime.now(),
}
@@ -334,19 +337,41 @@ class IOSoftware(Software):
)
return False
def remove_connection(self, connection_id: str) -> bool:
def terminate_connection(self, connection_id: str, send_disconnect: bool = True) -> bool:
"""
Remove a connection from this service.
Terminates a connection from this service.
Returns true if connection successfully removed
:param: connection_id: UUID of the connection to create
:param send_disconnect: If True, sends a disconnect payload to the ip address of the associated session.
:type: string
"""
if self.connections.get(connection_id):
self._connections.pop(connection_id)
self.sys_log.info(f"{self.name}: Connection {connection_id=} closed.")
return True
connection_dict = self._connections.pop(connection_id)
if send_disconnect:
self.software_manager.send_payload_to_session_manager(
payload={"type": "disconnect", "connection_id": connection_id},
session_id=connection_dict["session_id"],
)
self.sys_log.info(f"{self.name}: Connection {connection_id=} terminated")
return True
return False
def show_connections(self, markdown: bool = False):
"""
Display the connections in tabular format.
:param markdown: Whether to display the table in Markdown format or not. Default is `False`.
"""
table = PrettyTable(["IP Address", "Connection ID", "Creation Timestamp"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} {self.name} Connections"
for connection_id, data in self.connections.items():
table.add_row([data["ip_address"], connection_id, str(data["time"])])
print(table.get_string(sortby="Creation Timestamp"))
def clear_connections(self):
"""Clears all the connections from the software."""