Fixed PR: Used the same lock for each variable & changed locks names

This commit is contained in:
oriagmon
2018-10-17 11:32:50 +03:00
parent 086a403559
commit a8762a4ada
5 changed files with 53 additions and 42 deletions

View File

@@ -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__':

View File

@@ -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)

View File

@@ -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):

View File

@@ -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)

View File

@@ -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