From 8e19e05f570b57fee370190eb59384bda6adc77f Mon Sep 17 00:00:00 2001
From: Marek Wolan <marek.wolan@methods.co.uk>
Date: Sun, 21 Jan 2024 17:29:19 +0000
Subject: [PATCH] Fix acl actions for blue agent.

---
 .../config/_package_data/example_config.yaml  | 60 ++++++++++++-------
 src/primaite/game/agent/actions.py            | 15 ++++-
 2 files changed, 51 insertions(+), 24 deletions(-)

diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml
index ee0eb7ff..7393f5a3 100644
--- a/src/primaite/config/_package_data/example_config.yaml
+++ b/src/primaite/config/_package_data/example_config.yaml
@@ -304,63 +304,63 @@ agents:
             action: "NODE_RESET"
             options:
                 node_id: 5
-          22:
+          22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite)
             action: "NETWORK_ACL_ADDRULE"
             options:
                 position: 1
                 permission: 2
-                source_ip_id: 7
-                dest_ip_id: 1
+                source_ip_id: 7 # client 1
+                dest_ip_id: 1 # ALL
                 source_port_id: 1
                 dest_port_id: 1
                 protocol_id: 1
-          23:
+          23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite)
             action: "NETWORK_ACL_ADDRULE"
             options:
-                position: 1
+                position: 2
                 permission: 2
-                source_ip_id: 8
-                dest_ip_id: 1
+                source_ip_id: 8 # client 2
+                dest_ip_id: 1 # ALL
                 source_port_id: 1
                 dest_port_id: 1
                 protocol_id: 1
-          24:
+          24: # block tcp traffic from client 1 to web app
             action: "NETWORK_ACL_ADDRULE"
             options:
-                position: 1
+                position: 3
                 permission: 2
-                source_ip_id: 7
-                dest_ip_id: 3
+                source_ip_id: 7 # client 1
+                dest_ip_id: 3 # web server
                 source_port_id: 1
                 dest_port_id: 1
                 protocol_id: 3
-          25:
+          25: # block tcp traffic from client 2 to web app
             action: "NETWORK_ACL_ADDRULE"
             options:
-                position: 1
+                position: 4
                 permission: 2
-                source_ip_id: 8
-                dest_ip_id: 3
+                source_ip_id: 8 # client 2
+                dest_ip_id: 3 # web server
                 source_port_id: 1
                 dest_port_id: 1
                 protocol_id: 3
           26:
             action: "NETWORK_ACL_ADDRULE"
             options:
-                position: 1
+                position: 5
                 permission: 2
-                source_ip_id: 7
-                dest_ip_id: 4
+                source_ip_id: 7 # client 1
+                dest_ip_id: 4 # database
                 source_port_id: 1
                 dest_port_id: 1
                 protocol_id: 3
           27:
             action: "NETWORK_ACL_ADDRULE"
             options:
-                position: 1
+                position: 6
                 permission: 2
-                source_ip_id: 8
-                dest_ip_id: 4
+                source_ip_id: 8 # client 2
+                dest_ip_id: 4 # database
                 source_port_id: 1
                 dest_port_id: 1
                 protocol_id: 3
@@ -504,6 +504,24 @@ agents:
         max_services_per_node: 2
         max_nics_per_node: 8
         max_acl_rules: 10
+        ip_address_order:
+        - node_ref: domain_controller
+          nic_num: 1
+        - node_ref: web_server
+          nic_num: 1
+        - node_ref: database_server
+          nic_num: 1
+        - node_ref: backup_server
+          nic_num: 1
+        - node_ref: security_suite
+          nic_num: 1
+        - node_ref: client_1
+          nic_num: 1
+        - node_ref: client_2
+          nic_num: 1
+        - node_ref: security_suite
+          nic_num: 2
+
 
     reward_function:
       reward_components:
diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py
index 585e2dfa..6b15c5f8 100644
--- a/src/primaite/game/agent/actions.py
+++ b/src/primaite/game/agent/actions.py
@@ -470,13 +470,13 @@ class NetworkACLAddRuleAction(AbstractAction):
             dst_ip = "ALL"
             return ["do_nothing"]  # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS
         else:
-            dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id)
+            dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2)
             # subtract 2 to account for UNUSED=0, and ALL=1
 
         if dest_port_id == 1:
             dst_port = "ALL"
         else:
-            dst_port = self.manager.get_port_by_idx(dest_port_id)
+            dst_port = self.manager.get_port_by_idx(dest_port_id - 2)
             # subtract 2 to account for UNUSED=0, and ALL=1
 
         return [
@@ -924,6 +924,15 @@ class ActionManager:
         :return: The constructed ActionManager.
         :rtype: ActionManager
         """
+        ip_address_order = cfg["options"].pop("ip_address_order", {})
+        ip_address_list = []
+        for entry in ip_address_order:
+            node_ref = entry["node_ref"]
+            nic_num = entry["nic_num"]
+            node_obj = game.simulation.network.get_node_by_hostname(node_ref)
+            ip_address = node_obj.ethernet_port[nic_num].ip_address
+            ip_address_list.append(ip_address)
+
         obj = cls(
             game=game,
             actions=cfg["action_list"],
@@ -931,7 +940,7 @@ class ActionManager:
             **cfg["options"],
             protocols=game.options.protocols,
             ports=game.options.ports,
-            ip_address_list=None,
+            ip_address_list=ip_address_list or None,
             act_map=cfg.get("action_map"),
         )