From 1cd44832e60ba94bd02bd5225c93c5c87c14a735 Mon Sep 17 00:00:00 2001
From: Michael Cherny <cherny.business@gmail.com>
Date: Thu, 7 Mar 2019 14:45:26 +0200
Subject: [PATCH] Fixes #99 - pod local vulnerabilities are now reported as
 "Local to Pod" ( <pod name> )

Event  can now implement 'location()' method that return string representing  events logical location.
In events chain, the 'newest' event available location method will be used. This is because we compose (chain) events.
Core changed to support it.
Added 'location()' method to relevant event classes.
Reports are now using vulnerability.location() to retrieve location.
---
 src/core/events/types/common.py | 25 +++++++++++++++++++++++--
 src/modules/discovery/hosts.py  |  8 ++++++++
 src/modules/report/base.py      |  2 +-
 src/modules/report/plain.py     |  2 +-
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/src/core/events/types/common.py b/src/core/events/types/common.py
index ffe60a4..2ec8acf 100644
--- a/src/core/events/types/common.py
+++ b/src/core/events/types/common.py
@@ -16,6 +16,17 @@ class Event(object):
             if name in event.__dict__:
                 return event.__dict__[name]
 
+    # Event's logical location to be used mainly for reports.
+    # If event don't implement it check previous event
+    # This is because events are composed (previous -> previous ...)
+    # and not inheritted 
+    def location(self):
+        location = None
+        if self.previous:
+            location = self.previous.location()
+
+        return location
+
     # returns the event history ordered from newest to oldest
     @property
     def history(self):
@@ -85,7 +96,10 @@ class NewHostEvent(Event):
 
     def __str__(self):
         return str(self.host)
-
+    
+    # Event's logical location to be used mainly for reports.
+    def location(self):
+        return str(self.host)
 
 class OpenPortEvent(Event):
     def __init__(self, port):
@@ -93,7 +107,14 @@ class OpenPortEvent(Event):
     
     def __str__(self):
         return str(self.port)
-
+    
+    # Event's logical location to be used mainly for reports.
+    def location(self):
+        if self.host:
+            location = str(self.host) + ":" + str(self.port)
+        else:
+            location = str(self.port)
+        return location
 
 class HuntFinished(Event):
     pass
diff --git a/src/modules/discovery/hosts.py b/src/modules/discovery/hosts.py
index b942ce3..afce667 100644
--- a/src/modules/discovery/hosts.py
+++ b/src/modules/discovery/hosts.py
@@ -22,6 +22,14 @@ class RunningAsPodEvent(Event):
         self.auth_token = self.get_auth_token()
         self.client_cert = self.get_client_cert()
         
+    # Event's logical location to be used mainly for reports.
+    def location(self):
+        location = "Local to Pod"
+        if 'HOSTNAME' in os.environ:
+            location += "(" + os.environ['HOSTNAME'] + ")"
+        
+        return location
+
     def get_auth_token(self):
         try:
             with open("/var/run/secrets/kubernetes.io/serviceaccount/token") as token_file:
diff --git a/src/modules/report/base.py b/src/modules/report/base.py
index 36dda55..786cf26 100644
--- a/src/modules/report/base.py
+++ b/src/modules/report/base.py
@@ -24,7 +24,7 @@ class BaseReporter(object):
 
     def get_vulnerabilities(self):
         vulnerabilities_lock.acquire()
-        vulnerabilities_data = [{"location": "{}:{}".format(vuln.host, vuln.port) if vuln.host else "",
+        vulnerabilities_data = [{"location": vuln.location(),
                  "category": vuln.category.name,
                  "vulnerability": vuln.get_name(),
                  "description": vuln.explain(),
diff --git a/src/modules/report/plain.py b/src/modules/report/plain.py
index 282035a..70b88eb 100644
--- a/src/modules/report/plain.py
+++ b/src/modules/report/plain.py
@@ -82,7 +82,7 @@ class PlainReporter(object):
 
         vulnerabilities_lock.acquire()
         for vuln in vulnerabilities:
-            row = ["{}:{}".format(vuln.host, vuln.port) if vuln.host else "", vuln.category.name, vuln.get_name(), vuln.explain()]
+            row = [vuln.location(), vuln.category.name, vuln.get_name(), vuln.explain()]
             evidence = str(vuln.evidence)[:EVIDENCE_PREVIEW] + "..." if len(str(vuln.evidence)) > EVIDENCE_PREVIEW else str(vuln.evidence)
             row.append(evidence)
             vuln_table.add_row(row)