if script fails, report to telegram too
This commit is contained in:
parent
37ec4be9af
commit
5644c92731
@ -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__':
|
||||||
|
Loading…
x
Reference in New Issue
Block a user