qd.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# vim: set et sw=4 ts=4 sts=4 ff=unix fenc=utf8:
# Author: Binux<i@binux.me>
#         http://binux.me
# Created on 2014-08-18 12:17:21

import asyncio
import json
import socket
import sys

import config
from libs.fetcher import Fetcher
from libs.log import Log
from run import start_server

config.display_import_warning = False

logger_qd = Log('QD').getlogger()

# 判断 端口 是否被占用


def check_port(port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.connect(('127.0.0.1', port))
        s.shutdown(2)
        logger_qd.debug('Port %s is used' , port)
        return False
    except Exception:
        logger_qd.debug('Port %s is available' , port)
        return True


def usage():
    print(f"{sys.argv[0]} tpl.har [--key=value] [env.json]")
    sys.exit(1)


if __name__ == '__main__':
    if len(sys.argv) < 3:
        usage()

    # load tpl
    tpl_file = sys.argv[1]
    try:
        # deepcode ignore PT: tpl_file is a file
        tpl = json.load(open(tpl_file, encoding='utf-8'))
    except Exception as e:
        logger_qd.error(e, exc_info=config.traceback_print)
        usage()

    # load env
    variables = {}
    env = {}
    ENV_FILE = None
    for each in sys.argv[2:]:
        if each.startswith('--'):
            key, value = each.split('=', 1)
            key = key.lstrip('--')
            variables[key] = value
        else:
            ENV_FILE = each
    if ENV_FILE:
        try:
            # deepcode ignore PT: env_file is a file
            env = json.load(open(ENV_FILE, encoding='utf-8'))
        except Exception as e:
            logger_qd.error(e, exc_info=config.traceback_print)
            usage()
    if 'variables' not in env or not isinstance(env['variables'], dict) \
            or 'session' not in env:
        env = {
            'variables': env,
            'session': [],
        }
    env['variables'].update(variables)

    MANUAL_START = check_port(config.port)
    if MANUAL_START:
        logger_qd.info('QD service is not running on port %s' , config.port)
        logger_qd.info('QD service will be started on port %s' , config.port)
        # 创建新进程, 以执行 run 中的 main 异步函数
        import multiprocessing
        p = multiprocessing.Process(target=start_server)
        p.start()
        # 循环检测端口是否被占用, 如果被占用, 则继续下一步
        while True:
            if not check_port(config.port):
                break
            else:
                import time
                time.sleep(1)
    else:
        logger_qd.info('QD service is running on port %s' , config.port)

    # do fetch
    ioloop = asyncio.new_event_loop()
    asyncio.set_event_loop(ioloop)
    result: asyncio.Task = asyncio.ensure_future(Fetcher().do_fetch(tpl, env), loop=ioloop)
    logger_qd.info('QD start to do fetch: %s', tpl_file)
    ioloop.run_until_complete(result)
    ioloop.stop()

    try:
        env, _ = result.result()
    except Exception as e:
        logger_qd.error('QD failed: %s', e, exc_info=config.traceback_print)
    else:
        logger_qd.info('QD success! Results:\n %s', env.get('variables', {}).get('__log__', '').replace('\\r\\n', '\r\n'))

    if MANUAL_START:
        p.terminate()
        p.join()
        logger_qd.info('QD service is ended. ')