Skip to content

Commit

Permalink
feat: add threading for check and exploit function
Browse files Browse the repository at this point in the history
  • Loading branch information
@TuuuNya
TuuuNya committed Mar 11, 2019
1 parent 61121bc commit 835f148
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 55 deletions.
8 changes: 5 additions & 3 deletions lib/BaseExploit.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,19 +33,21 @@ def get_info(self):
info[field_name] = getattr(self, field_name)
return info

def register_tcp_target(self, port_value=None, timeout_value=10):
def register_tcp_target(self, port_value=None, timeout_value=5, threads_value=1):
self.target_type = "tcp"
self.register_options([
ExploitOption(name="HOST", required=True, description="The IP address to be tested"),
ExploitOption(name="PORT", required=True, description="The port to be tested", value=port_value),
ExploitOption(name="TIMEOUT", required=True, description="Connection timeout", value=timeout_value)
ExploitOption(name="TIMEOUT", required=True, description="Connection timeout", value=timeout_value),
ExploitOption(name="THREADS", required=True, description="The number of threads", value=threads_value)
])

def register_http_target(self, timeout_value=10):
def register_http_target(self, timeout_value=5, threads_value=1):
self.target_type = "http"
self.register_options([
ExploitOption(name="URL", required=True, description="The url to be tested"),
ExploitOption(name="TIMEOUT", required=True, description="Connection timeout", value=timeout_value),
ExploitOption(name="THREADS", required=True, description="The number of threads", value=threads_value)
])

def update_info(self, info):
Expand Down
144 changes: 92 additions & 52 deletions lib/Pocket.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import argparse
from threading import Thread
from queue import Queue
from lib.cmd2 import Cmd, with_category, with_argparser
from art import text2art, art
Expand Down Expand Up @@ -28,6 +29,7 @@ class Pocket(Cmd, Database):
def __init__(self):
super(Pocket, self).__init__()
Database.__init__(self)
self.thread_pool = list()
self.prompt = self.console_prompt + self.console_prompt_end
self.hidden_commands.extend(['alias', 'edit', 'macro', 'py', 'pyscript', 'shell', 'shortcuts', 'load'])
self.do_banner(None)
Expand Down Expand Up @@ -196,6 +198,33 @@ def do_show(self, content):
def do_run(self, args):
self.do_exploit(args=args)

def exploit_thread(self, targets_queue, target_type):
while not targets_queue.empty():
target = None
target_field = None
port = None

if target_type == "tcp":
[target, port] = module.parse_ip_port(targets_queue.get())
target_field = "HOST"
elif target_type == "http":
target = targets_queue.get()
target_field = "URL"
exp = self.module_class.Exploit()
exp.options.set_option(target_field, target)
exp.options.set_option("TIMEOUT", self.module_instance.options.get_option("TIMEOUT"))
if port:
exp.options.set_option("PORT", port)
else:
exp.options.set_option("PORT", self.module_instance.options.get_option("PORT"))

exploit_result = exp.exploit()

if exploit_result.status:
self._print_item(exploit_result.success_message)
else:
self._print_item(exploit_result.error_message, color=Fore.RED)

@with_category(CMD_MODULE)
def do_exploit(self, args):
if not self.module_instance:
Expand Down Expand Up @@ -238,35 +267,29 @@ def do_exploit(self, args):

# 处理tcp类型的多目标
while not targets_queue.empty() and target_type == "tcp":
[target, port] = module.parse_ip_port(targets_queue.get())
thread_count = int(self.module_instance.options.get_option("THREADS"))

exp = self.module_class.Exploit()
exp.options.set_option(target_field, target)
exp.options.set_option("TIMEOUT", self.module_instance.options.get_option("TIMEOUT"))
if port:
exp.options.set_option("PORT", port)
else:
exp.options.set_option("PORT", self.module_instance.options.get_option("PORT"))
for i in range(thread_count):
_thread = Thread(target=self.exploit_thread, args=(targets_queue, target_type))
_thread.start()
self.thread_pool.append(_thread)

exploit_result = exp.exploit()
if exploit_result.status:
self._print_item(exploit_result.success_message)
else:
self._print_item(exploit_result.error_message, color=Fore.RED)
for th in self.thread_pool:
th.join()
self.thread_pool.clear()

# 处理http类型的多目标
while not targets_queue.empty() and target_type == "http":
target = targets_queue.get()
thread_count = int(self.module_instance.options.get_option("THREADS"))

exp = self.module_class.Exploit()
exp.options.set_option(target_field, target)
exp.options.set_option("TIMEOUT", self.module_instance.options.get_option("TIMEOUT"))
for i in range(thread_count):
_thread = Thread(target=self.exploit_thread, args=(targets_queue, target_type))
_thread.start()
self.thread_pool.append(_thread)

exploit_result = exp.exploit()
if exploit_result.status:
self._print_item(exploit_result.success_message)
else:
self._print_item(exploit_result.error_message, color=Fore.RED)
for th in self.thread_pool:
th.join()
self.thread_pool.clear()

self.poutput("{style}[*]{style_end} module execution completed".format(
style=Fore.BLUE + Style.BRIGHT,
Expand All @@ -286,6 +309,37 @@ def do_exploit(self, args):
style_end=Style.RESET_ALL
))

def check_thread(self, targets_queue, target_type):
while not targets_queue.empty():
target = None
target_field = None
port = None

if target_type == "tcp":
[target, port] = module.parse_ip_port(targets_queue.get())
target_field = "HOST"
elif target_type == "http":
target = targets_queue.get()
target_field = "URL"
exp = self.module_class.Exploit()
exp.options.set_option(target_field, target)
exp.options.set_option("TIMEOUT", self.module_instance.options.get_option("TIMEOUT"))
if port:
exp.options.set_option("PORT", port)
else:
exp.options.set_option("PORT", self.module_instance.options.get_option("PORT"))

exploit_result = exp.check()

if exploit_result is None:
self._print_item("Check Error: check function no results returned")
return None

if exploit_result.status:
self._print_item(exploit_result.success_message)
else:
self._print_item(exploit_result.error_message, color=Fore.RED)

@with_category(CMD_MODULE)
def do_check(self, args):
if not self.module_instance:
Expand Down Expand Up @@ -328,43 +382,29 @@ def do_check(self, args):

# 处理TCP类型的多个目标
while not targets_queue.empty() and target_type == "tcp":
[target, port] = module.parse_ip_port(targets_queue.get())
exp = self.module_class.Exploit()
exp.options.set_option(target_field, target)
exp.options.set_option("TIMEOUT", self.module_instance.options.get_option("TIMEOUT"))
if port:
exp.options.set_option("PORT", port)
else:
exp.options.set_option("PORT", self.module_instance.options.get_option("PORT"))

exploit_result = exp.check()
thread_count = int(self.module_instance.options.get_option("THREADS"))

if exploit_result is None:
self._print_item("Check Error: check function no results returned")
return None
for i in range(thread_count):
_thread = Thread(target=self.check_thread, args=(targets_queue, target_type))
_thread.start()
self.thread_pool.append(_thread)

if exploit_result.status:
self._print_item(exploit_result.success_message)
else:
self._print_item(exploit_result.error_message, color=Fore.RED)
for th in self.thread_pool:
th.join()
self.thread_pool.clear()

# 处理http类型的多个目标
while not targets_queue.empty() and target_type == "http":
target = targets_queue.get()
exp = self.module_class.Exploit()
exp.options.set_option(target_field, target)
exp.options.set_option("TIMEOUT", self.module_instance.options.get_option("TIMEOUT"))

exploit_result = exp.check()
thread_count = int(self.module_instance.options.get_option("THREADS"))

if exploit_result is None:
self._print_item("Check Error: check function no results returned")
return None
for i in range(thread_count):
_thread = Thread(target=self.check_thread, args=(targets_queue, target_type))
_thread.start()
self.thread_pool.append(_thread)

if exploit_result.status:
self._print_item(exploit_result.success_message)
else:
self._print_item(exploit_result.error_message, color=Fore.RED)
for th in self.thread_pool:
th.join()
self.thread_pool.clear()

self.poutput("{style}[*]{style_end} module execution completed".format(
style=Fore.BLUE + Style.BRIGHT,
Expand Down

0 comments on commit 835f148

Please sign in to comment.