Feat/tests fixes improvments

.github/workflows/python-app.yml

@@ -10,27 +10,42 @@ permissions:
   contents: read
 
 jobs:
-  lint-and-format:
+  lint:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
-
     - name: Set up Python 3.12
       uses: actions/setup-python@v5
       with:
         python-version: "3.12"
-
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install ruff
-
     - name: Check formatting with ruff
-      run: |
-        # Fails the build if any files need to be formatted
-        ruff format --check .
-
+      run: ruff format --check .
     - name: Lint with ruff
+      run: ruff check .
+
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.12
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
+    - name: Install dependencies
       run: |
-        # Check for unused imports and undefined variables
-        ruff check .
+        python -m pip install --upgrade pip
+        pip install -e .
+        pip install pytest pytest-cov scikit-learn gymnasium pettingzoo networkx pyyaml
+    - name: Run tests with pytest
+      run: |
+        pytest tests/ -v --cov=netforge_rl --cov-report=xml --cov-fail-under=70
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v4
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        fail_ci_if_error: false
+

netforge_rl/actions/blue/identity.py

@@ -4,7 +4,7 @@
 from netforge_rl.core.registry import action_registry
 
 
-@action_registry.register('RotateKerberos', 'blue')
+@action_registry.register('blue_commander', 0)
 class RotateKerberos(BaseAction):
     """
     Apex Zero-Trust Action: Rotates Domain Kerberos TGT Keys globally.

netforge_rl/actions/blue/mitigation.py

@@ -7,7 +7,7 @@
 )
 
 
-@action_registry.register('blue_operator', 0)
+@action_registry.register('blue', 0)
 class IsolateHost(BaseAction):
     """Disconnects a compromised host completely from the network
 
@@ -56,7 +56,7 @@ def execute(self, global_state) -> ActionEffect:
         )
 
 
-@action_registry.register('blue_operator', 1)
+@action_registry.register('blue', 1)
 class RestoreHost(BaseAction):
     """Re-establishes network connectivity for a previously isolated host.
 
@@ -102,7 +102,7 @@ def execute(self, global_state) -> ActionEffect:
         )
 
 
-@action_registry.register('blue_operator', 4)
+@action_registry.register('blue', 4)
 class Remove(BaseAction):
     """Evicts unauthorized threat actors from a compromised element.
 
@@ -149,7 +149,7 @@ def execute(self, global_state) -> ActionEffect:
         )
 
 
-@action_registry.register('blue_operator', 5)
+@action_registry.register('blue', 5)
 class RestoreFromBackup(BaseAction):
     """Executes a bare-metal imaging recovery to purge advanced persistent
 
@@ -199,7 +199,7 @@ def execute(self, global_state) -> ActionEffect:
         )
 
 
-@action_registry.register('blue_operator', 6)
+@action_registry.register('blue', 6)
 class ConfigureACL(BaseAction):
     """
     Dynamically modifies the implicit routing Firewall to block specific port
@@ -211,7 +211,7 @@ class ConfigureACL(BaseAction):
         port (int): The destination port to drop (e.g., 445).
     """
 
-    def __init__(self, agent_id: str, target_subnet: str, port: int):
+    def __init__(self, agent_id: str, target_subnet: str, port: int = 445):
         super().__init__(agent_id, target_ip=target_subnet, cost=2)
         self.port = port
 
@@ -234,7 +234,7 @@ def execute(self, global_state) -> ActionEffect:
         )
 
 
-@action_registry.register('blue_operator', 7)
+@action_registry.register('blue', 7)
 class SecurityAwarenessTraining(BaseAction):
     """
     Deploys rapid, intensive anti-phishing training to a targeted subnet.

netforge_rl/actions/red/exploits.py

@@ -6,7 +6,7 @@
 )
 
 
-@action_registry.register('red_operator', 0)
+@action_registry.register('red', 0)
 class ExploitRemoteService(BaseAction):
     """Attempts to weaponize a generic remote code execution vulnerability on a
 
@@ -108,7 +108,7 @@ def execute(self, global_state) -> ActionEffect:
         )
 
 
-@action_registry.register('red_operator', 3)
+@action_registry.register('red', 3)
 class ExploitBlueKeep(BaseAction):
     """Executes the CVE-2019-0708 (BlueKeep) vulnerability against Remote
 
@@ -198,14 +198,15 @@ def execute(self, global_state) -> ActionEffect:
                 f'hosts/{self.target_ip}/compromised_by': self.agent_id,
             },
             observation_data={
-                'exploit': 'BlueKeep success',
+                'exploit': self.target_ip,
+                'status': 'BlueKeep success',
                 'sim2real_stdout': hw_result.stdout if hw_result else None,
                 'sim2real_reward_delta': reward_delta,
             },
         )
 
 
-@action_registry.register('red_operator', 4)
+@action_registry.register('red', 4)
 class ExploitEternalBlue(BaseAction):
     """Executes the MS17-010 (EternalBlue) exploit targeting poorly configured
 
@@ -285,11 +286,14 @@ def execute(self, global_state) -> ActionEffect:
                 f'hosts/{self.target_ip}/privilege': 'User',
                 f'hosts/{self.target_ip}/compromised_by': self.agent_id,
             },
-            observation_data={'exploit': 'EternalBlue success'},
+            observation_data={
+                'exploit': self.target_ip,
+                'status': 'EternalBlue success',
+            },
         )
 
 
-@action_registry.register('red_operator', 5)
+@action_registry.register('red', 5)
 class ExploitHTTP_RFI(BaseAction):
     """Simulates a Remote File Inclusion (RFI) web application attack vector
 

netforge_rl/actions/red/post_exploitation.py

@@ -2,7 +2,7 @@
 from netforge_rl.core.registry import action_registry
 
 
-@action_registry.register('DumpLSASS', 'red')
+@action_registry.register('red', 7)
 class DumpLSASS(BaseAction):
     """
     Advanced Post-Exploitation Action: Scrapes memory for Active Directory tokens.
@@ -63,7 +63,7 @@ def execute(self, state):
         )
 
 
-@action_registry.register('PassTheTicket', 'red')
+@action_registry.register('red', 8)
 class PassTheTicket(BaseAction):
     """
     Lateral Movement via Identity validation bypassing CVE exploits explicitly.

netforge_rl/actions/red/privilege_escalation.py

@@ -2,7 +2,7 @@
 from netforge_rl.core.registry import action_registry
 
 
-@action_registry.register('red_operator', 1)
+@action_registry.register('red', 1)
 class PrivilegeEscalate(BaseAction):
     """Executes a generic local privilege escalation exploit on a compromised
 

netforge_rl/actions/red/reconnaissance.py

@@ -112,6 +112,7 @@ def execute(self, global_state) -> ActionEffect:
 
 
 @action_registry.register('red_commander', 2)
+@action_registry.register('red', 2)
 class DiscoverNetworkServices(BaseAction):
     """Executes an intrusive port scan against a specific host to enumerate
 

netforge_rl/core/observation.py

@@ -66,9 +66,12 @@ def update_from_state(self, global_state: Any, action_effects: List[Any]):
             # Pull SIEM logs that have arrived (arrival_tick <= current_tick)
             if hasattr(global_state, 'siem_log_buffer'):
                 for log in global_state.siem_log_buffer:
-                    if log.get('arrival_tick', 0) <= getattr(
-                        global_state, 'current_tick', 0
-                    ):
+                    # Logs can be raw strings or telemetry dictionaries.
+                    arrival_tick = 0
+                    if isinstance(log, dict):
+                        arrival_tick = log.get('arrival_tick', 0)
+
+                    if arrival_tick <= getattr(global_state, 'current_tick', 0):
                         self.siem_alerts.append(log)
 
             self.network_telemetry['global_alert_level'] = np.random.uniform(0, 1)

netforge_rl/environment/parallel_env.py

@@ -273,10 +273,17 @@ def step(
             self.global_state.siem_log_buffer.append(anomaly)
 
         # 4. RESOLVE MATURE EVENTS
+        intended_effects = {}
+        action_metadata = {}
         remaining_events = []
         for event in self.event_queue:
             if self.current_tick >= event['completion_tick']:
-                intended_effects[event['agent']] = event['effect']
+                agent = event['agent']
+                intended_effects[agent] = event['effect']
+                action_metadata[agent] = {
+                    'name': type(event['action']).__name__,
+                    'target_ip': event.get('target_ip'),
+                }
             else:
                 remaining_events.append(event)
         self.event_queue = remaining_events
@@ -287,28 +294,11 @@ def step(
 
         # NLP-SIEM: generate structured event logs from resolved action effects
         for res_agent, res_effect in resolved_effects.items():
-            action_name = type(
-                next(
-                    (
-                        e['action']
-                        for e in self.event_queue
-                        if e.get('agent') == res_agent
-                    ),
-                    None,
-                )
-                or type('', (), {})()
-            ).__name__
-            # Prefer fetching name from the event that just resolved
-            for ev in list(self.event_queue) + [
-                e
-                for e in [
-                    {'agent': k, 'action': type('_A', (), {'__name__': 'Unknown'})()}
-                    for k in resolved_effects
-                ]
-            ]:
-                if ev.get('agent') == res_agent:
-                    action_name = type(ev.get('action', object())).__name__
-                    break
+            meta = action_metadata.get(res_agent, {})
+            action_name = meta.get('name', 'UnknownAction')
+            target_ip = meta.get('target_ip') or res_effect.observation_data.get(
+                'exploit'
+            )
             self.siem_logger.log_action(
                 action_name=action_name,
                 effect=res_effect,
@@ -458,7 +448,11 @@ def _extract_agent_infos(self, observations: dict, resolved_effects: dict) -> di
             hosts_isolated = 0
             services_restored = 0
 
-            if agent_effect and agent_effect.success:
+            if (
+                agent_effect
+                and agent_effect.success
+                and isinstance(agent_effect.state_deltas, dict)
+            ):
                 for delta_key, delta_val in agent_effect.state_deltas.items():
                     if 'status' in delta_key and delta_val == 'isolated':
                         hosts_isolated += 1

netforge_rl/nlp/log_encoder.py

@@ -18,6 +18,7 @@
 import hashlib
 import json
 import logging
+import random
 from pathlib import Path
 from typing import Literal
 
@@ -128,6 +129,11 @@ def _build_tfidf(self):
 
         def encode_fn(text: str) -> np.ndarray:
             vec = pipeline.transform([text])[0]
+            # Ensure fixed output dimension even if SVD capped out
+            if vec.shape[0] < EMBEDDING_DIM:
+                padded = np.zeros(EMBEDDING_DIM, dtype=np.float32)
+                padded[: vec.shape[0]] = vec
+                return padded
             return vec.astype(np.float32)
 
         return encode_fn
@@ -208,14 +214,28 @@ def _build_training_corpus(self) -> list[str]:
         for src, tgt in zip(sample_ips, reversed(sample_ips)):
             for fn in [evid_4624, evid_4625, evid_4648, evid_4776]:
                 corpus.append(fn(src, tgt))
-            corpus.append(evid_4688(src, process='mimikatz.exe'))
-            corpus.append(evid_4688(src, process='powershell.exe'))
+            # Add more variations to ensure > 128 samples
+            for proc in [
+                'cmd.exe',
+                'powershell.exe',
+                'mimikatz.exe',
+                'procdump.exe',
+                'net.exe',
+            ]:
+                corpus.append(evid_4688(src, process=proc))
+                corpus.append(sysmon_1(src, process=proc))
             corpus.append(evid_4768(src, tgt))
-            corpus.append(sysmon_1(src, process='powershell.exe'))
             corpus.append(sysmon_3(src, tgt, dst_port=445))
+            corpus.append(sysmon_3(src, tgt, dst_port=3389))
             corpus.append(sysmon_10(src))
             corpus.append(sysmon_22(src))
 
+        # Add 50 unique random noise strings to guarantee diversity
+        for i in range(50):
+            corpus.append(
+                f'Synthetic noise event {i} for dimension stability - {random.random()}'
+            )
+
         if not corpus:
             # Ultimate fallback — at least something to fit on
             corpus = [

netforge_rl/scenarios/ransomware.py

@@ -55,7 +55,12 @@ def _red_reward(
 
         # ── ONE-TIME action bonuses (only on success) ─────────
         if effect.success and effect.state_deltas:
-            for delta_key, delta_val in effect.state_deltas.items():
+            deltas = (
+                effect.state_deltas.items()
+                if isinstance(effect.state_deltas, dict)
+                else []
+            )
+            for delta_key, delta_val in deltas:
                 # Initial compromise (None → User)
                 if 'privilege' in delta_key and delta_val == 'User':
                     reward += 3.0
@@ -118,9 +123,19 @@ def _blue_reward(
     ) -> float:
         reward = 0.0
 
-        # ── ONE-TIME action bonuses ───────────────────────────
+        # ONE-TIME action bonuses
         if effect and effect.success and effect.state_deltas:
-            for delta_key, delta_val in effect.state_deltas.items():
+            # We iterate differently based on whether it's a dict or a list
+            deltas = (
+                effect.state_deltas.items()
+                if isinstance(effect.state_deltas, dict)
+                else []
+            )
+
+            # If it's a list (e.g. IdentityFlush), we don't have key/val pairs easily
+            # but we can look for specific attributes if needed.
+            # For now, we only reward dict-based state changes which are common for most actions.
+            for delta_key, delta_val in deltas:
                 # Successful isolation
                 if 'status' in delta_key and delta_val == 'isolated':
                     ip = delta_key.split('/')[1] if '/' in delta_key else None