From 482a7479e06c607c583acedacab68c7817a609b8 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Wed, 27 Jul 2022 15:03:29 +0200
Subject: [PATCH] BB: Add starting GCP machines of single tests

---
 envs/monkey_zoo/blackbox/conftest.py          |  5 ++
 .../blackbox/gcp_test_machine_list.py         | 67 +++++++++++++++++++
 envs/monkey_zoo/blackbox/test_blackbox.py     | 14 ++--
 3 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/envs/monkey_zoo/blackbox/conftest.py b/envs/monkey_zoo/blackbox/conftest.py
index a73e6fe8b..d844bc403 100644
--- a/envs/monkey_zoo/blackbox/conftest.py
+++ b/envs/monkey_zoo/blackbox/conftest.py
@@ -32,6 +32,11 @@ def no_gcp(request):
     return request.config.getoption("--no-gcp")
 
 
+@pytest.fixture(scope="session")
+def machines_to_start(request):
+    return request.config.getoption("-k")
+
+
 def pytest_runtest_setup(item):
     if "skip_powershell_reuse" in item.keywords and item.config.getoption(
         "--skip-powershell-reuse"
diff --git a/envs/monkey_zoo/blackbox/gcp_test_machine_list.py b/envs/monkey_zoo/blackbox/gcp_test_machine_list.py
index 866e69c3e..9aa5ef41f 100644
--- a/envs/monkey_zoo/blackbox/gcp_test_machine_list.py
+++ b/envs/monkey_zoo/blackbox/gcp_test_machine_list.py
@@ -26,3 +26,70 @@ GCP_TEST_MACHINE_LIST = {
         "log4j-tomcat-52",
     ],
 }
+
+DEPTH_2_A = {
+    "europe-west3-a": [
+        "sshkeys-11",
+        "sshkeys-12",
+    ]
+}
+
+
+DEPTH_1_A = {
+    "europe-west3-a": ["hadoop-2", "hadoop-3", "mssql-16", "mimikatz-14", "mimikatz-15"],
+    "europe-west1-b": [
+        "log4j-logstash-55",
+        "log4j-logstash-56",
+        "log4j-solr-49",
+        "log4j-solr-50",
+        "log4j-tomcat-51",
+        "log4j-tomcat-52",
+    ],
+}
+
+DEPTH_3_A = {
+    "europe-west3-a": [
+        "tunneling-9",
+        "tunneling-10",
+        "tunneling-11",
+        "tunneling-12",
+        "mimikatz-15",
+    ],
+    "europe-west1-b": [
+        "powershell-3-45",
+        "powershell-3-46",
+        "powershell-3-47",
+        "powershell-3-48",
+    ],
+}
+
+POWERSHELL_EXPLOITER_REUSE = {
+    "europe-west1-b": [
+        "powershell-3-46",
+    ]
+}
+
+ZEROLOGON = {
+    "europe-west3-a": [
+        "zerologon-25",
+    ],
+}
+
+WMI_AND_MIMIKATZ = {
+    "europe-west3-a": [
+        "mimikatz-14",
+        "mimikatz-15",
+    ]
+}
+
+SMB_PTH = {"europe-west3-a": ["mimikatz-15"]}
+
+GCP_SINGLE_TEST_LIST = {
+    "test_depth_2_a": DEPTH_2_A,
+    "test_depth_1_a": DEPTH_1_A,
+    "test_depth_3_a": DEPTH_3_A,
+    "test_powershell_exploiter_credentials_reuse": POWERSHELL_EXPLOITER_REUSE,
+    "test_zerologon_exploiter": ZEROLOGON,
+    "test_wmi_and_mimikatz_exploiters": WMI_AND_MIMIKATZ,
+    "test_smb_pth": SMB_PTH,
+}
diff --git a/envs/monkey_zoo/blackbox/test_blackbox.py b/envs/monkey_zoo/blackbox/test_blackbox.py
index 961d094fb..f9fa72c9b 100644
--- a/envs/monkey_zoo/blackbox/test_blackbox.py
+++ b/envs/monkey_zoo/blackbox/test_blackbox.py
@@ -6,7 +6,10 @@ import pytest
 
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
 from envs.monkey_zoo.blackbox.analyzers.zerologon_analyzer import ZerologonAnalyzer
-from envs.monkey_zoo.blackbox.gcp_test_machine_list import GCP_TEST_MACHINE_LIST
+from envs.monkey_zoo.blackbox.gcp_test_machine_list import (
+    GCP_SINGLE_TEST_LIST,
+    GCP_TEST_MACHINE_LIST,
+)
 from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
 from envs.monkey_zoo.blackbox.island_client.test_configuration_parser import get_target_ips
 from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
@@ -35,18 +38,21 @@ LOGGER = logging.getLogger(__name__)
 
 
 @pytest.fixture(autouse=True, scope="session")
-def GCPHandler(request, no_gcp):
+def GCPHandler(request, no_gcp, machines_to_start):
     if not no_gcp:
+        list_machines = GCP_TEST_MACHINE_LIST
+        if machines_to_start:
+            list_machines = GCP_SINGLE_TEST_LIST[machines_to_start]
         try:
             initialize_gcp_client()
-            start_machines(GCP_TEST_MACHINE_LIST)
+            start_machines(list_machines)
         except Exception as e:
             LOGGER.error("GCP Handler failed to initialize: %s." % e)
             pytest.exit("Encountered an error while starting GCP machines. Stopping the tests.")
         wait_machine_bootup()
 
         def fin():
-            stop_machines(GCP_TEST_MACHINE_LIST)
+            stop_machines(list_machines)
 
         request.addfinalizer(fin)