#!/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()