#3062 - rename identifier to discriminator
This commit is contained in:
@@ -172,7 +172,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# no applications exist yet so we will create our own.\n",
|
||||
"class MSPaint(Application, identifier=\"MSPaint\"):\n",
|
||||
"class MSPaint(Application, discriminator=\"MSPaint\"):\n",
|
||||
" def describe_state(self):\n",
|
||||
" return super().describe_state()"
|
||||
]
|
||||
|
||||
@@ -244,7 +244,7 @@ class SimComponent(BaseModel):
|
||||
|
||||
..code::python
|
||||
|
||||
class WebBrowser(Application, identifier="WebBrowser"):
|
||||
class WebBrowser(Application, discriminator="WebBrowser"):
|
||||
def _init_request_manager(self) -> RequestManager:
|
||||
rm = super()._init_request_manager() # all requests generic to any Application get initialised
|
||||
rm.add_request(...) # initialise any requests specific to the web browser
|
||||
|
||||
@@ -180,7 +180,7 @@ class Network(SimComponent):
|
||||
table.align = "l"
|
||||
table.title = "Nodes"
|
||||
for node in self.nodes.values():
|
||||
table.add_row((node.hostname, type(node)._identifier, node.operating_state.name))
|
||||
table.add_row((node.hostname, type(node)._discriminator, node.operating_state.name))
|
||||
print(table)
|
||||
|
||||
if ip_addresses:
|
||||
|
||||
@@ -22,7 +22,7 @@ class NetworkNodeAdder(BaseModel):
|
||||
|
||||
Here is a template that users can use to define custom node adders:
|
||||
```
|
||||
class YourNodeAdder(NetworkNodeAdder, identifier="your_name"):
|
||||
class YourNodeAdder(NetworkNodeAdder, discriminator="your_name"):
|
||||
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
|
||||
property_1 : str
|
||||
property_2 : int
|
||||
@@ -40,8 +40,8 @@ class NetworkNodeAdder(BaseModel):
|
||||
"""
|
||||
Base schema for node adders.
|
||||
|
||||
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The identifier is used
|
||||
by the from_config method to select the correct node adder at runtime.
|
||||
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The discriminator is
|
||||
used by the from_config method to select the correct node adder at runtime.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
@@ -50,20 +50,20 @@ class NetworkNodeAdder(BaseModel):
|
||||
|
||||
_registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {}
|
||||
|
||||
def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None:
|
||||
def __init_subclass__(cls, discriminator: Optional[str], **kwargs: Any) -> None:
|
||||
"""
|
||||
Register a network node adder class.
|
||||
|
||||
:param identifier: Unique name for the node adder to use for matching against primaite config entries.
|
||||
:type identifier: str
|
||||
:param discriminator: Unique name for the node adder to use for matching against primaite config entries.
|
||||
:type discriminator: str
|
||||
:raises ValueError: When attempting to register a name that is already reserved.
|
||||
"""
|
||||
super().__init_subclass__(**kwargs)
|
||||
if identifier is None:
|
||||
if discriminator is None:
|
||||
return
|
||||
if identifier in cls._registry:
|
||||
raise ValueError(f"Duplicate node adder {identifier}")
|
||||
cls._registry[identifier] = cls
|
||||
if discriminator in cls._registry:
|
||||
raise ValueError(f"Duplicate node adder {discriminator}")
|
||||
cls._registry[discriminator] = cls
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
@@ -99,7 +99,7 @@ class NetworkNodeAdder(BaseModel):
|
||||
adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network)
|
||||
|
||||
|
||||
class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"):
|
||||
class OfficeLANAdder(NetworkNodeAdder, discriminator="office_lan"):
|
||||
"""Creates an office LAN."""
|
||||
|
||||
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
|
||||
|
||||
@@ -824,7 +824,7 @@ class User(SimComponent):
|
||||
return self.model_dump()
|
||||
|
||||
|
||||
class UserManager(Service, identifier="UserManager"):
|
||||
class UserManager(Service, discriminator="UserManager"):
|
||||
"""
|
||||
Manages users within the PrimAITE system, handling creation, authentication, and administration.
|
||||
|
||||
@@ -1137,7 +1137,7 @@ class RemoteUserSession(UserSession):
|
||||
return state
|
||||
|
||||
|
||||
class UserSessionManager(Service, identifier="UserSessionManager"):
|
||||
class UserSessionManager(Service, discriminator="UserSessionManager"):
|
||||
"""
|
||||
Manages user sessions on a Node, including local and remote sessions.
|
||||
|
||||
@@ -1483,7 +1483,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
|
||||
return self.local_session is not None
|
||||
|
||||
|
||||
class Node(SimComponent):
|
||||
class Node(SimComponent, ABC):
|
||||
"""
|
||||
A basic Node class that represents a node on the network.
|
||||
|
||||
@@ -1556,25 +1556,26 @@ class Node(SimComponent):
|
||||
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
|
||||
"""Registry of application types. Automatically populated when subclasses are defined."""
|
||||
|
||||
_identifier: ClassVar[str] = "unknown"
|
||||
"""Identifier for this particular class, used for printing and logging. Each subclass redefines this."""
|
||||
# TODO: this should not be set for abstract classes.
|
||||
_discriminator: ClassVar[str]
|
||||
"""discriminator for this particular class, used for printing and logging. Each subclass redefines this."""
|
||||
|
||||
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
|
||||
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
|
||||
"""
|
||||
Register a node type.
|
||||
|
||||
:param identifier: Uniquely specifies an node class by name. Used for finding items by config.
|
||||
:type identifier: str
|
||||
:param discriminator: Uniquely specifies an node class by name. Used for finding items by config.
|
||||
:type discriminator: str
|
||||
:raises ValueError: When attempting to register an node with a name that is already allocated.
|
||||
"""
|
||||
super().__init_subclass__(**kwargs)
|
||||
if identifier is None:
|
||||
if discriminator is None:
|
||||
return
|
||||
identifier = identifier.lower()
|
||||
if identifier in cls._registry:
|
||||
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
|
||||
cls._registry[identifier] = cls
|
||||
cls._identifier = identifier
|
||||
discriminator = discriminator.lower()
|
||||
if discriminator in cls._registry:
|
||||
raise ValueError(f"Tried to define new node {discriminator}, but this name is already reserved.")
|
||||
cls._registry[discriminator] = cls
|
||||
cls._discriminator = discriminator
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""
|
||||
|
||||
@@ -5,7 +5,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
|
||||
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
|
||||
|
||||
|
||||
class Computer(HostNode, identifier="computer"):
|
||||
class Computer(HostNode, discriminator="computer"):
|
||||
"""
|
||||
A basic Computer class.
|
||||
|
||||
|
||||
@@ -262,7 +262,7 @@ class NIC(IPWiredNetworkInterface):
|
||||
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}"
|
||||
|
||||
|
||||
class HostNode(Node, identifier="HostNode"):
|
||||
class HostNode(Node, discriminator="HostNode"):
|
||||
"""
|
||||
Represents a host node in the network.
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
|
||||
|
||||
|
||||
class Server(HostNode, identifier="server"):
|
||||
class Server(HostNode, discriminator="server"):
|
||||
"""
|
||||
A basic Server class.
|
||||
|
||||
@@ -31,7 +31,7 @@ class Server(HostNode, identifier="server"):
|
||||
"""
|
||||
|
||||
|
||||
class Printer(HostNode, identifier="printer"):
|
||||
class Printer(HostNode, discriminator="printer"):
|
||||
"""Printer? I don't even know her!."""
|
||||
|
||||
# TODO: Implement printer-specific behaviour
|
||||
|
||||
@@ -27,7 +27,7 @@ DMZ_PORT_ID: Final[int] = 3
|
||||
"""The Firewall port ID of the DMZ port."""
|
||||
|
||||
|
||||
class Firewall(Router, identifier="firewall"):
|
||||
class Firewall(Router, discriminator="firewall"):
|
||||
"""
|
||||
A Firewall class that extends the functionality of a Router.
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ from primaite.simulator.network.transmission.data_link_layer import Frame
|
||||
from primaite.simulator.system.services.arp.arp import ARP
|
||||
|
||||
|
||||
class NetworkNode(Node, identifier="NetworkNode"):
|
||||
class NetworkNode(Node, discriminator="NetworkNode"):
|
||||
"""
|
||||
Represents an abstract base class for a network node that can receive and process network frames.
|
||||
|
||||
|
||||
@@ -1184,7 +1184,7 @@ class RouterSessionManager(SessionManager):
|
||||
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
|
||||
|
||||
|
||||
class Router(NetworkNode, identifier="router"):
|
||||
class Router(NetworkNode, discriminator="router"):
|
||||
"""
|
||||
Represents a network router, managing routing and forwarding of IP packets across network interfaces.
|
||||
|
||||
|
||||
@@ -87,7 +87,7 @@ class SwitchPort(WiredNetworkInterface):
|
||||
return False
|
||||
|
||||
|
||||
class Switch(NetworkNode, identifier="switch"):
|
||||
class Switch(NetworkNode, discriminator="switch"):
|
||||
"""
|
||||
A class representing a Layer 2 network switch.
|
||||
|
||||
|
||||
@@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface):
|
||||
)
|
||||
|
||||
|
||||
class WirelessRouter(Router, identifier="wireless_router"):
|
||||
class WirelessRouter(Router, discriminator="wireless_router"):
|
||||
"""
|
||||
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.
|
||||
|
||||
|
||||
@@ -53,20 +53,20 @@ class Application(IOSoftware, ABC):
|
||||
_registry: ClassVar[Dict[str, Type["Application"]]] = {}
|
||||
"""Registry of application types. Automatically populated when subclasses are defined."""
|
||||
|
||||
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
|
||||
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
|
||||
"""
|
||||
Register an application type.
|
||||
|
||||
:param identifier: Uniquely specifies an application class by name. Used for finding items by config.
|
||||
:type identifier: Optional[str]
|
||||
:param discriminator: Uniquely specifies an application class by name. Used for finding items by config.
|
||||
:type discriminator: Optional[str]
|
||||
:raises ValueError: When attempting to register an application with a name that is already allocated.
|
||||
"""
|
||||
super().__init_subclass__(**kwargs)
|
||||
if identifier is None:
|
||||
if discriminator is None:
|
||||
return
|
||||
if identifier in cls._registry:
|
||||
raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.")
|
||||
cls._registry[identifier] = cls
|
||||
if discriminator in cls._registry:
|
||||
raise ValueError(f"Tried to define new application {discriminator}, but this name is already reserved.")
|
||||
cls._registry[discriminator] = cls
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict) -> "Application":
|
||||
|
||||
@@ -61,7 +61,7 @@ class DatabaseClientConnection(BaseModel):
|
||||
return str(self)
|
||||
|
||||
|
||||
class DatabaseClient(Application, identifier="DatabaseClient"):
|
||||
class DatabaseClient(Application, discriminator="DatabaseClient"):
|
||||
"""
|
||||
A DatabaseClient application.
|
||||
|
||||
|
||||
@@ -44,7 +44,7 @@ class PortScanPayload(SimComponent):
|
||||
return state
|
||||
|
||||
|
||||
class NMAP(Application, identifier="NMAP"):
|
||||
class NMAP(Application, discriminator="NMAP"):
|
||||
"""
|
||||
A class representing the NMAP application for network scanning.
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
|
||||
from primaite.utils.validation.port import Port, PORT_LOOKUP
|
||||
|
||||
|
||||
class C2Beacon(AbstractC2, identifier="C2Beacon"):
|
||||
class C2Beacon(AbstractC2, discriminator="C2Beacon"):
|
||||
"""
|
||||
C2 Beacon Application.
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ from primaite.simulator.system.applications.red_applications.c2 import (
|
||||
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
|
||||
|
||||
|
||||
class C2Server(AbstractC2, identifier="C2Server"):
|
||||
class C2Server(AbstractC2, discriminator="C2Server"):
|
||||
"""
|
||||
C2 Server Application.
|
||||
|
||||
|
||||
@@ -40,7 +40,7 @@ class DataManipulationAttackStage(IntEnum):
|
||||
"Signifies that the attack has failed."
|
||||
|
||||
|
||||
class DataManipulationBot(Application, identifier="DataManipulationBot"):
|
||||
class DataManipulationBot(Application, discriminator="DataManipulationBot"):
|
||||
"""A bot that simulates a script which performs a SQL injection attack."""
|
||||
|
||||
class ConfigSchema(Application.ConfigSchema):
|
||||
|
||||
@@ -32,7 +32,7 @@ class DoSAttackStage(IntEnum):
|
||||
"Attack is completed."
|
||||
|
||||
|
||||
class DoSBot(DatabaseClient, identifier="DoSBot"):
|
||||
class DoSBot(DatabaseClient, discriminator="DoSBot"):
|
||||
"""A bot that simulates a Denial of Service attack."""
|
||||
|
||||
class ConfigSchema(DatabaseClient.ConfigSchema):
|
||||
|
||||
@@ -14,7 +14,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
|
||||
from primaite.utils.validation.port import PORT_LOOKUP
|
||||
|
||||
|
||||
class RansomwareScript(Application, identifier="RansomwareScript"):
|
||||
class RansomwareScript(Application, discriminator="RansomwareScript"):
|
||||
"""Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
|
||||
|
||||
:ivar payload: The attack stage query payload. (Default ENCRYPT)
|
||||
|
||||
@@ -23,7 +23,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class WebBrowser(Application, identifier="WebBrowser"):
|
||||
class WebBrowser(Application, discriminator="WebBrowser"):
|
||||
"""
|
||||
Represents a web browser in the simulation environment.
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
|
||||
from primaite.utils.validation.port import PORT_LOOKUP
|
||||
|
||||
|
||||
class ARP(Service, identifier="ARP"):
|
||||
class ARP(Service, discriminator="ARP"):
|
||||
"""
|
||||
The ARP (Address Resolution Protocol) Service.
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class DatabaseService(Service, identifier="DatabaseService"):
|
||||
class DatabaseService(Service, discriminator="DatabaseService"):
|
||||
"""
|
||||
A class for simulating a generic SQL Server service.
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ if TYPE_CHECKING:
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class DNSClient(Service, identifier="DNSClient"):
|
||||
class DNSClient(Service, discriminator="DNSClient"):
|
||||
"""Represents a DNS Client as a Service."""
|
||||
|
||||
class ConfigSchema(Service.ConfigSchema):
|
||||
|
||||
@@ -14,7 +14,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class DNSServer(Service, identifier="DNSServer"):
|
||||
class DNSServer(Service, discriminator="DNSServer"):
|
||||
"""Represents a DNS Server as a Service."""
|
||||
|
||||
class ConfigSchema(Service.ConfigSchema):
|
||||
|
||||
@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class FTPClient(FTPServiceABC, identifier="FTPClient"):
|
||||
class FTPClient(FTPServiceABC, discriminator="FTPClient"):
|
||||
"""
|
||||
A class for simulating an FTP client service.
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class FTPServer(FTPServiceABC, identifier="FTPServer"):
|
||||
class FTPServer(FTPServiceABC, discriminator="FTPServer"):
|
||||
"""
|
||||
A class for simulating an FTP server service.
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class ICMP(Service, identifier="ICMP"):
|
||||
class ICMP(Service, discriminator="ICMP"):
|
||||
"""
|
||||
The Internet Control Message Protocol (ICMP) service.
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class NTPClient(Service, identifier="NTPClient"):
|
||||
class NTPClient(Service, discriminator="NTPClient"):
|
||||
"""Represents a NTP client as a service."""
|
||||
|
||||
class ConfigSchema(Service.ConfigSchema):
|
||||
|
||||
@@ -13,7 +13,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class NTPServer(Service, identifier="NTPServer"):
|
||||
class NTPServer(Service, discriminator="NTPServer"):
|
||||
"""Represents a NTP server as a service."""
|
||||
|
||||
class ConfigSchema(Service.ConfigSchema):
|
||||
|
||||
@@ -61,22 +61,22 @@ class Service(IOSoftware):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
|
||||
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
|
||||
"""
|
||||
Register a hostnode type.
|
||||
|
||||
:param identifier: Uniquely specifies an hostnode class by name. Used for finding items by config.
|
||||
:type identifier: str
|
||||
:param discriminator: Uniquely specifies an hostnode class by name. Used for finding items by config.
|
||||
:type discriminator: str
|
||||
:raises ValueError: When attempting to register an hostnode with a name that is already allocated.
|
||||
"""
|
||||
super().__init_subclass__(**kwargs)
|
||||
if identifier is None:
|
||||
if discriminator is None:
|
||||
return
|
||||
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
|
||||
identifier = identifier.lower()
|
||||
if identifier in cls._registry:
|
||||
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.")
|
||||
cls._registry[identifier] = cls
|
||||
discriminator = discriminator.lower()
|
||||
if discriminator in cls._registry:
|
||||
raise ValueError(f"Tried to define new hostnode {discriminator}, but this name is already reserved.")
|
||||
cls._registry[discriminator] = cls
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict) -> "Service":
|
||||
|
||||
@@ -129,7 +129,7 @@ class RemoteTerminalConnection(TerminalClientConnection):
|
||||
return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id)
|
||||
|
||||
|
||||
class Terminal(Service, identifier="Terminal"):
|
||||
class Terminal(Service, discriminator="Terminal"):
|
||||
"""Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH."""
|
||||
|
||||
class ConfigSchema(Service.ConfigSchema):
|
||||
|
||||
@@ -21,7 +21,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class WebServer(Service, identifier="WebServer"):
|
||||
class WebServer(Service, discriminator="WebServer"):
|
||||
"""Class used to represent a Web Server Service in simulation."""
|
||||
|
||||
class ConfigSchema(Service.ConfigSchema):
|
||||
|
||||
Reference in New Issue
Block a user