From 229347e9fa85497f35170d9c8143552ada356503 Mon Sep 17 00:00:00 2001 From: oriagmon Date: Tue, 16 Oct 2018 16:53:07 +0300 Subject: [PATCH] Attempting to solve the threading bug, I will do more checking to be sure its gone --- kube-hunter.py | 13 +++++++++++-- src/core/events/handler.py | 2 +- src/core/events/types/common.py | 1 + src/modules/report/collector.py | 1 + src/modules/report/plain.py | 13 +++++++++++++ src/modules/report/yaml.py | 19 +++++++++++++++---- 6 files changed, 42 insertions(+), 7 deletions(-) diff --git a/kube-hunter.py b/kube-hunter.py index 0db99a4..08a954d 100755 --- a/kube-hunter.py +++ b/kube-hunter.py @@ -3,6 +3,7 @@ from __future__ import print_function import argparse import logging +import threading try: raw_input # Python 2 @@ -87,8 +88,10 @@ def list_hunters(): name, docs = parse_docs(hunter, docs) print("* {}\n {}\n".format( name, docs)) - +tlock3 = threading.Lock +tlock3.acquire() hunt_started = False +tlock3.release() def main(): global hunt_started scan_options = [ @@ -104,8 +107,10 @@ def main(): if not any(scan_options): if not interactive_set_config(): return - + tlock = threading.Lock + tlock.acquire() hunt_started = True + tlock.release() handler.publish_event(HuntStarted()) handler.publish_event(HostScanEvent()) @@ -117,11 +122,15 @@ def main(): except EOFError: logging.error("\033[0;31mPlease run again with -it\033[0m") finally: + tlock2 = threading.Lock() + tlock2.acquire() if hunt_started: handler.publish_event(HuntFinished()) handler.join() handler.free() logging.debug("Cleaned Queue") + tlock2.release() + if __name__ == '__main__': main() diff --git a/src/core/events/handler.py b/src/core/events/handler.py index 051087d..8e3b109 100644 --- a/src/core/events/handler.py +++ b/src/core/events/handler.py @@ -10,9 +10,9 @@ from __main__ import config from ..types import ActiveHunter, Hunter from ...core.events.types import HuntFinished +import threading working_count = 0 -lock = Lock() # Inherits Queue object, handles events asynchronously class EventQueue(Queue, object): diff --git a/src/core/events/types/common.py b/src/core/events/types/common.py index f123b71..014953c 100644 --- a/src/core/events/types/common.py +++ b/src/core/events/types/common.py @@ -1,6 +1,7 @@ import logging import requests import json +import threading class Event(object): def __init__(self): diff --git a/src/modules/report/collector.py b/src/modules/report/collector.py index d316439..db6b77f 100644 --- a/src/modules/report/collector.py +++ b/src/modules/report/collector.py @@ -3,6 +3,7 @@ import logging from __main__ import config from src.core.events import handler from src.core.events.types import Event, Service, Vulnerability, HuntFinished, HuntStarted +import threading services = list() vulnerabilities = list() diff --git a/src/modules/report/plain.py b/src/modules/report/plain.py index 5f7b5d3..64d2dda 100644 --- a/src/modules/report/plain.py +++ b/src/modules/report/plain.py @@ -4,6 +4,7 @@ from prettytable import ALL, PrettyTable from __main__ import config from collector import services, vulnerabilities +import threading EVIDENCE_PREVIEW = 40 MAX_TABLE_WIDTH = 20 @@ -14,6 +15,8 @@ class PlainReporter(object): def get_report(self): """generates report tables""" output = "" + tlock = threading.Lock + tlock.acquire() if len(services): output += self.nodes_table() if not config.mapping: @@ -24,6 +27,7 @@ class PlainReporter(object): output += "\nNo vulnerabilities were found" else: print("\nKube Hunter couldn't find any clusters") + tlock.release() # print("\nKube Hunter couldn't find any clusters. {}".format("Maybe try with --active?" if not config.active else "")) return output @@ -37,10 +41,13 @@ class PlainReporter(object): nodes_table.header_style = "upper" # TODO: replace with sets id_memory = list() + tlock = threading.Lock + tlock.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() return "\nNodes\n{}\n".format(nodes_table) def services_table(self): @@ -51,8 +58,11 @@ class PlainReporter(object): services_table.sortby = "Service" services_table.reversesort = True services_table.header_style = "upper" + tlock = threading.Lock + tlock.acquire() for service in services: services_table.add_row([service.get_name(), "{}:{}{}".format(service.host, service.port, service.get_path()), service.explain()]) + tlock.release() return "\nDetected Services\n{}\n".format(services_table) def vulns_table(self): @@ -64,9 +74,12 @@ class PlainReporter(object): vuln_table.reversesort = True vuln_table.padding_width = 1 vuln_table.header_style = "upper" + tlock = threading.Lock + tlock.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() return "\nVulnerabilities\n{}\n".format(vuln_table) diff --git a/src/modules/report/yaml.py b/src/modules/report/yaml.py index 07418b6..2a0b372 100644 --- a/src/modules/report/yaml.py +++ b/src/modules/report/yaml.py @@ -3,7 +3,7 @@ import StringIO from ruamel.yaml import YAML from collector import services, vulnerabilities - +import threading class YAMLReporter(object): def get_report(self): @@ -19,24 +19,35 @@ class YAMLReporter(object): def get_nodes(self): nodes = list() - node_locations = set() + node_locations = set() + tlock = threading.Lock + tlock.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() return nodes def get_services(self): - return [{"service": service.get_name(), + tlock = threading.Lock + tlock.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() + return services_data def get_vulenrabilities(self): - return [{"location": "{}:{}".format(vuln.host, vuln.port) if vuln.host else "", + tlock = threading.Lock + tlock.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() + return vulnerabilities_data