/
threads.py
103 lines (85 loc) · 3.39 KB
/
threads.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import time
import threading
import traceback
from pocsuite3.lib.core.data import conf
from pocsuite3.lib.core.data import kb
from pocsuite3.lib.core.data import logger
from pocsuite3.lib.core.exception import PocsuiteConnectionException
from pocsuite3.lib.core.exception import PocsuiteThreadException
from pocsuite3.lib.core.exception import PocsuiteUserQuitException
from pocsuite3.lib.core.exception import PocsuiteValueException
from pocsuite3.lib.core.settings import MAX_NUMBER_OF_THREADS
def exception_handled_function(thread_function, args=(), silent=False):
try:
thread_function(*args)
except KeyboardInterrupt:
kb.thread_continue = False
kb.thread_exception = True
raise
except Exception as ex:
if not silent:
logger.error("thread {0}: {1}".format(threading.currentThread().getName(), str(ex)))
if conf.verbose > 1:
traceback.print_exc()
def run_threads(num_threads, thread_function, args: tuple = (), forward_exception=True, start_msg=True):
threads = []
kb.multi_thread_mode = True
kb.thread_continue = True
kb.thread_exception = False
try:
if num_threads > 1:
if start_msg:
info_msg = "staring {0} threads".format(num_threads)
logger.info(info_msg)
if num_threads > MAX_NUMBER_OF_THREADS:
warn_msg = ""
logger.warn(warn_msg)
else:
thread_function()
return
# Start the threads
for num_threads in range(num_threads):
thread = threading.Thread(target=exception_handled_function, name=str(num_threads),
args=(thread_function, args))
thread.setDaemon(True)
try:
thread.start()
except Exception as ex:
err_msg = "error occurred while starting new thread ('{0}')".format(str(ex))
logger.critical(err_msg)
break
threads.append(thread)
# And wait for them to all finish
alive = True
while alive:
alive = False
for thread in threads:
if thread.isAlive():
alive = True
time.sleep(0.1)
except (KeyboardInterrupt, PocsuiteUserQuitException) as ex:
kb.thread_continue = False
kb.thread_exception = True
if num_threads > 1:
logger.info("waiting for threads to finish{0}".format(
" (Ctrl+C was pressed)" if isinstance(ex, KeyboardInterrupt) else ""))
try:
while threading.activeCount() > 1:
pass
except KeyboardInterrupt:
raise PocsuiteThreadException("user aborted (Ctrl+C was pressed multiple times)")
if forward_exception:
raise
except (PocsuiteConnectionException, PocsuiteValueException) as ex:
kb.thread_exception = True
logger.error("thread {0}: {1}".format(threading.currentThread().getName(), str(ex)))
if conf.verbose > 1:
traceback.print_exc()
except Exception as ex:
kb.thread_exception = True
logger.error("thread {0}: {1}".format(threading.currentThread().getName(), str(ex)))
traceback.print_exc()
finally:
kb.multi_thread_mode = False
kb.thread_continue = True
kb.thread_exception = False