104 lines
3.5 KiB
Python
Executable File
104 lines
3.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import logging
|
|
import yaml
|
|
import math
|
|
import html
|
|
import traceback
|
|
|
|
from argparse import ArgumentParser
|
|
from lib.worker import Worker
|
|
from lib.results import Results
|
|
from ch1p import telegram_notify
|
|
|
|
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
|
|
|
|
try:
|
|
# 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'],
|
|
opened=data['opened'],
|
|
ignore=data['ignore'] if 'ignore' in data else None,
|
|
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'])
|
|
except KeyboardInterrupt:
|
|
pass
|
|
except Exception as e:
|
|
if not args.no_telegram:
|
|
telegram_notify(html.escape(traceback.format_exc()),
|
|
parse_mode='html',
|
|
chat_id=config['telegram']['chat-id'],
|
|
token=config['telegram']['token'])
|
|
else:
|
|
traceback.print_exc()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|