main.py

#!/usr/bin/env python3

# should be called before everything else
# it's very fucking important!
def init_uvloop():
    import uvloop
    uvloop.install()


init_uvloop()

from proxy_py import settings
from processor import Processor
from server.proxy_provider_server import ProxyProviderServer
from statistics import statistics
from checkers.base_checker import BaseChecker
from tools import test_collector
import materialized_view_updater
import asyncio
import logging
import argparse
import collectors_list
import subprocess
import sys

test_collector_path = None
main_logger = None


def process_cmd_arguments():
    global test_collector_path

    def str_to_bool(value):
        if value.lower() in ("yes", "true", "t", "y", "1"):
            return True
        elif value.lower() in ("no", "false", "f", "n", "0"):
            return False
        else:
            raise argparse.ArgumentTypeError("Boolean value expected.")

    cmd_parser = argparse.ArgumentParser()
    cmd_parser.add_argument("--debug", type=str_to_bool,
                            help="override settings' debug value")
    cmd_parser.add_argument(
        "--proxy-checking-timeout", type=float, help="override settings' proxy checking timeout"
    )
    cmd_parser.add_argument(
        "--test-collector", help="test collector with a given path")

    args = cmd_parser.parse_args()

    if args.debug is not None:
        settings.DEBUG = args.debug

    if args.proxy_checking_timeout is not None:
        if args.proxy_checking_timeout < 0:
            raise ValueError("--proxy-checking-timeout should be positive")

        settings.PROXY_CHECKING_TIMEOUT = args.proxy_checking_timeout

    test_collector_path = args.test_collector


def prepare_loggers():
    global main_logger

    asyncio_logger = logging.getLogger("asyncio")
    asyncio_logger.setLevel(logging.DEBUG if settings.DEBUG else logging.INFO)
    asyncio_logger_handler = logging.StreamHandler(sys.stdout)
    asyncio_logger_handler.setLevel(logging.DEBUG if settings.DEBUG else logging.INFO)
    asyncio_logger_handler.setFormatter(
        logging.Formatter(settings.LOG_FORMAT_STRING)
    )
    asyncio_logger.addHandler(asyncio_logger_handler)
    asyncio.get_event_loop().set_debug(settings.DEBUG)

    main_logger = logging.getLogger("proxy_py/main")
    main_logger.setLevel(logging.DEBUG if settings.DEBUG else logging.INFO)
    logger_handler = logging.StreamHandler(sys.stdout)
    logger_handler.setFormatter(logging.Formatter(settings.LOG_FORMAT_STRING))
    logger_handler.setLevel(logging.DEBUG if settings.DEBUG else logging.INFO)

    main_logger.addHandler(logger_handler)


async def core():
    process_cmd_arguments()
    prepare_loggers()

    if test_collector_path is not None:
        return await test_collector.run(test_collector_path)

    proxy_processor = Processor.get_instance()

    try:
        code = await asyncio.gather(*[
            proxy_processor.worker(),
            statistics.worker(),
            materialized_view_updater.worker(),
        ])
        BaseChecker.clean()
        return code
    except KeyboardInterrupt:
        pass
    except BaseException as ex:
        main_logger.exception(ex)
        print("critical error happened, see logs/main.log")
        return 1

    return 0


async def print_collectors():
    for collector_name in collectors_list.collectors.keys():
        print(collector_name)


def server():
    proxy_provider_server = ProxyProviderServer(
        settings.PROXY_PROVIDER_SERVER_ADDRESS['HOST'],
        settings.PROXY_PROVIDER_SERVER_ADDRESS['PORT'],
    )

    return proxy_provider_server.start(asyncio.get_event_loop())


def print_help():
    print("""Usage: ./main.py [COMMAND] [OPTION]...
Runs proxy_py

The commands are:

    "core" - runs core part of the project (proxies parsing and processing)
    "print_collectors" - prints collectors
    "server" - runs server for providing API
    "" - runs both

use ./main.py COMMAND --help to get more information

Project's page: https://github.com/DevAlone/proxy_py
""")


def main():
    if len(sys.argv) < 2:
        # run default configuration
        # server
        p = subprocess.Popen(["python3", sys.argv[0], 'server'])

        # and core
        code = asyncio.get_event_loop().run_until_complete(core())
        p.wait()
        return code

    command = sys.argv[1].strip()
    sys.argv = sys.argv[1:]
    try:
        return {
            'core': lambda: asyncio.get_event_loop().run_until_complete(core()),
            'print_collectors': lambda: asyncio.get_event_loop().run_until_complete(print_collectors()),
            'server': server,
        }[command]()
    except KeyError:
        print_help()
        return 0


if __name__ == "__main__":
    exit(main())