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