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:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user