From a8762a4adabe30ae1983ba239b883d527b04229a Mon Sep 17 00:00:00 2001 From: oriagmon Date: Wed, 17 Oct 2018 11:32:50 +0300 Subject: [PATCH] Fixed PR: Used the same lock for each variable & changed locks names --- kube-hunter.py | 19 +++++++++---------- src/core/events/types/common.py | 13 +++++++------ src/modules/report/collector.py | 20 +++++++++++++++++--- src/modules/report/plain.py | 23 ++++++++++------------- src/modules/report/yaml.py | 20 ++++++++++---------- 5 files changed, 53 insertions(+), 42 deletions(-) diff --git a/kube-hunter.py b/kube-hunter.py index 77fa06c..90b7c99 100755 --- a/kube-hunter.py +++ b/kube-hunter.py @@ -82,7 +82,7 @@ def list_hunters(): print("\nPassive Hunters:\n----------------") for i, (hunter, docs) in enumerate(handler.passive_hunters.items()): name, docs = parse_docs(hunter, docs) - print("* {}\n {}\n".format( name, docs)) + print("* {}\n {}\n".format(name, docs)) if config.active: print("\n\nActive Hunters:\n---------------") @@ -91,10 +91,11 @@ def list_hunters(): print("* {}\n {}\n".format( name, docs)) -tlock3 = threading.Lock() -tlock3.acquire() +global hunt_started_lock +hunt_started_lock = threading.Lock() +hunt_started_lock.acquire() hunt_started = False -tlock3.release() +hunt_started_lock.release() def main(): @@ -112,10 +113,9 @@ def main(): if not any(scan_options): if not interactive_set_config(): return - tlock = threading.Lock() - tlock.acquire() + hunt_started_lock.acquire() hunt_started = True - tlock.release() + hunt_started_lock.release() handler.publish_event(HuntStarted()) handler.publish_event(HostScanEvent()) @@ -127,14 +127,13 @@ def main(): except EOFError: logging.error("\033[0;31mPlease run again with -it\033[0m") finally: - tlock2 = threading.Lock() - tlock2.acquire() + hunt_started_lock.acquire() if hunt_started: handler.publish_event(HuntFinished()) handler.join() handler.free() logging.debug("Cleaned Queue") - tlock2.release() + hunt_started_lock.release() if __name__ == '__main__': diff --git a/src/core/events/types/common.py b/src/core/events/types/common.py index d19dc2f..1822f12 100644 --- a/src/core/events/types/common.py +++ b/src/core/events/types/common.py @@ -65,10 +65,12 @@ class Vulnerability(object): def explain(self): return self.__doc__ -tlock1 = threading.Lock() -tlock1.acquire() + +global event_id_count_lock +event_id_count_lock = threading.Lock() +event_id_count_lock.acquire() event_id_count = 0 -tlock1.release() +event_id_count_lock.release() """ Discovery/Hunting Events """ @@ -78,11 +80,10 @@ class NewHostEvent(Event): global event_id_count self.host = host self.cloud = cloud - tlock = threading.Lock() - tlock.acquire() + event_id_count_lock.acquire() self.event_id = event_id_count event_id_count += 1 - tlock.release() + event_id_count_lock.release() def __str__(self): return str(self.host) diff --git a/src/modules/report/collector.py b/src/modules/report/collector.py index fa50cb0..5bf0a25 100644 --- a/src/modules/report/collector.py +++ b/src/modules/report/collector.py @@ -5,8 +5,21 @@ from src.core.events import handler from src.core.events.types import Event, Service, Vulnerability, HuntFinished, HuntStarted import threading + +global services_lock +services_lock = threading.Lock() +services_lock.acquire() services = list() +services_lock.release() + +global vulnerabilities_lock +vulnerabilities_lock = threading.Lock() +vulnerabilities_lock.acquire() vulnerabilities = list() +vulnerabilities_lock.release() + + + def console_trim(text, prefix=' '): @@ -38,12 +51,12 @@ class Collector(object): def execute(self): """function is called only when collecting data""" - tlock = threading.Lock() - tlock.acquire() global services, vulnerabilities bases = self.event.__class__.__mro__ if Service in bases: + services_lock.acquire() services.append(self.event) + services_lock.release() import datetime logging.info("|\n| {name}:\n| type: open service\n| service: {name}\n|_ host: {host}:{port}".format( host=self.event.host, @@ -53,7 +66,9 @@ class Collector(object): )) elif Vulnerability in bases: + vulnerabilities_lock.acquire() vulnerabilities.append(self.event) + vulnerabilities_lock.release() logging.info( "|\n| {name}:\n| type: vulnerability\n| host: {host}:{port}\n| description: \n{desc}".format( name=self.event.get_name(), @@ -61,7 +76,6 @@ class Collector(object): port=self.event.port, desc=wrap_last_line(console_trim(self.event.explain(), '| ')) )) - tlock.release() class TablesPrinted(Event): diff --git a/src/modules/report/plain.py b/src/modules/report/plain.py index 64d2dda..0464e6e 100644 --- a/src/modules/report/plain.py +++ b/src/modules/report/plain.py @@ -3,7 +3,7 @@ from __future__ import print_function from prettytable import ALL, PrettyTable from __main__ import config -from collector import services, vulnerabilities +from collector import services, vulnerabilities,services_lock, vulnerabilities_lock import threading EVIDENCE_PREVIEW = 40 @@ -15,8 +15,7 @@ class PlainReporter(object): def get_report(self): """generates report tables""" output = "" - tlock = threading.Lock - tlock.acquire() + services_lock.acquire() if len(services): output += self.nodes_table() if not config.mapping: @@ -27,7 +26,7 @@ class PlainReporter(object): output += "\nNo vulnerabilities were found" else: print("\nKube Hunter couldn't find any clusters") - tlock.release() + services_lock.release() # print("\nKube Hunter couldn't find any clusters. {}".format("Maybe try with --active?" if not config.active else "")) return output @@ -41,13 +40,12 @@ class PlainReporter(object): nodes_table.header_style = "upper" # TODO: replace with sets id_memory = list() - tlock = threading.Lock - tlock.acquire() + services_lock.acquire() for service in services: if service.event_id not in id_memory: nodes_table.add_row(["Node/Master", service.host]) id_memory.append(service.event_id) - tlock.release() + services_lock.acquire() return "\nNodes\n{}\n".format(nodes_table) def services_table(self): @@ -58,11 +56,10 @@ class PlainReporter(object): services_table.sortby = "Service" services_table.reversesort = True services_table.header_style = "upper" - tlock = threading.Lock - tlock.acquire() + services_lock.acquire() for service in services: services_table.add_row([service.get_name(), "{}:{}{}".format(service.host, service.port, service.get_path()), service.explain()]) - tlock.release() + services_lock.release() return "\nDetected Services\n{}\n".format(services_table) def vulns_table(self): @@ -74,12 +71,12 @@ class PlainReporter(object): vuln_table.reversesort = True vuln_table.padding_width = 1 vuln_table.header_style = "upper" - tlock = threading.Lock - tlock.acquire() + + 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()] 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) - tlock.release() + vulnerabilities_lock.release() return "\nVulnerabilities\n{}\n".format(vuln_table) diff --git a/src/modules/report/yaml.py b/src/modules/report/yaml.py index 2a0b372..80ae1fd 100644 --- a/src/modules/report/yaml.py +++ b/src/modules/report/yaml.py @@ -2,17 +2,20 @@ import StringIO from ruamel.yaml import YAML -from collector import services, vulnerabilities +from collector import services, vulnerabilities, services_lock, vulnerabilities_lock import threading + class YAMLReporter(object): def get_report(self): yaml = YAML() + vulnerabilities_lock.acquire() report = { "nodes": self.get_nodes(), "services": self.get_services(), "vulnerabilities": self.get_vulenrabilities() } + vulnerabilities_lock.release() output = StringIO.StringIO() yaml.dump(report, output) return output.getvalue() @@ -20,34 +23,31 @@ class YAMLReporter(object): def get_nodes(self): nodes = list() node_locations = set() - tlock = threading.Lock - tlock.acquire() + services_lock.acquire() for service in services: node_location = str(service.host) if node_location not in node_locations: nodes.append({"type": "Node/Master", "location": str(service.host)}) node_locations.add(node_location) - tlock.release() + services_lock.release() return nodes def get_services(self): - tlock = threading.Lock - tlock.acquire() + services_lock.acquire() services_data = [{"service": service.get_name(), "location": "{}:{}{}".format(service.host, service.port, service.get_path()), "description": service.explain()} for service in services] - tlock.release() + services_lock.release() return services_data def get_vulenrabilities(self): - tlock = threading.Lock - tlock.acquire() + vulnerabilities_lock.acquire() vulnerabilities_data = [{"location": "{}:{}".format(vuln.host, vuln.port) if vuln.host else "", "category": vuln.category.name, "vulnerability": vuln.get_name(), "description": vuln.explain(), "evidence": str(vuln.evidence)} for vuln in vulnerabilities] - tlock.release() + vulnerabilities_lock.release() return vulnerabilities_data