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
|
#!/usr/bin/env python3
import logging
import yaml
import math
from argparse import ArgumentParser
from lib.worker import Worker
from lib.results import Results
logger = logging.getLogger(__name__)
def main():
parser = ArgumentParser()
parser.add_argument('--config', type=str, required=True,
help='path to config file in yaml format')
parser.add_argument('--verbose', action='store_true',
help='set logging level to DEBUG')
parser.add_argument('--concurrency', default=200, type=int,
help='default number of threads per target')
parser.add_argument('--timeout', default=5, type=int,
help='default timeout')
parser.add_argument('--threads-limit', default=0, type=int,
help='global threads limit')
parser.add_argument('--no-telegram', action='store_true',
help='just print results, don\'t send to telegram')
args = parser.parse_args()
# setup logging
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=(logging.DEBUG if args.verbose else logging.INFO))
# load config
with open(args.config, 'r') as f:
config = yaml.safe_load(f)
# pprint(config)
assert isinstance(config, dict)
assert 'servers' in config
if not args.no_telegram:
assert 'telegram' in config
# let's go
results = Results()
max_threads = math.inf if args.threads_limit == 0 else args.threads_limit
active_threads = 1
def get_active_threads():
n = active_threads
if workers:
n += workers[0].concurrency
return n
workers = []
for name, data in config['servers'].items():
w = Worker(name, data['host'], data['opened'], data['ignore'],
concurrency=int(data['concurrency']) if 'concurrency' in data else args.concurrency,
timeout=int(data['timeout']) if 'timeout' in data else args.timeout)
workers.append(w)
current_workers = []
while workers:
w = workers.pop(0)
active_threads += w.concurrency+1
current_workers.append(w)
w.start()
while current_workers and get_active_threads() >= max_threads:
for cw in current_workers:
cw.join(timeout=0.1)
if not cw.is_alive():
results.add(cw)
current_workers.remove(cw)
active_threads -= cw.concurrency+1
for cw in current_workers:
cw.join()
results.add(cw)
if results.has_warnings() and not args.no_telegram:
results.notify(chat_id=config['telegram']['chat-id'],
token=config['telegram']['token'])
if __name__ == '__main__':
main()
|
