if script fails, report to telegram too

This commit is contained in:
Evgeny Zinoviev 2022-03-18 04:48:55 +03:00
parent 37ec4be9af
commit 5644c92731

View File

@ -2,10 +2,12 @@
import logging import logging
import yaml import yaml
import math import math
import html
from argparse import ArgumentParser from argparse import ArgumentParser
from lib.worker import Worker from lib.worker import Worker
from lib.results import Results from lib.results import Results
from ch1p import telegram_notify
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -40,47 +42,56 @@ def main():
if not args.no_telegram: if not args.no_telegram:
assert 'telegram' in config assert 'telegram' in config
# let's go try:
results = Results() # let's go
max_threads = math.inf if args.threads_limit == 0 else args.threads_limit results = Results()
active_threads = 1 max_threads = math.inf if args.threads_limit == 0 else args.threads_limit
active_threads = 1
def get_active_threads(): def get_active_threads():
n = active_threads n = active_threads
if workers: if workers:
n += workers[0].concurrency n += workers[0].concurrency
return n return n
workers = [] workers = []
for name, data in config['servers'].items(): for name, data in config['servers'].items():
w = Worker(name, data['host'], data['opened'], data['ignore'], w = Worker(name, data['host'], data['opened'], data['ignore'],
concurrency=int(data['concurrency']) if 'concurrency' in data else args.concurrency, concurrency=int(data['concurrency']) if 'concurrency' in data else args.concurrency,
timeout=int(data['timeout']) if 'timeout' in data else args.timeout) timeout=int(data['timeout']) if 'timeout' in data else args.timeout)
workers.append(w) workers.append(w)
current_workers = [] current_workers = []
while workers: while workers:
w = workers.pop(0) w = workers.pop(0)
active_threads += w.concurrency+1 active_threads += w.concurrency+1
current_workers.append(w) current_workers.append(w)
w.start() w.start()
while current_workers and get_active_threads() >= max_threads: while current_workers and get_active_threads() >= max_threads:
for cw in current_workers: for cw in current_workers:
cw.join(timeout=0.1) cw.join(timeout=0.1)
if not cw.is_alive(): if not cw.is_alive():
results.add(cw) results.add(cw)
current_workers.remove(cw) current_workers.remove(cw)
active_threads -= cw.concurrency+1 active_threads -= cw.concurrency+1
for cw in current_workers: for cw in current_workers:
cw.join() cw.join()
results.add(cw) results.add(cw)
if results.has_warnings() and not args.no_telegram: if results.has_warnings() and not args.no_telegram:
results.notify(chat_id=config['telegram']['chat-id'], results.notify(chat_id=config['telegram']['chat-id'],
token=config['telegram']['token']) token=config['telegram']['token'])
except KeyboardInterrupt:
pass
except Exception as e:
if not args.no_telegram:
telegram_notify(html.escape(str(e)),
parse_mode='html',
chat_id=config['telegram']['chat-id'],
token=config['telegram']['token'])
if __name__ == '__main__': if __name__ == '__main__':