Merge remote-tracking branch 'origin/dev' into feature/2417-observation-refactor

tests/e2e_integration_tests/test_primaite_session.py

@@ -1,6 +1,8 @@
 import pydantic
 import pytest
 
+from primaite.simulator.network.hardware.nodes.host.server import Printer
+from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
 from tests import TEST_ASSETS_ROOT
 from tests.conftest import TempPrimaiteSession
 
@@ -30,7 +32,11 @@ class TestPrimaiteSession:
             assert session.env
 
             assert session.env.game.simulation.network
-            assert len(session.env.game.simulation.network.nodes) == 10
+            assert len(session.env.game.simulation.network.nodes) == 12
+            wireless = session.env.game.simulation.network.get_node_by_hostname("router_2")
+            assert isinstance(wireless, WirelessRouter)
+            printer = session.env.game.simulation.network.get_node_by_hostname("HP_LaserJet_Pro_4102fdn_printer")
+            assert isinstance(printer, Printer)
 
     @pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
     @pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)

tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py

@@ -1,8 +1,13 @@
+import yaml
+
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteGymEnv
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.host.server import Server
 from primaite.simulator.system.applications.database_client import DatabaseClient
 from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
 from primaite.simulator.system.services.database.database_service import DatabaseService
+from tests import TEST_ASSETS_ROOT
 
 
 def test_data_manipulation(uc2_network):
@@ -32,3 +37,41 @@ def test_data_manipulation(uc2_network):
 
     # Now check that the DB client on the web_server can successfully query the users table on the database
     assert db_client.query("SELECT")
+
+
+def test_application_install_uninstall_on_uc2():
+    """Test Application install and uninstall via agent actions mid episode."""
+    with open(TEST_ASSETS_ROOT / "configs/test_application_install.yaml", "r") as f:
+        cfg = yaml.safe_load(f)
+
+    env = PrimaiteGymEnv(game_config=cfg)
+    env.agent.flatten_obs = False
+    env.reset()
+
+    _, _, _, _, _ = env.step(0)
+    domcon = env.game.simulation.network.get_node_by_hostname("domain_controller")
+
+    # Test we cannot execute the DoSBot app as it is not installed yet
+    _, _, _, _, info = env.step(81)
+    assert info["agent_actions"]["defender"].response.status == "unreachable"
+
+    # Test we can Install the DoSBot app
+    _, _, _, _, info = env.step(78)
+    assert "DoSBot" in domcon.software_manager.software
+
+    # Test we can now execute the DoSBot app
+    _, _, _, _, info = env.step(81)
+    assert info["agent_actions"]["defender"].response.status == "success"
+
+    # Test we can Uninstall the DoSBot app
+    _, _, _, _, info = env.step(79)
+    assert "DoSBot" not in domcon.software_manager.software
+
+    # Test we cannot execute the DoSBot app as it was uninstalled
+    _, _, _, _, info = env.step(81)
+    assert info["agent_actions"]["defender"].response.status == "unreachable"
+
+    # Test we can uninstall one of the default apps (WebBrowser)
+    assert "WebBrowser" in domcon.software_manager.software
+    _, _, _, _, info = env.step(80)
+    assert "WebBrowser" not in domcon.software_manager.software