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

#3029 - Add string-ip validator, improve validation, fix minor bugs in pulling schema data

Browse Source
This commit is contained in:
Marek Wolan
2025-01-31 12:18:52 +00:00
parent 09bfa7b700
commit 4fb54c9492
19 changed files with 116 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/primaite/simulator/network/creation.py
View File

@@ -3,7 +3,7 @@ from abc import ABC, abstractmethod
from ipaddress import IPv4Address
from typing import Any, ClassVar, Dict, Literal, Optional, Type
from pydantic import BaseModel, model_validator
from pydantic import BaseModel, ConfigDict, model_validator
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.computer import Computer
@@ -44,6 +44,7 @@ class NetworkNodeAdder(BaseModel):
by the from_config method to select the correct node adder at runtime.
"""
model_config = ConfigDict(extra="forbid")
type: str
"""Uniquely identifies the node adder class to use for adding nodes to network."""

4
src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py
View File

@@ -127,8 +127,8 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
self.configure(
c2_server_ip_address=c2_remote_ip,
keep_alive_frequency=frequency,
masquerade_protocol=PROTOCOL_LOOKUP[protocol],
masquerade_port=PORT_LOOKUP[port],
masquerade_protocol=protocol,
masquerade_port=port,
)
)

8
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -10,8 +10,8 @@ from primaite.game.science import simulate_trial
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
from primaite.utils.validation.ipv4_address import ipv4_validator, IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP, port_validator
_LOGGER = getLogger(__name__)
@@ -106,9 +106,9 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
:rtype: RequestResponse
"""
if "target_ip_address" in request[-1]:
request[-1]["target_ip_address"] = IPv4Address(request[-1]["target_ip_address"])
request[-1]["target_ip_address"] = ipv4_validator(request[-1]["target_ip_address"])
if "target_port" in request[-1]:
request[-1]["target_port"] = PORT_LOOKUP[request[-1]["target_port"]]
request[-1]["target_port"] = port_validator(request[-1]["target_port"])
return RequestResponse.from_bool(self.configure(**request[-1]))
rm.add_request("configure", request_type=RequestType(func=_configure))

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

@@ -31,6 +31,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
type: str = "DatabaseService"
backup_server_ip: Optional[IPv4Address] = None
db_password: Optional[str] = None
config: "DatabaseService.ConfigSchema" = Field(default_factory=lambda: DatabaseService.ConfigSchema())
@@ -53,6 +54,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
super().__init__(**kwargs)
self._create_db_file()
self.backup_server_ip = self.config.backup_server_ip
self.password = self.config.db_password
def install(self):
"""

60
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Dict, Optional
from typing import Dict, Optional, TYPE_CHECKING
from pydantic import Field
@@ -9,8 +9,12 @@ from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest
from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.service import Service
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
if TYPE_CHECKING:
from primaite.simulator.network.hardware.base import Node
_LOGGER = getLogger(__name__)
@@ -21,6 +25,7 @@ class DNSClient(Service, identifier="DNSClient"):
"""ConfigSchema for DNSClient."""
type: str = "DNSClient"
dns_server: Optional[IPV4Address] = None
config: "DNSClient.ConfigSchema" = Field(default_factory=lambda: DNSClient.ConfigSchema())
dns_cache: Dict[str, IPv4Address] = {}
@@ -36,6 +41,7 @@ class DNSClient(Service, identifier="DNSClient"):
# TCP for now
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.dns_server = self.config.dns_server
self.start()
def describe_state(self) -> Dict:
@@ -79,6 +85,14 @@ class DNSClient(Service, identifier="DNSClient"):
if not self._can_perform_action():
return False
# check if the domain is already in the DNS cache
if target_domain in self.dns_cache:
self.sys_log.info(
f"{self.name}: Domain lookup for {target_domain} successful,"
f"resolves to {self.dns_cache[target_domain]}"
)
return True
# check if DNS server is configured
if self.dns_server is None:
self.sys_log.warning(f"{self.name}: DNS Server is not configured")
@@ -87,31 +101,23 @@ class DNSClient(Service, identifier="DNSClient"):
# check if the target domain is in the client's DNS cache
payload = DNSPacket(dns_request=DNSRequest(domain_name_request=target_domain))
# check if the domain is already in the DNS cache
if target_domain in self.dns_cache:
self.sys_log.info(
f"{self.name}: Domain lookup for {target_domain} successful,"
f"resolves to {self.dns_cache[target_domain]}"
)
return True
# return False if already reattempted
if is_reattempt:
self.sys_log.warning(f"{self.name}: Domain lookup for {target_domain} failed")
return False
else:
# return False if already reattempted
if is_reattempt:
self.sys_log.warning(f"{self.name}: Domain lookup for {target_domain} failed")
return False
else:
# send a request to check if domain name exists in the DNS Server
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=self.dns_server, dest_port=PORT_LOOKUP["DNS"]
)
# send a request to check if domain name exists in the DNS Server
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=self.dns_server, dest_port=PORT_LOOKUP["DNS"]
)
# recursively re-call the function passing is_reattempt=True
return self.check_domain_exists(
target_domain=target_domain,
session_id=session_id,
is_reattempt=True,
)
# recursively re-call the function passing is_reattempt=True
return self.check_domain_exists(
target_domain=target_domain,
session_id=session_id,
is_reattempt=True,
)
def send(
self,
@@ -168,3 +174,9 @@ class DNSClient(Service, identifier="DNSClient"):
self.sys_log.warning(f"Failed to resolve domain name {payload.dns_request.domain_name_request}")
return False
def install(self) -> None:
"""Set the DNS server to be the node's DNS server unless a different one was already provided."""
self.parent: Node
if self.parent and not self.dns_server:
self.dns_server = self.parent.dns_server

1
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -36,6 +36,7 @@ class DNSServer(Service, identifier="DNSServer"):
# TCP for now
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.dns_table = self.config.domain_mapping
self.start()
def describe_state(self) -> Dict:

7
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -6,7 +6,6 @@ from pydantic import Field
from primaite import getLogger
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
from primaite.simulator.system.services.service import Service
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
@@ -22,14 +21,13 @@ class FTPServer(FTPServiceABC, identifier="FTPServer"):
"""
config: "FTPServer.ConfigSchema" = Field(default_factory=lambda: FTPServer.ConfigSchema())
server_password: Optional[str] = None
"""Password needed to connect to FTP server. Default is None."""
class ConfigSchema(Service.ConfigSchema):
class ConfigSchema(FTPServiceABC.ConfigSchema):
"""ConfigSchema for FTPServer."""
type: str = "FTPServer"
server_password: Optional[str] = None
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
@@ -37,6 +35,7 @@ class FTPServer(FTPServiceABC, identifier="FTPServer"):
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.start()
self.server_password = self.config.server_password
def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
"""

3
src/primaite/simulator/system/services/ntp/ntp_client.py
View File

@@ -9,6 +9,7 @@ from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
@@ -21,6 +22,7 @@ class NTPClient(Service, identifier="NTPClient"):
"""ConfigSchema for NTPClient."""
type: str = "NTPClient"
ntp_server_ip: Optional[IPV4Address] = None
config: "NTPClient.ConfigSchema" = Field(default_factory=lambda: NTPClient.ConfigSchema())
@@ -33,6 +35,7 @@ class NTPClient(Service, identifier="NTPClient"):
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)
self.ntp_server = self.config.ntp_server_ip
self.start()
def configure(self, ntp_server_ip_address: IPv4Address) -> None: