#2689 Updated documentation and docustrings following PR comments.

This commit is contained in:
Archer Bowen
2024-08-15 17:47:33 +01:00
parent 7d086ec35e
commit e5be392ea8
7 changed files with 75 additions and 76 deletions

View File

@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field, field_validator, ValidationInfo
from primaite.interface.request import RequestFormat
class Command_Opts(BaseModel):
class CommandOpts(BaseModel):
"""A C2 Pydantic Schema acting as a base class for all C2 Commands."""
@field_validator("payload", "exfiltration_folder_name", "ip_address", mode="before", check_fields=False)
@@ -18,7 +18,7 @@ class Command_Opts(BaseModel):
return v
class Ransomware_Opts(Command_Opts):
class RansomwareOpts(CommandOpts):
"""A Pydantic Schema for the Ransomware Configuration command options."""
server_ip_address: str
@@ -28,7 +28,7 @@ class Ransomware_Opts(Command_Opts):
"""The malicious payload to be used to attack the target database."""
class Remote_Opts(Command_Opts):
class RemoteOpts(CommandOpts):
"""A base C2 Pydantic Schema for all C2 Commands that require a terminal connection."""
ip_address: Optional[str] = Field(default=None)
@@ -41,7 +41,7 @@ class Remote_Opts(Command_Opts):
"""A Password of a valid user account. Used to login into both remote and local hosts."""
class Exfil_Opts(Remote_Opts):
class ExfilOpts(RemoteOpts):
"""A Pydantic Schema for the C2 Data Exfiltration command options."""
target_ip_address: str
@@ -54,10 +54,10 @@ class Exfil_Opts(Remote_Opts):
"""The name of the remote folder which contains the target file."""
exfiltration_folder_name: str = Field(default="exfiltration_folder")
""""""
"""The name of C2 Suite folder used to store the target file. Defaults to ``exfiltration_folder``"""
class Terminal_Opts(Remote_Opts):
class TerminalOpts(RemoteOpts):
"""A Pydantic Schema for the C2 Terminal command options."""
commands: Union[list[RequestFormat], RequestFormat]

View File

@@ -277,8 +277,10 @@ class AbstractC2(Application, identifier="AbstractC2"):
"""Abstract Method: Used in C2 beacon to parse and handle commands received from the c2 server."""
pass
@abstractmethod
def _handle_keep_alive(self, payload: C2Packet, session_id: Optional[str]) -> bool:
"""Abstract Method: The C2 Server and the C2 Beacon handle the KEEP ALIVEs differently."""
"""Abstract Method: Each C2 suite handles ``C2Payload.KEEP_ALIVE`` differently."""
pass
# from_network_interface=from_network_interface
def receive(self, payload: any, session_id: Optional[str] = None, **kwargs) -> bool:

View File

@@ -11,7 +11,7 @@ from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.masquerade import C2Packet
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.red_applications.c2 import Exfil_Opts, Ransomware_Opts, Terminal_Opts
from primaite.simulator.system.applications.red_applications.c2 import ExfilOpts, RansomwareOpts, TerminalOpts
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
from primaite.simulator.system.services.terminal.terminal import Terminal, TerminalClientConnection
@@ -30,7 +30,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
Extends the Abstract C2 application to include the following:
1. Receiving commands from the C2 Server (Command input)
2. Leveraging the terminal application to execute requests (dependant on the command given)
2. Leveraging the terminal application to execute requests (dependent on the command given)
3. Sending the RequestResponse back to the C2 Server (Command output)
Please refer to the Command-&-Control notebook for an in-depth example of the C2 Suite.
@@ -156,7 +156,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:type c2_server_ip_address: IPv4Address
:param keep_alive_frequency: The frequency (timesteps) at which the C2 beacon will send keep alive(s).
:type keep_alive_frequency: Int
:param masquerade_protocol: The Protocol that C2 Traffic will masquerade as. Defaults as TCP.
:param masquerade_protocol: The Protocol that C2 Traffic will masquerade as. Defaults to TCP.
:type masquerade_protocol: Enum (IPProtocol)
:param masquerade_port: The Port that the C2 Traffic will masquerade as. Defaults to FTP.
:type masquerade_port: Enum (Port)
@@ -294,7 +294,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:return: Returns the Request Response returned by the Terminal execute method.
:rtype: Request Response
"""
command_opts = Ransomware_Opts.model_validate(payload.payload)
command_opts = RansomwareOpts.model_validate(payload.payload)
if self._host_ransomware_script is None:
return RequestResponse(
status="failure",
@@ -352,7 +352,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
data={"Reason": "Cannot find any instances of both a FTP Server & Client. Are they installed?"},
)
command_opts = Exfil_Opts.model_validate(payload.payload)
command_opts = ExfilOpts.model_validate(payload.payload)
# Setting up the terminal session and the ftp server
if not self._set_terminal_session(
@@ -401,7 +401,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
Attempts to exfiltrate a target file from a target using the parameters given.
Uses the current terminal_session to send a command to the
remote host's FTP Client passing the exfil_opts as command options.
remote host's FTP Client passing the ExfilOpts as command options.
This will instruct the FTP client to send the target file to the
dest_ip_address's destination folder.
@@ -411,8 +411,8 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
2. The target has a functioning FTP Client Service.
:exfil_opts: A Pydantic model containing the require configuration options
:type exfil_opts: Exfil_Opts
:ExfilOpts: A Pydantic model containing the require configuration options
:type ExfilOpts: ExfilOpts
:return: Returns a tuple containing a success boolean and a Request Response..
:rtype: tuple[bool, RequestResponse
"""
@@ -473,7 +473,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:return: Returns the Request Response returned by the Terminal execute method.
:rtype: Request Response
"""
command_opts = Terminal_Opts.model_validate(payload.payload)
command_opts = TerminalOpts.model_validate(payload.payload)
if self._host_terminal is None:
return RequestResponse(

View File

@@ -8,10 +8,10 @@ from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.masquerade import C2Packet
from primaite.simulator.system.applications.red_applications.c2 import (
Command_Opts,
Exfil_Opts,
Ransomware_Opts,
Terminal_Opts,
CommandOpts,
ExfilOpts,
RansomwareOpts,
TerminalOpts,
)
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
@@ -264,7 +264,7 @@ class C2Server(AbstractC2, identifier="C2Server"):
)
return self.current_command_output
def _command_setup(self, given_command: C2Command, command_options: dict) -> tuple[bool, Command_Opts]:
def _command_setup(self, given_command: C2Command, command_options: dict) -> tuple[bool, CommandOpts]:
"""
Performs any necessary C2 Server setup needed to perform certain commands.
@@ -289,13 +289,13 @@ class C2Server(AbstractC2, identifier="C2Server"):
:param command_options: The relevant command parameters.
:type command_options: Dict
:returns: Tuple containing a success bool if the setup was successful and the validated c2 opts.
:rtype: tuple[bool, Command_Opts]
:rtype: tuple[bool, CommandOpts]
"""
server_setup_success: bool = True
if given_command == C2Command.DATA_EXFILTRATION: # Data exfiltration setup
# Validating command options
command_options = Exfil_Opts.model_validate(command_options)
command_options = ExfilOpts.model_validate(command_options)
if self._host_ftp_server is None:
self.sys_log.warning(f"{self.name}: Unable to setup the FTP Server for data exfiltration")
server_setup_success = False
@@ -306,15 +306,15 @@ class C2Server(AbstractC2, identifier="C2Server"):
if given_command == C2Command.TERMINAL:
# Validating command options
command_options = Terminal_Opts.model_validate(command_options)
command_options = TerminalOpts.model_validate(command_options)
if given_command == C2Command.RANSOMWARE_CONFIGURE:
# Validating command options
command_options = Ransomware_Opts.model_validate(command_options)
command_options = RansomwareOpts.model_validate(command_options)
if given_command == C2Command.RANSOMWARE_LAUNCH:
# Validating command options
command_options = Command_Opts.model_validate(command_options)
command_options = CommandOpts.model_validate(command_options)
return [server_setup_success, command_options]