diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 9a5fedc9..d2752459 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1071,6 +1071,103 @@ class NodeNetworkServiceReconAction(AbstractAction): ] +class ConfigureC2BeaconAction(AbstractAction): + """Action which configures a C2 Beacon based on the parameters given.""" + + class _Opts(BaseModel): + """Schema for options that can be passed to this action.""" + + c2_server_ip_address: str + keep_alive_frequency: int = Field(default=5, ge=1) + masquerade_protocol: str = Field(default="TCP") + masquerade_port: str = Field(default="HTTP") + + @field_validator( + "c2_server_ip_address", + "keep_alive_frequency", + "masquerade_protocol", + "masquerade_port", + mode="before", + ) + @classmethod + def not_none(cls, v: str, info: ValidationInfo) -> int: + """If None is passed, use the default value instead.""" + if v is None: + return cls.model_fields[info.field_name].default + return v + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, config: Dict) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + config = ConfigureC2BeaconAction._Opts( + c2_server_ip_address=config["c2_server_ip_address"], + keep_alive_frequency=config["keep_alive_frequency"], + masquerade_port=config["masquerade_protocol"], + masquerade_protocol=config["masquerade_port"], + ) + + ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema + + return ["network", "node", node_name, "application", "C2Beacon", "configure", config.__dict__] + + +class RansomwareConfigureC2ServerAction(AbstractAction): + """Action which sends a command from the C2 Server to the C2 Beacon which configures a local RansomwareScript.""" + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, config: Dict) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + # Using the ransomware scripts model to validate. + ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema + return ["network", "node", node_name, "application", "C2Server", "ransomware_configure", config] + + +class RansomwareLaunchC2ServerAction(AbstractAction): + """Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript.""" + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + # Not options needed for this action. + return ["network", "node", node_name, "application", "C2Server", "ransomware_launch"] + + +class TerminalC2ServerAction(AbstractAction): + """Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed.""" + + class _Opts(BaseModel): + """Schema for options that can be passed to this action.""" + + model_config = ConfigDict(extra="forbid") + commands: RequestFormat + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, config: Dict) -> RequestFormat: + """Return the action formatted as a request that can be ingested by the simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None: + return ["do_nothing"] + TerminalC2ServerAction._Opts.model_validate(config) # check that options adhere to schema + return ["network", "node", node_name, "application", "C2Server", "terminal_command", config] + + class ActionManager: """Class which manages the action space for an agent.""" @@ -1122,6 +1219,10 @@ class ActionManager: "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, "CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction, "CONFIGURE_DOSBOT": ConfigureDoSBotAction, + "CONFIGURE_C2_BEACON": ConfigureC2BeaconAction, + "C2_SERVER_RANSOMWARE_LAUNCH": RansomwareLaunchC2ServerAction, + "C2_SERVER_RANSOMWARE_CONFIGURE": RansomwareConfigureC2ServerAction, + "C2_SERVER_TERMINAL_COMMAND": TerminalC2ServerAction, } """Dictionary which maps action type strings to the corresponding action class.""" diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 5ef8c14c..831dab2b 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -31,6 +31,8 @@ from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.sim_container import Simulation from primaite.simulator.system.applications.application import Application from primaite.simulator.system.applications.database_client import DatabaseClient # noqa: F401 +from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon # noqa: F401 +from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Server # noqa: F401 from primaite.simulator.system.applications.red_applications.data_manipulation_bot import ( # noqa: F401 DataManipulationBot, ) diff --git a/src/primaite/notebooks/Command-&-Control-E2E-Demonstration.ipynb b/src/primaite/notebooks/Command-&-Control-E2E-Demonstration.ipynb new file mode 100644 index 00000000..60ea756d --- /dev/null +++ b/src/primaite/notebooks/Command-&-Control-E2E-Demonstration.ipynb @@ -0,0 +1,367 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Command and Control Application Suite E2E Demonstration\n", + "\n", + "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "\n", + "This notebook demonstrates the current implementation of the command and control (C2) server and beacon applications in primAITE." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Imports\n", + "from primaite.config.load import data_manipulation_config_path\n", + "from primaite.session.environment import PrimaiteGymEnv\n", + "from primaite.game.agent.interface import AgentHistoryItem\n", + "import yaml\n", + "from pprint import pprint\n", + "from primaite.simulator.network.container import Network\n", + "from primaite.game.game import PrimaiteGame\n", + "from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon\n", + "from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Server\n", + "from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import C2Command, C2Payload\n", + "from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript\n", + "from primaite.simulator.network.hardware.nodes.host.computer import Computer\n", + "from primaite.simulator.network.hardware.nodes.host.server import Server" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## **Notebook Setup** | **Network Configuration:**\n", + "\n", + "This notebook uses the same network setup as UC2. Please refer to the main [UC2-E2E-Demo notebook for further reference](./Data-Manipulation-E2E-Demonstration.ipynb).\n", + "\n", + "However, this notebook will replaces with the red agent used in UC2 with a custom proxy red agent built for this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "custom_c2_agent = \"\"\"\n", + " - ref: CustomC2Agent\n", + " team: RED\n", + " type: ProxyAgent\n", + " observation_space: null\n", + " action_space:\n", + " action_list:\n", + " - type: DONOTHING\n", + " - type: NODE_APPLICATION_INSTALL\n", + " - type: NODE_APPLICATION_EXECUTE\n", + " - type: CONFIGURE_C2_BEACON\n", + " - type: C2_SERVER_RANSOMWARE_LAUNCH\n", + " - type: C2_SERVER_RANSOMWARE_CONFIGURE\n", + " - type: C2_SERVER_TERMINAL_COMMAND\n", + " options:\n", + " nodes:\n", + " - node_name: client_1\n", + " applications: \n", + " - application_name: C2Beacon\n", + " - node_name: domain_controller\n", + " applications: \n", + " - application_name: C2Server\n", + " max_folders_per_node: 1\n", + " max_files_per_folder: 1\n", + " max_services_per_node: 2\n", + " max_nics_per_node: 8\n", + " max_acl_rules: 10\n", + " ip_list:\n", + " - 192.168.1.10\n", + " - 192.168.1.14\n", + " action_map:\n", + " 0:\n", + " action: DONOTHING\n", + " options: {}\n", + " 1:\n", + " action: NODE_APPLICATION_INSTALL\n", + " options:\n", + " node_id: 0\n", + " application_name: C2Beacon\n", + " 2:\n", + " action: CONFIGURE_C2_BEACON\n", + " options:\n", + " node_id: 0\n", + " config:\n", + " c2_server_ip_address: 192.168.1.10\n", + " keep_alive_frequency:\n", + " masquerade_protocol:\n", + " masquerade_port:\n", + " 3:\n", + " action: NODE_APPLICATION_EXECUTE\n", + " options:\n", + " node_id: 0\n", + " application_id: 0 \n", + " 4:\n", + " action: NODE_APPLICATION_INSTALL\n", + " options:\n", + " node_id: 0\n", + " application_name: RansomwareScript \n", + " 5:\n", + " action: C2_SERVER_RANSOMWARE_CONFIGURE\n", + " options:\n", + " node_id: 1\n", + " config:\n", + " server_ip_address: 192.168.1.14\n", + " payload: ENCRYPT\n", + " 6:\n", + " action: C2_SERVER_RANSOMWARE_LAUNCH\n", + " options:\n", + " node_id: 1\n", + " 7:\n", + " action: C2_SERVER_TERMINAL_COMMAND\n", + " options:\n", + " node_id: 1\n", + " application_id: 0 \n", + "\n", + " reward_function:\n", + " reward_components:\n", + " - type: DUMMY\n", + "\"\"\"\n", + "c2_agent_yaml = yaml.safe_load(custom_c2_agent)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with open(data_manipulation_config_path()) as f:\n", + " cfg = yaml.safe_load(f)\n", + " # removing all agents & adding the custom agent.\n", + " cfg['agents'] = {}\n", + " cfg['agents'] = c2_agent_yaml\n", + " \n", + "\n", + "env = PrimaiteGymEnv(env_config=cfg)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## **Notebook Setup** | Network Prerequisites\n", + "\n", + "Before the Red Agent is able to perform any C2 specific actions, the C2 Server needs to be installed and run before the episode begins.\n", + "\n", + "This is because higher fidelity environments (and the real-world) a C2 server would not be accessible by private network blue agent and the C2 Server would already be in place before the an adversary (Red Agent) before the narrative of the use case.\n", + "\n", + "The cells below installs and runs the C2 Server on the domain controller server directly via the simulation API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "domain_controller: Server = env.game.simulation.network.get_node_by_hostname(\"domain_controller\")\n", + "domain_controller.software_manager.install(C2Server)\n", + "c2_server: C2Server = domain_controller.software_manager.software[\"C2Server\"]\n", + "c2_server.run()\n", + "domain_controller.software_manager.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## **Command and Control** | C2 Beacon Actions\n", + "\n", + "Before the Red Agent is able to perform any C2 Server commands, it must first establish connection with a C2 beacon.\n", + "\n", + "This can be done by installing, configuring and then executing a C2 Beacon. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "c2_red_agent = env.game.agents[\"CustomC2Agent\"]\n", + "client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### **Command and Control** | C2 Beacon Actions | Installation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.step(1)\n", + "client_1.software_manager.show()\n", + "c2_beacon: C2Beacon = client_1.software_manager.software[\"C2Beacon\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### **Command and Control** | C2 Beacon Actions | Configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.step(2)\n", + "c2_beacon.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### **Command and Control** | C2 Beacon Actions | Establishing Connection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.step(3)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "c2_beacon.show()\n", + "c2_server.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## **Command and Control** | C2 Server Actions" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### **Command and Control** | C2 Server Actions | Configuring Ransomware" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.step(4)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.step(5)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ransomware_script: RansomwareScript = client_1.software_manager.software[\"RansomwareScript\"]\n", + "client_1.software_manager.show()\n", + "ransomware_script.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### **Command and Control** | C2 Server Actions | Launching Ransomware" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.step(6)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "database_server: Server = env.game.simulation.network.get_node_by_hostname(\"database_server\")\n", + "database_server.software_manager.file_system.show(full=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### **Command and Control** | C2 Server Actions | Executing Terminal Commands" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# TODO: Post Terminal.\n", + "#env.step(7)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.8" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py b/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py index 16420164..c73799da 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py @@ -5,6 +5,7 @@ from typing import Dict, Optional # from primaite.simulator.system.services.terminal.terminal import Terminal from prettytable import MARKDOWN, PrettyTable +from pydantic import validate_call from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -17,7 +18,7 @@ from primaite.simulator.system.applications.red_applications.ransomware_script i from primaite.simulator.system.software import SoftwareHealthState -class C2Beacon(AbstractC2, identifier="C2 Beacon"): +class C2Beacon(AbstractC2, identifier="C2Beacon"): """ C2 Beacon Application. @@ -94,16 +95,16 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): ) c2_remote_ip = IPv4Address(c2_remote_ip) - # TODO: validation. frequency = request[-1].get("keep_alive_frequency") protocol = request[-1].get("masquerade_protocol") port = request[-1].get("masquerade_port") + return RequestResponse.from_bool( self.configure( c2_server_ip_address=c2_remote_ip, keep_alive_frequency=frequency, - masquerade_protocol=protocol, - masquerade_port=port, + masquerade_protocol=IPProtocol[protocol], + masquerade_port=Port[port], ) ) @@ -114,12 +115,13 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): kwargs["name"] = "C2Beacon" super().__init__(**kwargs) + @validate_call def configure( self, c2_server_ip_address: IPv4Address = None, - keep_alive_frequency: Optional[int] = 5, - masquerade_protocol: Optional[Enum] = IPProtocol.TCP, - masquerade_port: Optional[Enum] = Port.HTTP, + keep_alive_frequency: int = 5, + masquerade_protocol: Enum = IPProtocol.TCP, + masquerade_port: Enum = Port.HTTP, ) -> bool: """ Configures the C2 beacon to communicate with the C2 server with following additional parameters. @@ -278,8 +280,7 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): :return: Returns the Request Response returned by the Terminal execute method. :rtype: Request Response """ - # TODO: replace and use terminal - return RequestResponse(status="success", data={"Reason": "Placeholder."}) + return RequestResponse.from_bool(self._host_ransomware_script.attack()) def _command_terminal(self, payload: MasqueradePacket) -> RequestResponse: """ @@ -295,7 +296,6 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): """ # TODO: uncomment and replace (uses terminal) return RequestResponse(status="success", data={"Reason": "Placeholder."}) - # return self._host_terminal.execute(command) def _handle_keep_alive(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: """ @@ -421,10 +421,23 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): ``Keep Alive Frequency``: How often should the C2 Beacon attempt a keep alive? + ``Current Masquerade Protocol``: + The current protocol that the C2 Traffic is using. (e.g TCP/UDP) + + ``Current Masquerade Port``: + The current port that the C2 Traffic is using. (e.g HTTP (Port 80)) + :param markdown: If True, outputs the table in markdown format. Default is False. """ table = PrettyTable( - ["C2 Connection Active", "C2 Remote Connection", "Keep Alive Inactivity", "Keep Alive Frequency"] + [ + "C2 Connection Active", + "C2 Remote Connection", + "Keep Alive Inactivity", + "Keep Alive Frequency", + "Current Masquerade Protocol", + "Current Masquerade Port", + ] ) if markdown: table.set_style(MARKDOWN) @@ -436,6 +449,8 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): self.c2_remote_connection, self.keep_alive_inactivity, self.keep_alive_frequency, + self.current_masquerade_protocol, + self.current_masquerade_port, ] ) print(table) diff --git a/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py b/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py index d01cd412..c381403e 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py @@ -10,7 +10,7 @@ from primaite.simulator.network.protocols.masquerade import MasqueradePacket from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload -class C2Server(AbstractC2, identifier="C2 Server"): +class C2Server(AbstractC2, identifier="C2Server"): """ C2 Server Application. @@ -74,8 +74,8 @@ class C2Server(AbstractC2, identifier="C2 Server"): :rtype: RequestResponse """ # TODO: Parse the parameters from the request to get the parameters - placeholder: dict = {} - return self._send_command(given_command=C2Command.TERMINAL, command_options=placeholder) + terminal_commands = {"commands": request[-1].get("commands")} + return self._send_command(given_command=C2Command.TERMINAL, command_options=terminal_commands) rm.add_request( name="ransomware_configure", @@ -250,14 +250,29 @@ class C2Server(AbstractC2, identifier="C2 Server"): ``C2 Remote Connection``: The IP of the C2 Beacon. (Configured by upon receiving a keep alive.) + ``Current Masquerade Protocol``: + The current protocol that the C2 Traffic is using. (e.g TCP/UDP) + + ``Current Masquerade Port``: + The current port that the C2 Traffic is using. (e.g HTTP (Port 80)) + :param markdown: If True, outputs the table in markdown format. Default is False. """ - table = PrettyTable(["C2 Connection Active", "C2 Remote Connection"]) + table = PrettyTable( + ["C2 Connection Active", "C2 Remote Connection", "Current Masquerade Protocol", "Current Masquerade Port"] + ) if markdown: table.set_style(MARKDOWN) table.align = "l" table.title = f"{self.name} Running Status" - table.add_row([self.c2_connection_active, self.c2_remote_connection]) + table.add_row( + [ + self.c2_connection_active, + self.c2_remote_connection, + self.current_masquerade_protocol, + self.current_masquerade_port, + ] + ) print(table) # Abstract method inherited from abstract C2 - Not currently utilised.