From 1b849947fadad387dbb0bb4181f8dca16ddf50ea Mon Sep 17 00:00:00 2001
From: Liz Rice <liz@lizrice.com>
Date: Fri, 8 Mar 2019 16:27:52 +0000
Subject: [PATCH] Use a predicate and avoid a whole extra event

---
 src/modules/discovery/apiserver.py     |  5 -----
 src/modules/hunting/apiserver.py       |  5 +++--
 tests/discovery/test_apiserver.py      | 17 ++++-------------
 tests/hunting/test_apiserver_hunter.py |  2 +-
 4 files changed, 8 insertions(+), 21 deletions(-)

diff --git a/src/modules/discovery/apiserver.py b/src/modules/discovery/apiserver.py
index eb3a715..658802e 100644
--- a/src/modules/discovery/apiserver.py
+++ b/src/modules/discovery/apiserver.py
@@ -29,13 +29,8 @@ class ApiServerDiscovery(Hunter):
 
     def execute(self):
         logging.debug("Attempting to discover an API server")
-
-        # We can discover the API Server with or without the use of a service account token
         main_request = requests.get("https://{}:{}".format(self.event.host, self.event.port), verify=False).text
         if '"code"' in main_request:
             self.event.role = "Master"
             self.publish_event(ApiServer())
 
-            # But if we have a service account token we will try additional checks
-            if self.event.auth_token:
-                self.publish_event(ApiServerWithServiceAccountToken())
diff --git a/src/modules/hunting/apiserver.py b/src/modules/hunting/apiserver.py
index 6950f96..ee85712 100644
--- a/src/modules/hunting/apiserver.py
+++ b/src/modules/hunting/apiserver.py
@@ -186,7 +186,8 @@ class ApiServerPassiveHunterFinished(Event):
         self.namespaces = namespaces
 
 
-# Passive Hunter
+# This Hunter checks what happens if we try to access the API Server without a service account token
+# If we have a service account token we'll also trigger AccessApiServerWithToken below
 @handler.subscribe(ApiServer)
 class AccessApiServer(Hunter):
     """ API Server Hunter
@@ -269,7 +270,7 @@ class AccessApiServer(Hunter):
         # the token
         self.publish_event(ApiServerPassiveHunterFinished(namespaces))
 
-@handler.subscribe(ApiServerWithServiceAccountToken)
+@handler.subscribe(ApiServer, predicate=lambda x: x.auth_token)
 class AccessApiServerWithToken(AccessApiServer):
     """ API Server Hunter
     Accessing the API server using the service account token obtained from a compromised pod
diff --git a/tests/discovery/test_apiserver.py b/tests/discovery/test_apiserver.py
index 4c46154..7652afc 100644
--- a/tests/discovery/test_apiserver.py
+++ b/tests/discovery/test_apiserver.py
@@ -1,7 +1,7 @@
 import requests_mock
 import time
 
-from src.modules.discovery.apiserver import ApiServer, ApiServerDiscovery, ApiServerWithServiceAccountToken
+from src.modules.discovery.apiserver import ApiServer, ApiServerDiscovery
 from src.core.events.types import Event
 from src.core.events import handler
 
@@ -46,19 +46,18 @@ def test_ApiServerWithServiceAccountToken():
         time.sleep(0.1)
         assert counter == 1
 
-        # If we have a token this should create two events
         e.auth_token = "very_secret"
         a = ApiServerDiscovery(e)
         a.execute()
         time.sleep(0.1)
-        assert counter == 3
+        assert counter == 2
 
         # But we shouldn't generate an event if we don't see an error code
         e.host = 'mockOther'
         a = ApiServerDiscovery(e)
         a.execute()
         time.sleep(0.1)
-        assert counter == 3
+        assert counter == 2
 
 
 # We should only generate an ApiServer event for a response that looks like it came from a Kubernetes node
@@ -67,12 +66,4 @@ class testApiServer(object):
     def __init__(self, event):
         assert event.host == 'mockKubernetes'
         global counter
-        counter += 1
-
-@handler.subscribe(ApiServerWithServiceAccountToken)
-class testApiServerWithServiceAccountToken(object):
-    def __init__(self, event):
-        assert event.host == 'mockKubernetes'
-        assert event.auth_token == "very_secret"
-        global counter 
-        counter += 1
+        counter += 1
\ No newline at end of file
diff --git a/tests/hunting/test_apiserver_hunter.py b/tests/hunting/test_apiserver_hunter.py
index a7d88f4..07bcdad 100644
--- a/tests/hunting/test_apiserver_hunter.py
+++ b/tests/hunting/test_apiserver_hunter.py
@@ -135,7 +135,7 @@ class test_ServerApiAccess(object):
             assert event.auth_token is None
         else:
             assert event.category == InformationDisclosure
-            assert event.auth_token is not None
+            assert event.auth_token == "so-secret"
         global counter
         counter += 1