"IndexError: list index out of range" caused by _process_priv_esc

[Anth0rx]

I am trying to train and evaluate a PPO agent based on stable_baselines3 for Scenario2 running CybORG 2.1.

Notice: I use a reduced scenario file with only a red agent in a flat network configuration without Sleep and Impact actions for testing purposes:

Agents:
  Red:
    AllowedSubnets:
      - Flat
    INT:
      Hosts:
        User0:
          Interfaces: All
          System info: All
    actions:
    - DiscoverRemoteSystems
    - DiscoverNetworkServices
    - ExploitRemoteService
    - BlueKeep
    - EternalBlue
    - FTPDirectoryTraversal
    - HarakaRCE
    - HTTPRFI
    - HTTPSRFI
    - SQLInjection
    - PrivilegeEscalate
    - SSHBruteForce
    agent_type: SleepAgent
    reward_calculator_type: HybridImpactPwn
    starting_sessions:
    - hostname: User0
      name: RedPhish
      type: RedAbstractSession
      username: SYSTEM
    wrappers: []
Hosts:
  User0:
    AWS_Info: []
    image: windows_user_host1
    info:
      User0:
        Interfaces: All
      User1:
        Interfaces: IP Address
      User2:
        Interfaces: IP Address
      User3:
        Interfaces: IP Address
      User4:
        Interfaces: IP Address
      Enterprise0:
        Interfaces: IP Address
      Op_Server0:
        Interfaces: IP Address
    ConfidentialityValue: None
    AvailabilityValue: None
  User1:
    AWS_Info: []
    image: windows_user_host1
    info:
      User0:
        Interfaces: IP Address
      User1:
        Interfaces: All
      User2:
        Interfaces: IP Address
      User3:
        Interfaces: IP Address
      User4:
        Interfaces: IP Address
      Enterprise0:
        Interfaces: IP Address
      Op_Server0:
        Interfaces: IP Address
    ConfidentialityValue: Low
    AvailabilityValue: None
  User2:
    AWS_Info: []
    image: windows_user_host2
    info:
      User0:
        Interfaces: IP Address
      User1:
        Interfaces: IP Address
      User2:
        Interfaces: All
      User3:
        Interfaces: IP Address
      User4:
        Interfaces: IP Address
      Enterprise0:
        Interfaces: IP Address
      Op_Server0:
        Interfaces: IP Address
    ConfidentialityValue: Low
    AvailabilityValue: None
  User3:
    AWS_Info: []
    image: linux_user_host1
    info:
      User0:
        Interfaces: IP Address
      User1:
        Interfaces: IP Address
      User2:
        Interfaces: IP Address
      User3:
        Interfaces: All
      User4:
        Interfaces: IP Address
      Enterprise0:
        Interfaces: IP Address
      Op_Server0:
        Interfaces: IP Address
    ConfidentialityValue: Low
    AvailabilityValue: None
  User4:
    AWS_Info: []
    image: linux_user_host2
    info:
      User0:
        Interfaces: IP Address
      User1:
        Interfaces: IP Address
      User2:
        Interfaces: IP Address
      User3:
        Interfaces: IP Address
      User4:
        Interfaces: All
      Enterprise0:
        Interfaces: IP Address
      Op_Server0:
        Interfaces: IP Address
    ConfidentialityValue: Low
    AvailabilityValue: None
  Enterprise0:
    AWS_Info: []
    image: Gateway
    info:
      User0:
        Interfaces: IP Address
      User1:
        Interfaces: IP Address
      User2:
        Interfaces: IP Address
      User3:
        Interfaces: IP Address
      User4:
        Interfaces: IP Address
      Enterprise0:
        Interfaces: All
      Op_Server0:
        Interfaces: IP Address
    ConfidentialityValue: Medium
    AvailabilityValue: Medium
  Op_Server0:
    AWS_Info: []
    image: OP_Server
    info:
      User0:
        Interfaces: IP Address
      User1:
        Interfaces: IP Address
      User2:
        Interfaces: IP Address
      User3:
        Interfaces: IP Address
      User4:
        Interfaces: IP Address
      Enterprise0:
        Interfaces: IP Address
      Op_Server0:
        Interfaces: All
        Services:
        - OTService
    ConfidentialityValue: High
    AvailabilityValue: High
Subnets:
  Flat:
    Hosts:
    - User0
    - User1
    - User2
    - User3
    - User4
    - Enterprise0
    - Op_Server0
    NACLs:
      all:
        in: all
        out: all
    Size: 7

Currently, I am receiving the following error:

Traceback (most recent call last):
  File "/usr/lib64/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib64/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "[...]/cage-challenge-2/venv-3.10/lib64/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py", line 30, in _worker
    observation, reward, done, info = env.step(data)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/OpenAIGymWrapper.py", line 27, in step
    result = self.env.step(self.agent_name, action)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/EnumActionWrapper.py", line 20, in step
    return super().step(agent, action)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/BaseWrapper.py", line 16, in step
    result = self.env.step(agent, action)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/BaseWrapper.py", line 17, in step
    result.observation = self.observation_change(result.observation)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/RedTableWrapper.py", line 48, in observation_change
    self._update_red_info(observation)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/RedTableWrapper.py", line 85, in _update_red_info
    self._process_priv_esc(obs, hostname)
  File "[...]/cage-challenge-2/CybORG/CybORG/Agents/Wrappers/RedTableWrapper.py", line 132, in _process_priv_esc
    [info for info in self.red_info.values() if info[2] == hostname][0][4] = 'None'
IndexError: list index out of range

The problem seems to stem from the list comprehension inside the _process_priv_esc() function:

[info for info in self.red_info.values() if info[2] == hostname][0][4] = 'None'

After adding an additional check on the list length the agent model learns without errors:

def _process_priv_esc(self, obs, hostname):
        if obs['success'] == False:
            red_info = [info for info in self.red_info.values() if info[2] == hostname]
            if len(red_info) > 0:
                red_info[0][4] = 'None'
[...]