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

Merge remote-tracking branch 'origin/dev' into feature/2769-implement-user-account-action-space

Browse Source
This commit is contained in:
Marek Wolan
2024-08-12 09:23:55 +01:00
parent 3df55a708d 31c46950ad
commit 5956721804
12 changed files with 299 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/primaite/game/game.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
"""PrimAITE game - Encapsulates the simulation and agents."""
from ipaddress import IPv4Address
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
import numpy as np
from pydantic import BaseModel, ConfigDict
@@ -44,8 +44,10 @@ from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.services.terminal.terminal import Terminal
from primaite.simulator.system.services.web_server.web_server import WebServer
from primaite.simulator.system.software import Software
_LOGGER = getLogger(__name__)
@@ -328,6 +330,20 @@ class PrimaiteGame:
user_manager: UserManager = new_node.software_manager.software["UserManager"] # noqa
for user_cfg in node_cfg["users"]:
user_manager.add_user(**user_cfg, bypass_can_perform_action=True)
def _set_software_listen_on_ports(software: Union[Software, Service], software_cfg: Dict):
"""Set listener ports on software."""
listen_on_ports = []
for port_id in set(software_cfg.get("options", {}).get("listen_on_ports", [])):
port = None
if isinstance(port_id, int):
port = Port(port_id)
elif isinstance(port_id, str):
port = Port[port_id]
if port:
listen_on_ports.append(port)
software.listen_on_ports = set(listen_on_ports)
if "services" in node_cfg:
for service_cfg in node_cfg["services"]:
new_service = None
@@ -341,6 +357,7 @@ class PrimaiteGame:
if "fix_duration" in service_cfg.get("options", {}):
new_service.fixing_duration = service_cfg["options"]["fix_duration"]
_set_software_listen_on_ports(new_service, service_cfg)
# start the service
new_service.start()
else:
@@ -390,6 +407,8 @@ class PrimaiteGame:
_LOGGER.error(msg)
raise ValueError(msg)
_set_software_listen_on_ports(new_application, application_cfg)
# run the application
new_application.run()

28
src/primaite/simulator/system/core/software_manager.py
View File

@@ -1,4 +1,5 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from copy import deepcopy
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
@@ -76,6 +77,8 @@ class SoftwareManager:
for software in self.port_protocol_mapping.values():
if software.operating_state in {ApplicationOperatingState.RUNNING, ServiceOperatingState.RUNNING}:
open_ports.append(software.port)
if software.listen_on_ports:
open_ports += list(software.listen_on_ports)
return open_ports
def check_port_is_open(self, port: Port, protocol: IPProtocol) -> bool:
@@ -223,7 +226,9 @@ class SoftwareManager:
frame: Frame,
):
"""
Receive a payload from the SessionManager and forward it to the corresponding service or application.
Receive a payload from the SessionManager and forward it to the corresponding service or applications.
This function handles both software assigned a specific port, and software listening in on other ports.
:param payload: The payload being received.
:param session: The transport session the payload originates from.
@@ -231,14 +236,25 @@ class SoftwareManager:
if payload.__class__.__name__ == "PortScanPayload":
self.software.get("NMAP").receive(payload=payload, session_id=session_id)
return
receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None)
if receiver:
receiver.receive(
main_receiver = self.port_protocol_mapping.get((port, protocol), None)
if main_receiver:
main_receiver.receive(
payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame
)
else:
listening_receivers = [
software
for software in self.software.values()
if port in software.listen_on_ports and software != main_receiver
]
for receiver in listening_receivers:
receiver.receive(
payload=deepcopy(payload),
session_id=session_id,
from_network_interface=from_network_interface,
frame=frame,
)
if not main_receiver and not listening_receivers:
self.sys_log.warning(f"No service or application found for port {port} and protocol {protocol}")
pass
def show(self, markdown: bool = False):
"""

2
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -385,6 +385,8 @@ class DatabaseService(Service):
)
else:
result = {"status_code": 401, "type": "sql"}
else:
self.sys_log.info(f"{self.name}: Ignoring payload as it is not a Database payload")
self.send(payload=result, session_id=session_id)
return True

5
src/primaite/simulator/system/software.py
View File

@@ -4,9 +4,10 @@ from abc import abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, TYPE_CHECKING, Union
from typing import Any, Dict, Optional, Set, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
@@ -252,6 +253,8 @@ class IOSoftware(Software):
"Indicates if the software uses UDP protocol for communication. Default is True."
port: Port
"The port to which the software is connected."
listen_on_ports: Set[Port] = Field(default_factory=set)
"The set of ports to listen on."
protocol: IPProtocol
"The IP Protocol the Software operates on."
_connections: Dict[str, Dict] = {}