adding cage-challenge-2

CybORG/Agents/SimpleAgents/B_line.py

@@ -1,3 +1,5 @@
+import random
+
 from CybORG.Agents import BaseAgent
 from CybORG.Shared import Results
 from CybORG.Shared.Actions import PrivilegeEscalate, ExploitRemoteService, DiscoverRemoteSystems, Impact, \
@@ -33,40 +35,45 @@ def get_action(self, observation, action_space):
 
             # Discover Remote Systems
             elif self.action == 0:
+                self.initial_ip = observation['User0']['Interface'][0]['IP Address']
                 self.last_subnet = observation['User0']['Interface'][0]['Subnet']
                 action = DiscoverRemoteSystems(session=session, agent='Red', subnet=self.last_subnet)
             # Discover Network Services- new IP address found
             elif self.action == 1:
-                self.last_ip_address = [value for key, value in observation.items() if key != 'success'][1]['Interface'][0]['IP Address']
+                hosts = [value for key, value in observation.items() if key != 'success']
+                get_ip = lambda x : x['Interface'][0]['IP Address']
+                interfaces = [get_ip(x) for x in hosts if get_ip(x)!= self.initial_ip]
+                self.last_ip_address = random.choice(interfaces)
                 action =DiscoverNetworkServices(session=session, agent='Red', ip_address=self.last_ip_address)
 
             # Exploit User1
             elif self.action == 2:
                  action = ExploitRemoteService(session=session, agent='Red', ip_address=self.last_ip_address)
 
-            # Privilege escalation on User1
+            # Privilege escalation on User Host
             elif self.action == 3:
                 hostname = [value for key, value in observation.items() if key != 'success' and 'System info' in value][0]['System info']['Hostname']
                 action = PrivilegeEscalate(agent='Red', hostname=hostname, session=session)
 
             # Discover Network Services- new IP address found
             elif self.action == 4:
-                self.last_ip_address = observation['Enterprise1']['Interface'][0]['IP Address']
+                self.enterprise_host = [x for x in observation if 'Enterprise' in x][0]
+                self.last_ip_address = observation[self.enterprise_host]['Interface'][0]['IP Address']
                 action = DiscoverNetworkServices(session=session, agent='Red', ip_address=self.last_ip_address)
 
-            # Exploit- Enterprise1
+            # Exploit- Enterprise Host
             elif self.action == 5:
                 self.target_ip_address = [value for key, value in observation.items() if key != 'success'][0]['Interface'][0]['IP Address']
                 action = ExploitRemoteService(session=session, agent='Red', ip_address=self.target_ip_address)
 
-            # Privilege escalation on Enterprise1
+            # Privilege escalation on Enterprise Host
             elif self.action == 6:
                 hostname = [value for key, value in observation.items() if key != 'success' and 'System info' in value][0]['System info']['Hostname']
                 action = PrivilegeEscalate(agent='Red', hostname=hostname, session=session)
 
             # Scanning the new subnet found.
             elif self.action == 7:
-                self.last_subnet = observation['Enterprise1']['Interface'][0]['Subnet']
+                self.last_subnet = observation[self.enterprise_host]['Interface'][0]['Subnet']
                 action = DiscoverRemoteSystems(subnet=self.last_subnet, agent='Red', session=session)
 
             # Discover Network Services- Enterprise2

CybORG/Agents/SimpleAgents/BlueLoadAgent.py

@@ -8,6 +8,7 @@
 from CybORG.Agents.Wrappers.FixedFlatWrapper import FixedFlatWrapper
 from CybORG.Agents.Wrappers.OpenAIGymWrapper import OpenAIGymWrapper
 from CybORG.Agents.Wrappers.ReduceActionSpaceWrapper import ReduceActionSpaceWrapper
+from CybORG.Agents.Wrappers import ChallengeWrapper
 
 class BlueLoadAgent(BaseAgent):
     # agent that loads a StableBaselines3 PPO model file
@@ -31,7 +32,7 @@ def get_action(self, observation, action_space):
         if self.model is None:
             path = str(inspect.getfile(CybORG))
             path = path[:-10] + '/Shared/Scenarios/Scenario1b.yaml'
-            cyborg = OpenAIGymWrapper('Blue', EnumActionWrapper(FixedFlatWrapper(ReduceActionSpaceWrapper(CybORG(path, 'sim')))))
+            cyborg = ChallengeWrapper(env=CybORG(path, 'sim'), agent_name='Blue')
             self.model = PPO('MlpPolicy', cyborg)
         action, _states = self.model.predict(observation)
         return action

CybORG/Agents/SimpleAgents/DebuggingAgent.py

@@ -3,8 +3,8 @@
 
 
 class DebuggingAgent(HeuristicRed):
-    def __init__(self, ip_list:list,session=0):
-        super().__init__(session=session)
+    def __init__(self, ip_list:list,session=0, priority=None):
+        super().__init__(session=session, priority=priority)
         self.ip_list = ip_list
         self.position = 0
         self.active_ip = self.ip_list[self.position]
@@ -16,5 +16,4 @@ def _choose_ip(self):
             self.position += 1 if self.position < len(self.ip_list) - 1 else 0
             ip = self.active_ip = self.ip_list[self.position]
 
-        assert ip in self.ip_status
         return ip

CybORG/Agents/SimpleAgents/HeuristicRed.py

@@ -2,7 +2,8 @@
 from CybORG.Shared.Actions import DiscoverRemoteSystems, DiscoverNetworkServices, ExploitRemoteService, PrivilegeEscalate, Impact
 
 class HeuristicRed():
-    def __init__(self, session=0):
+    def __init__(self, session=0, priority=None):
+        self.priority = priority
         self.parameters = {
                 'session':session,
                 'agent':'Red',
@@ -28,8 +29,9 @@ def get_action(self,obs):
         else:
             self._process_last_action_success() if self.last_action else None
             self._process_new_ips(obs)
-
+        
         action = self._advance_killchain()
+
         return action
 
     def _process_last_action_success(self):
@@ -88,8 +90,12 @@ def _advance_killchain(self):
             subnet = random.choice(list(self.unexplored_subnets))
             action = DiscoverRemoteSystems(subnet=subnet,**self.parameters)
         else:
+
             ip = self._choose_ip()
+
             action = self._choose_exploit(ip)
+            if ip not in self.ip_status:
+                self.ip_status[ip] = 0
 
         self.last_action = action
         self.history.append(action)
@@ -114,8 +120,10 @@ def _choose_ip(self):
     def _choose_exploit(self,ip):
         status = self.ip_status[ip]
         command = self.killchain[status]
-        if status < 2:
+        if status == 0:
             action = command(ip_address=ip,**self.parameters)
+        elif status == 1: 
+            action = command(ip_address=ip, priority=self.priority,**self.parameters)
         else:
             hostname = self.ip_map[ip]
             action = command(hostname=hostname,**self.parameters)

CybORG/Agents/SimpleAgents/KeyboardAgent.py

@@ -54,6 +54,8 @@ def _get_valid_commands(self,action_space):
             for parameter in parameter_list:
                 if parameter == 'self':
                     continue
+                if parameter == 'priority':
+                    continue
 
                 option_dict = action_space[parameter]
                 filter_f = lambda key : option_dict[key]

CybORG/Agents/SimpleAgents/Meander.py

@@ -1,3 +1,5 @@
+import random
+
 from CybORG.Agents.SimpleAgents.BaseAgent import BaseAgent
 from CybORG.Shared import Results
 from CybORG.Shared.Actions import PrivilegeEscalate, ExploitRemoteService, DiscoverRemoteSystems, Impact, \
@@ -21,41 +23,34 @@ def train(self, results: Results):
 
     def get_action(self, observation, action_space):
         """gets an action from the agent that should be performed based on the agent's internal state and provided observation and action space"""
-        if self.last_ip is not None:
-            if observation['success'] == True:
-                self.host_ip_map[[value['System info']['Hostname'] for key, value in observation.items()
-                                  if key != 'success' and 'System info' in value
-                                  and 'Hostname' in value['System info']][0]] = self.last_ip
-            else:
-                self.escalated_hosts = []
-            self.last_ip = None
-        if self.last_host is not None:
-            if observation['success'] == False:
-                if self.last_host in self.escalated_hosts:
-                    self.escalated_hosts.remove(self.last_host)
-                if self.last_host in self.host_ip_map and self.host_ip_map[self.last_host] in self.exploited_ips:
-                    self.exploited_ips.remove(self.host_ip_map[self.last_host])
-            self.last_host = None
+        self._process_success(observation)
 
         session = list(action_space['session'].keys())[0]
+
+        # Always impact if able
+        if 'Op_Server0' in self.escalated_hosts:
+            self.last_host = 'Op_Server0'
+            return Impact(agent='Red', hostname='Op_Server0', session=session)
+
         # start by scanning
         for subnet in action_space["subnet"]:
             if not action_space["subnet"][subnet] or subnet in self.scanned_subnets:
                 continue
             self.scanned_subnets.append(subnet)
             return DiscoverRemoteSystems(subnet=subnet, agent='Red', session=session)
         # discover network services
-        # act on ip addresses discovered in first subnet
+        # # act on ip addresses discovered in first subnet
         addresses = [i for i in action_space["ip_address"]]
+        random.shuffle(addresses)
         for address in addresses:
             if not action_space["ip_address"][address] or address in self.scanned_ips:
                 continue
             self.scanned_ips.append(address)
 
             return DiscoverNetworkServices(ip_address=address, agent='Red', session=session)
-
         # priv esc on owned hosts
-        hostnames = action_space['hostname']
+        hostnames = [x for x in action_space['hostname'].keys()]
+        random.shuffle(hostnames)
         for hostname in hostnames:
             # test if host is not known
             if not action_space["hostname"][hostname]:
@@ -79,8 +74,40 @@ def get_action(self, observation, action_space):
             self.last_ip = address
             return ExploitRemoteService(ip_address=address, agent='Red', session=session)
 
-        self.last_host = 'Op_Server0'
-        return Impact(agent='Red', hostname='Op_Server0', session=session)
+        raise NotImplementedError('Red Meander has run out of options!')
+
+
+
+    def _process_success(self, observation):
+        if self.last_ip is not None:
+            if observation['success'] == True:
+                self.host_ip_map[[value['System info']['Hostname'] for key, value in observation.items()
+                                  if key != 'success' and 'System info' in value
+                                  and 'Hostname' in value['System info']][0]] = self.last_ip
+            else:
+                self._process_failed_ip()
+            self.last_ip = None
+        if self.last_host is not None:
+            if observation['success'] == False:
+                if self.last_host in self.escalated_hosts:
+                    self.escalated_hosts.remove(self.last_host)
+                if self.last_host in self.host_ip_map and self.host_ip_map[self.last_host] in self.exploited_ips:
+                    self.exploited_ips.remove(self.host_ip_map[self.last_host])
+            self.last_host = None
+
+    def _process_failed_ip(self):
+        self.exploited_ips.remove(self.last_ip)
+        hosts_of_type = lambda y: [x for x in self.escalated_hosts if y in x]
+        if len(hosts_of_type('Op')) > 0:
+            for host in hosts_of_type('Op'):
+                self.escalated_hosts.remove(host)
+                ip = self.host_ip_map[host]
+                self.exploited_ips.remove(ip)
+        elif len(hosts_of_type('Ent')) > 0:
+            for host in hosts_of_type('Ent'):
+                self.escalated_hosts.remove(host)
+                ip = self.host_ip_map[host]
+                self.exploited_ips.remove(ip)
 
     def end_episode(self):
         self.scanned_subnets = []
@@ -92,4 +119,4 @@ def end_episode(self):
         self.last_ip = None
 
     def set_initial_values(self, action_space, observation):
-        pass
+        pass

CybORG/Agents/Wrappers/BlueTableWrapper.py

@@ -71,8 +71,7 @@ def observation_change(self,observation,baseline=False):
             raise NotImplementedError('Invalid output_mode for BlueTableWrapper')
 
     def _process_initial_obs(self, obs):
-        # TODO remove deepcopy replace with dict comprehension
-        obs = deepcopy(obs)
+        obs = obs.copy()
         self.baseline = obs
         del self.baseline['success']
         for hostid in obs:
@@ -157,19 +156,24 @@ def _process_anomalies(self,anomaly_dict):
 
     def _interpret_connections(self,activity:list):                
         num_connections = len(activity)
+
         ports = set([item['Connections'][0]['local_port'] \
             for item in activity if 'Connections' in item])
         port_focus = len(ports)
 
-        remote_ports = set([item['Connections'][0]['remote_port'] \
+        remote_ports = set([item['Connections'][0].get('remote_port') \
             for item in activity if 'Connections' in item])
+        if None in remote_ports:
+            remote_ports.remove(None)
 
         if num_connections >= 3 and port_focus >=3:
             anomaly = 'Scan'
         elif 4444 in remote_ports:
             anomaly = 'Exploit'
         elif num_connections >= 3 and port_focus == 1:
             anomaly = 'Exploit'
+        elif 'Service Name' in activity[0]:
+            anomaly = 'None'
         else:
             anomaly = 'Scan'
 

CybORG/Agents/Wrappers/ChallengeWrapper.py

@@ -58,3 +58,7 @@ def get_ip_map(self):
 
     def get_rewards(self):
         return self.get_attr('get_rewards')()
+
+    def get_reward_breakdown(self, agent: str):
+        return self.get_attr('get_reward_breakdown')(agent)
+

CybORG/Agents/Wrappers/EnumActionWrapper.py

@@ -33,6 +33,8 @@ def action_space_change(self, action_space: dict) -> int:
             param_dict = {}
             param_list = [{}]
             for p in self.action_signature[action]:
+                if p == 'priority':
+                    continue
                 temp[p] = []
                 if p not in params:
                     params.append(p)

CybORG/Agents/Wrappers/OpenAIGymWrapper.py

@@ -17,7 +17,7 @@ def __init__(self, agent_name: str, env: BaseWrapper = None, agent: BaseAgent =
             assert isinstance(self.get_action_space(self.agent_name), int)
             self.action_space = spaces.Discrete(self.get_action_space(self.agent_name))
         box_len = len(self.observation_change(self.env.reset(self.agent_name).observation))
-        self.observation_space = spaces.Box(-1.0, 3.0, shape=(box_len,), dtype=np.float32)
+        self.observation_space = spaces.Box(-1.0, 1.0, shape=(box_len,), dtype=np.float32)
         self.reward_range = (float('-inf'), float('inf'))
         self.metadata = {}
         self.action = None
@@ -28,13 +28,13 @@ def step(self, action: Union[int, List[int]] = None) -> (object, float, bool, di
         result.observation = self.observation_change(result.observation)
         result.action_space = self.action_space_change(result.action_space)
         info = vars(result)
-        return np.array(result.observation, dtype=np.float32), result.reward, result.done, info
+        return np.array(result.observation), result.reward, result.done, info
 
     def reset(self, agent=None):
         result = self.env.reset(self.agent_name)
         result.action_space = self.action_space_change(result.action_space)
         result.observation = self.observation_change(result.observation)
-        return np.array(result.observation, dtype=np.float32)
+        return np.array(result.observation)
 
     def render(self):
         # TODO: If FixedFlatWrapper it will error out!