Merge remote-tracking branch 'origin/dev' into feature/2736-instantaneous-rewards
This commit is contained in:
@@ -1071,6 +1071,83 @@ class NodeNetworkServiceReconAction(AbstractAction):
|
||||
]
|
||||
|
||||
|
||||
class NodeAccountsChangePasswordAction(AbstractAction):
|
||||
"""Action which changes the password for a user."""
|
||||
|
||||
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
||||
super().__init__(manager=manager)
|
||||
|
||||
def form_request(self, node_id: str, username: str, current_password: str, new_password: str) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
node_name = self.manager.get_node_name_by_idx(node_id)
|
||||
return [
|
||||
"network",
|
||||
"node",
|
||||
node_name,
|
||||
"service",
|
||||
"UserManager",
|
||||
"change_password",
|
||||
username,
|
||||
current_password,
|
||||
new_password,
|
||||
]
|
||||
|
||||
|
||||
class NodeSessionsRemoteLoginAction(AbstractAction):
|
||||
"""Action which performs a remote session login."""
|
||||
|
||||
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
||||
super().__init__(manager=manager)
|
||||
|
||||
def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
node_name = self.manager.get_node_name_by_idx(node_id)
|
||||
return [
|
||||
"network",
|
||||
"node",
|
||||
node_name,
|
||||
"service",
|
||||
"Terminal",
|
||||
"ssh_to_remote",
|
||||
username,
|
||||
password,
|
||||
remote_ip,
|
||||
]
|
||||
|
||||
|
||||
class NodeSessionsRemoteLogoutAction(AbstractAction):
|
||||
"""Action which performs a remote session logout."""
|
||||
|
||||
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
||||
super().__init__(manager=manager)
|
||||
|
||||
def form_request(self, node_id: str, remote_ip: str) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
node_name = self.manager.get_node_name_by_idx(node_id)
|
||||
return ["network", "node", node_name, "service", "Terminal", "remote_logoff", remote_ip]
|
||||
|
||||
|
||||
class NodeSendRemoteCommandAction(AbstractAction):
|
||||
"""Action which sends a terminal command to a remote node via SSH."""
|
||||
|
||||
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
||||
super().__init__(manager=manager)
|
||||
|
||||
def form_request(self, node_id: int, remote_ip: str, command: RequestFormat) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
node_name = self.manager.get_node_name_by_idx(node_id)
|
||||
return [
|
||||
"network",
|
||||
"node",
|
||||
node_name,
|
||||
"service",
|
||||
"Terminal",
|
||||
"send_remote_command",
|
||||
remote_ip,
|
||||
{"command": command},
|
||||
]
|
||||
|
||||
|
||||
class ActionManager:
|
||||
"""Class which manages the action space for an agent."""
|
||||
|
||||
@@ -1122,6 +1199,10 @@ class ActionManager:
|
||||
"CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction,
|
||||
"CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction,
|
||||
"CONFIGURE_DOSBOT": ConfigureDoSBotAction,
|
||||
"NODE_ACCOUNTS_CHANGE_PASSWORD": NodeAccountsChangePasswordAction,
|
||||
"SSH_TO_REMOTE": NodeSessionsRemoteLoginAction,
|
||||
"SESSIONS_REMOTE_LOGOFF": NodeSessionsRemoteLogoutAction,
|
||||
"NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction,
|
||||
}
|
||||
"""Dictionary which maps action type strings to the corresponding action class."""
|
||||
|
||||
|
||||
@@ -36,6 +36,8 @@ class AgentHistoryItem(BaseModel):
|
||||
|
||||
reward: Optional[float] = None
|
||||
|
||||
reward_info: Dict[str, Any] = {}
|
||||
|
||||
|
||||
class AgentStartSettings(BaseModel):
|
||||
"""Configuration values for when an agent starts performing actions."""
|
||||
|
||||
@@ -10,6 +10,7 @@ from primaite import getLogger
|
||||
from primaite.game.agent.observations.acl_observation import ACLObservation
|
||||
from primaite.game.agent.observations.nic_observations import PortObservation
|
||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
@@ -32,6 +33,8 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
"""List of protocols for encoding ACLs."""
|
||||
num_rules: Optional[int] = None
|
||||
"""Number of rules ACL rules to show."""
|
||||
include_users: Optional[bool] = True
|
||||
"""If True, report user session information."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -41,6 +44,7 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
port_list: List[int],
|
||||
protocol_list: List[str],
|
||||
num_rules: int,
|
||||
include_users: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Initialise a firewall observation instance.
|
||||
@@ -58,9 +62,13 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
:type protocol_list: List[str]
|
||||
:param num_rules: Number of rules configured in the firewall.
|
||||
:type num_rules: int
|
||||
:param include_users: If True, report user session information.
|
||||
:type include_users: bool
|
||||
"""
|
||||
self.where: WhereType = where
|
||||
|
||||
self.include_users: bool = include_users
|
||||
self.max_users: int = 3
|
||||
"""Maximum number of remote sessions observable, excess sessions are truncated."""
|
||||
self.ports: List[PortObservation] = [
|
||||
PortObservation(where=self.where + ["NICs", port_num]) for port_num in (1, 2, 3)
|
||||
]
|
||||
@@ -142,6 +150,9 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
:return: Observation containing the status of ports and ACLs for internal, DMZ, and external traffic.
|
||||
:rtype: ObsType
|
||||
"""
|
||||
firewall_state = access_from_nested_dict(state, self.where)
|
||||
if firewall_state is NOT_PRESENT_IN_STATE:
|
||||
return self.default_observation
|
||||
obs = {
|
||||
"PORTS": {i + 1: p.observe(state) for i, p in enumerate(self.ports)},
|
||||
"ACL": {
|
||||
@@ -159,6 +170,12 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
},
|
||||
},
|
||||
}
|
||||
if self.include_users:
|
||||
sess = firewall_state["services"]["UserSessionManager"]
|
||||
obs["users"] = {
|
||||
"local_login": 1 if sess["current_local_user"] else 0,
|
||||
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),
|
||||
}
|
||||
return obs
|
||||
|
||||
@property
|
||||
@@ -218,4 +235,5 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
port_list=config.port_list,
|
||||
protocol_list=config.protocol_list,
|
||||
num_rules=config.num_rules,
|
||||
include_users=config.include_users,
|
||||
)
|
||||
|
||||
@@ -52,6 +52,8 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
"""
|
||||
If True, files and folders must be scanned to update the health state. If False, true state is always shown.
|
||||
"""
|
||||
include_users: Optional[bool] = True
|
||||
"""If True, report user session information."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -69,6 +71,7 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
monitored_traffic: Optional[Dict],
|
||||
include_num_access: bool,
|
||||
file_system_requires_scan: bool,
|
||||
include_users: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Initialise a host observation instance.
|
||||
@@ -103,10 +106,15 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
:param file_system_requires_scan: If True, the files and folders must be scanned to update the health state.
|
||||
If False, the true state is always shown.
|
||||
:type file_system_requires_scan: bool
|
||||
:param include_users: If True, report user session information.
|
||||
:type include_users: bool
|
||||
"""
|
||||
self.where: WhereType = where
|
||||
|
||||
self.include_num_access = include_num_access
|
||||
self.include_users = include_users
|
||||
self.max_users: int = 3
|
||||
"""Maximum number of remote sessions observable, excess sessions are truncated."""
|
||||
|
||||
# Ensure lists have lengths equal to specified counts by truncating or padding
|
||||
self.services: List[ServiceObservation] = services
|
||||
@@ -165,6 +173,8 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
if self.include_num_access:
|
||||
self.default_observation["num_file_creations"] = 0
|
||||
self.default_observation["num_file_deletions"] = 0
|
||||
if self.include_users:
|
||||
self.default_observation["users"] = {"local_login": 0, "remote_sessions": 0}
|
||||
|
||||
def observe(self, state: Dict) -> ObsType:
|
||||
"""
|
||||
@@ -192,6 +202,12 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
if self.include_num_access:
|
||||
obs["num_file_creations"] = node_state["file_system"]["num_file_creations"]
|
||||
obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"]
|
||||
if self.include_users:
|
||||
sess = node_state["services"]["UserSessionManager"]
|
||||
obs["users"] = {
|
||||
"local_login": 1 if sess["current_local_user"] else 0,
|
||||
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),
|
||||
}
|
||||
return obs
|
||||
|
||||
@property
|
||||
@@ -216,6 +232,10 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
if self.include_num_access:
|
||||
shape["num_file_creations"] = spaces.Discrete(4)
|
||||
shape["num_file_deletions"] = spaces.Discrete(4)
|
||||
if self.include_users:
|
||||
shape["users"] = spaces.Dict(
|
||||
{"local_login": spaces.Discrete(2), "remote_sessions": spaces.Discrete(self.max_users + 1)}
|
||||
)
|
||||
return spaces.Dict(shape)
|
||||
|
||||
@classmethod
|
||||
@@ -273,4 +293,5 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
monitored_traffic=config.monitored_traffic,
|
||||
include_num_access=config.include_num_access,
|
||||
file_system_requires_scan=config.file_system_requires_scan,
|
||||
include_users=config.include_users,
|
||||
)
|
||||
|
||||
@@ -46,6 +46,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
|
||||
"""Flag to include the number of accesses."""
|
||||
file_system_requires_scan: bool = True
|
||||
"""If True, the folder must be scanned to update the health state. Tf False, the true state is always shown."""
|
||||
include_users: Optional[bool] = True
|
||||
"""If True, report user session information."""
|
||||
num_ports: Optional[int] = None
|
||||
"""Number of ports."""
|
||||
ip_list: Optional[List[str]] = None
|
||||
@@ -191,6 +193,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
|
||||
host_config.include_num_access = config.include_num_access
|
||||
if host_config.file_system_requires_scan is None:
|
||||
host_config.file_system_requires_scan = config.file_system_requires_scan
|
||||
if host_config.include_users is None:
|
||||
host_config.include_users = config.include_users
|
||||
|
||||
for router_config in config.routers:
|
||||
if router_config.num_ports is None:
|
||||
@@ -205,6 +209,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
|
||||
router_config.protocol_list = config.protocol_list
|
||||
if router_config.num_rules is None:
|
||||
router_config.num_rules = config.num_rules
|
||||
if router_config.include_users is None:
|
||||
router_config.include_users = config.include_users
|
||||
|
||||
for firewall_config in config.firewalls:
|
||||
if firewall_config.ip_list is None:
|
||||
@@ -217,6 +223,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
|
||||
firewall_config.protocol_list = config.protocol_list
|
||||
if firewall_config.num_rules is None:
|
||||
firewall_config.num_rules = config.num_rules
|
||||
if firewall_config.include_users is None:
|
||||
firewall_config.include_users = config.include_users
|
||||
|
||||
hosts = [HostObservation.from_config(config=c, parent_where=where) for c in config.hosts]
|
||||
routers = [RouterObservation.from_config(config=c, parent_where=where) for c in config.routers]
|
||||
|
||||
@@ -39,6 +39,8 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
||||
"""List of protocols for encoding ACLs."""
|
||||
num_rules: Optional[int] = None
|
||||
"""Number of rules ACL rules to show."""
|
||||
include_users: Optional[bool] = True
|
||||
"""If True, report user session information."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -46,6 +48,7 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
||||
ports: List[PortObservation],
|
||||
num_ports: int,
|
||||
acl: ACLObservation,
|
||||
include_users: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Initialise a router observation instance.
|
||||
@@ -59,12 +62,16 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
||||
:type num_ports: int
|
||||
:param acl: ACL observation representing the access control list of the router.
|
||||
:type acl: ACLObservation
|
||||
:param include_users: If True, report user session information.
|
||||
:type include_users: bool
|
||||
"""
|
||||
self.where: WhereType = where
|
||||
self.ports: List[PortObservation] = ports
|
||||
self.acl: ACLObservation = acl
|
||||
self.num_ports: int = num_ports
|
||||
|
||||
self.include_users: bool = include_users
|
||||
self.max_users: int = 3
|
||||
"""Maximum number of remote sessions observable, excess sessions are truncated."""
|
||||
while len(self.ports) < num_ports:
|
||||
self.ports.append(PortObservation(where=None))
|
||||
while len(self.ports) > num_ports:
|
||||
@@ -95,6 +102,12 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
||||
obs["ACL"] = self.acl.observe(state)
|
||||
if self.ports:
|
||||
obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)}
|
||||
if self.include_users:
|
||||
sess = router_state["services"]["UserSessionManager"]
|
||||
obs["users"] = {
|
||||
"local_login": 1 if sess["current_local_user"] else 0,
|
||||
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),
|
||||
}
|
||||
return obs
|
||||
|
||||
@property
|
||||
@@ -143,4 +156,4 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
||||
|
||||
ports = [PortObservation.from_config(config=c, parent_where=where) for c in config.ports]
|
||||
acl = ACLObservation.from_config(config=config.acl, parent_where=where)
|
||||
return cls(where=where, ports=ports, num_ports=config.num_ports, acl=acl)
|
||||
return cls(where=where, ports=ports, num_ports=config.num_ports, acl=acl, include_users=config.include_users)
|
||||
|
||||
@@ -47,7 +47,15 @@ class AbstractReward:
|
||||
|
||||
@abstractmethod
|
||||
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
|
||||
"""Calculate the reward for the current state."""
|
||||
"""Calculate the reward for the current state.
|
||||
|
||||
:param state: Current simulation state
|
||||
:type state: Dict
|
||||
:param last_action_response: Current agent history state
|
||||
:type last_action_response: AgentHistoryItem state
|
||||
:return: Reward value
|
||||
:rtype: float
|
||||
"""
|
||||
return 0.0
|
||||
|
||||
@classmethod
|
||||
@@ -67,7 +75,15 @@ class DummyReward(AbstractReward):
|
||||
"""Dummy reward function component which always returns 0."""
|
||||
|
||||
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
|
||||
"""Calculate the reward for the current state."""
|
||||
"""Calculate the reward for the current state.
|
||||
|
||||
:param state: Current simulation state
|
||||
:type state: Dict
|
||||
:param last_action_response: Current agent history state
|
||||
:type last_action_response: AgentHistoryItem state
|
||||
:return: Reward value
|
||||
:rtype: float
|
||||
"""
|
||||
return 0.0
|
||||
|
||||
@classmethod
|
||||
@@ -109,8 +125,12 @@ class DatabaseFileIntegrity(AbstractReward):
|
||||
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
|
||||
"""Calculate the reward for the current state.
|
||||
|
||||
:param state: The current state of the simulation.
|
||||
:param state: Current simulation state
|
||||
:type state: Dict
|
||||
:param last_action_response: Current agent history state
|
||||
:type last_action_response: AgentHistoryItem state
|
||||
:return: Reward value
|
||||
:rtype: float
|
||||
"""
|
||||
database_file_state = access_from_nested_dict(state, self.location_in_state)
|
||||
if database_file_state is NOT_PRESENT_IN_STATE:
|
||||
@@ -322,6 +342,12 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
|
||||
component will keep track of that information. In that case, it doesn't matter whether the last successful
|
||||
request returned was able to connect to the database server, because there has been an unsuccessful request
|
||||
since.
|
||||
:param state: Current simulation state
|
||||
:type state: Dict
|
||||
:param last_action_response: Current agent history state
|
||||
:type last_action_response: AgentHistoryItem state
|
||||
:return: Reward value
|
||||
:rtype: float
|
||||
"""
|
||||
request_attempted = last_action_response.request == [
|
||||
"network",
|
||||
@@ -333,10 +359,13 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
|
||||
]
|
||||
|
||||
if request_attempted: # if agent makes request, always recalculate fresh value
|
||||
last_action_response.reward_info = {"connection_attempt_status": last_action_response.response.status}
|
||||
self.reward = 1.0 if last_action_response.response.status == "success" else -1.0
|
||||
elif not self.sticky: # if no new request and not sticky, set reward to 0
|
||||
last_action_response.reward_info = {"connection_attempt_status": "n/a"}
|
||||
self.reward = 0.0
|
||||
else: # if no new request and sticky, reuse reward value from last step
|
||||
last_action_response.reward_info = {"connection_attempt_status": "n/a"}
|
||||
pass
|
||||
|
||||
return self.reward
|
||||
@@ -384,7 +413,15 @@ class SharedReward(AbstractReward):
|
||||
"""Method that retrieves an agent's current reward given the agent's name."""
|
||||
|
||||
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
|
||||
"""Simply access the other agent's reward and return it."""
|
||||
"""Simply access the other agent's reward and return it.
|
||||
|
||||
:param state: Current simulation state
|
||||
:type state: Dict
|
||||
:param last_action_response: Current agent history state
|
||||
:type last_action_response: AgentHistoryItem state
|
||||
:return: Reward value
|
||||
:rtype: float
|
||||
"""
|
||||
return self.callback(self.agent_name)
|
||||
|
||||
@classmethod
|
||||
@@ -417,7 +454,15 @@ class ActionPenalty(AbstractReward):
|
||||
self.do_nothing_penalty = do_nothing_penalty
|
||||
|
||||
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
|
||||
"""Calculate the penalty to be applied."""
|
||||
"""Calculate the penalty to be applied.
|
||||
|
||||
:param state: Current simulation state
|
||||
:type state: Dict
|
||||
:param last_action_response: Current agent history state
|
||||
:type last_action_response: AgentHistoryItem state
|
||||
:return: Reward value
|
||||
:rtype: float
|
||||
"""
|
||||
if last_action_response.action == "DONOTHING":
|
||||
return self.do_nothing_penalty
|
||||
else:
|
||||
|
||||
@@ -990,6 +990,7 @@ class UserManager(Service):
|
||||
if user and user.password == current_password:
|
||||
user.password = new_password
|
||||
self.sys_log.info(f"{self.name}: Password changed for {username}")
|
||||
self._user_session_manager._logout_user(user=user)
|
||||
return True
|
||||
self.sys_log.info(f"{self.name}: Password change failed for {username}")
|
||||
return False
|
||||
@@ -1027,6 +1028,10 @@ class UserManager(Service):
|
||||
self.sys_log.info(f"{self.name}: Failed to enable user: {username}")
|
||||
return False
|
||||
|
||||
@property
|
||||
def _user_session_manager(self) -> "UserSessionManager":
|
||||
return self.software_manager.software["UserSessionManager"] # noqa
|
||||
|
||||
|
||||
class UserSession(SimComponent):
|
||||
"""
|
||||
@@ -1260,7 +1265,8 @@ class UserSessionManager(Service):
|
||||
:return: A dictionary representing the current state.
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state["active_remote_logins"] = len(self.remote_sessions)
|
||||
state["current_local_user"] = None if not self.local_session else self.local_session.user.username
|
||||
state["active_remote_sessions"] = list(self.remote_sessions.keys())
|
||||
return state
|
||||
|
||||
@property
|
||||
@@ -1435,6 +1441,19 @@ class UserSessionManager(Service):
|
||||
"""
|
||||
return self._logout(local=False, remote_session_id=remote_session_id)
|
||||
|
||||
def _logout_user(self, user: Union[str, User]) -> bool:
|
||||
"""End a user session by username or user object."""
|
||||
if isinstance(user, str):
|
||||
user = self._user_manager.users[user] # grab user object from username
|
||||
for sess_id, session in self.remote_sessions.items():
|
||||
if session.user is user:
|
||||
self._logout(local=False, remote_session_id=sess_id)
|
||||
return True
|
||||
if self.local_user_logged_in and self.local_session.user is user:
|
||||
self.local_logout()
|
||||
return True
|
||||
return False
|
||||
|
||||
@property
|
||||
def local_user_logged_in(self) -> bool:
|
||||
"""
|
||||
|
||||
@@ -23,6 +23,8 @@ from primaite.simulator.system.core.software_manager import SoftwareManager
|
||||
from primaite.simulator.system.services.service import Service, ServiceOperatingState
|
||||
|
||||
|
||||
# TODO 2824: Since remote terminal connections and remote user sessions are the same thing, we could refactor
|
||||
# the terminal to leverage the user session manager's list. This way we avoid potential bugs and code ducplication
|
||||
class TerminalClientConnection(BaseModel):
|
||||
"""
|
||||
TerminalClientConnection Class.
|
||||
@@ -92,7 +94,7 @@ class LocalTerminalConnection(TerminalClientConnection):
|
||||
if not self.is_active:
|
||||
self.parent_terminal.sys_log.warning("Connection inactive, cannot execute")
|
||||
return None
|
||||
return self.parent_terminal.execute(command, connection_id=self.connection_uuid)
|
||||
return self.parent_terminal.execute(command)
|
||||
|
||||
|
||||
class RemoteTerminalConnection(TerminalClientConnection):
|
||||
@@ -162,22 +164,6 @@ class Terminal(Service):
|
||||
def _init_request_manager(self) -> RequestManager:
|
||||
"""Initialise Request manager."""
|
||||
rm = super()._init_request_manager()
|
||||
rm.add_request(
|
||||
"send",
|
||||
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())),
|
||||
)
|
||||
|
||||
def _login(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
login = self._process_local_login(username=request[0], password=request[1])
|
||||
if login:
|
||||
return RequestResponse(
|
||||
status="success",
|
||||
data={
|
||||
"ip_address": login.ip_address,
|
||||
},
|
||||
)
|
||||
else:
|
||||
return RequestResponse(status="failure", data={"reason": "Invalid login credentials"})
|
||||
|
||||
def _remote_login(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
login = self._send_remote_login(username=request[0], password=request[1], ip_address=request[2])
|
||||
@@ -191,10 +177,34 @@ class Terminal(Service):
|
||||
else:
|
||||
return RequestResponse(status="failure", data={})
|
||||
|
||||
rm.add_request(
|
||||
"ssh_to_remote",
|
||||
request_type=RequestType(func=_remote_login),
|
||||
)
|
||||
|
||||
def _remote_logoff(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
"""Logoff from remote connection."""
|
||||
ip_address = IPv4Address(request[0])
|
||||
remote_connection = self._get_connection_from_ip(ip_address=ip_address)
|
||||
if remote_connection:
|
||||
outcome = self._disconnect(remote_connection.connection_uuid)
|
||||
if outcome:
|
||||
return RequestResponse(
|
||||
status="success",
|
||||
data={},
|
||||
)
|
||||
else:
|
||||
return RequestResponse(
|
||||
status="failure",
|
||||
data={"reason": "No remote connection held."},
|
||||
)
|
||||
|
||||
rm.add_request("remote_logoff", request_type=RequestType(func=_remote_logoff))
|
||||
|
||||
def remote_execute_request(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
"""Execute an instruction."""
|
||||
command: str = request[0]
|
||||
ip_address: IPv4Address = IPv4Address(request[1])
|
||||
ip_address: IPv4Address = IPv4Address(request[0])
|
||||
command: str = request[1]["command"]
|
||||
remote_connection = self._get_connection_from_ip(ip_address=ip_address)
|
||||
if remote_connection:
|
||||
outcome = remote_connection.execute(command)
|
||||
@@ -209,30 +219,11 @@ class Terminal(Service):
|
||||
data={},
|
||||
)
|
||||
|
||||
def _logoff(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
"""Logoff from connection."""
|
||||
connection_uuid = request[0]
|
||||
self.parent.user_session_manager.local_logout(connection_uuid)
|
||||
self._disconnect(connection_uuid)
|
||||
return RequestResponse(status="success", data={})
|
||||
|
||||
rm.add_request(
|
||||
"Login",
|
||||
request_type=RequestType(func=_login),
|
||||
)
|
||||
|
||||
rm.add_request(
|
||||
"Remote Login",
|
||||
request_type=RequestType(func=_remote_login),
|
||||
)
|
||||
|
||||
rm.add_request(
|
||||
"Execute",
|
||||
"send_remote_command",
|
||||
request_type=RequestType(func=remote_execute_request),
|
||||
)
|
||||
|
||||
rm.add_request("Logoff", request_type=RequestType(func=_logoff))
|
||||
|
||||
return rm
|
||||
|
||||
def execute(self, command: List[Any]) -> Optional[RequestResponse]:
|
||||
@@ -280,13 +271,9 @@ class Terminal(Service):
|
||||
if self.operating_state != ServiceOperatingState.RUNNING:
|
||||
self.sys_log.warning(f"{self.name}: Cannot login as service is not running.")
|
||||
return None
|
||||
connection_request_id = str(uuid4())
|
||||
self._client_connection_requests[connection_request_id] = None
|
||||
if ip_address:
|
||||
# Assuming that if IP is passed we are connecting to remote
|
||||
return self._send_remote_login(
|
||||
username=username, password=password, ip_address=ip_address, connection_request_id=connection_request_id
|
||||
)
|
||||
return self._send_remote_login(username=username, password=password, ip_address=ip_address)
|
||||
else:
|
||||
return self._process_local_login(username=username, password=password)
|
||||
|
||||
@@ -313,6 +300,9 @@ class Terminal(Service):
|
||||
|
||||
def _check_client_connection(self, connection_id: str) -> bool:
|
||||
"""Check that client_connection_id is valid."""
|
||||
if not self.parent.user_session_manager.validate_remote_session_uuid(connection_id):
|
||||
self._disconnect(connection_id)
|
||||
return False
|
||||
return connection_id in self._connections
|
||||
|
||||
def _send_remote_login(
|
||||
@@ -320,32 +310,24 @@ class Terminal(Service):
|
||||
username: str,
|
||||
password: str,
|
||||
ip_address: IPv4Address,
|
||||
connection_request_id: str,
|
||||
connection_request_id: Optional[str] = None,
|
||||
is_reattempt: bool = False,
|
||||
) -> Optional[RemoteTerminalConnection]:
|
||||
"""Send a remote login attempt and connect to Node.
|
||||
|
||||
:param: username: Username used to connect to the remote node.
|
||||
:type: username: str
|
||||
|
||||
:param: password: Password used to connect to the remote node
|
||||
:type: password: str
|
||||
|
||||
:param: ip_address: Target Node IP address for login attempt.
|
||||
:type: ip_address: IPv4Address
|
||||
|
||||
:param: connection_request_id: Connection Request ID
|
||||
:type: connection_request_id: str
|
||||
|
||||
:param: connection_request_id: Connection Request ID, if not provided, a new one is generated
|
||||
:type: connection_request_id: Optional[str]
|
||||
:param: is_reattempt: True if the request has been reattempted. Default False.
|
||||
:type: is_reattempt: Optional[bool]
|
||||
|
||||
:return: RemoteTerminalConnection: Connection Object for sending further commands if successful, else False.
|
||||
|
||||
"""
|
||||
self.sys_log.info(
|
||||
f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}"
|
||||
)
|
||||
connection_request_id = connection_request_id or str(uuid4())
|
||||
if is_reattempt:
|
||||
valid_connection_request = self._validate_client_connection_request(connection_id=connection_request_id)
|
||||
if valid_connection_request:
|
||||
@@ -360,6 +342,9 @@ class Terminal(Service):
|
||||
self.sys_log.warning(f"{self.name}: Remote connection to {ip_address} declined.")
|
||||
return None
|
||||
|
||||
self.sys_log.info(
|
||||
f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}"
|
||||
)
|
||||
transport_message: SSHTransportMessage = SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST
|
||||
connection_message: SSHConnectionMessage = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA
|
||||
user_details: SSHUserCredentials = SSHUserCredentials(username=username, password=password)
|
||||
|
||||
Reference in New Issue
Block a user