Прокси сервер в отдельном процессе, полезен для тестирования (имитация разрыва соединения)

proxy_server.py

import asyncio
import logging
from multiprocessing import Process

g_proxy_addr = ''
g_proxy_port = 0

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


async def pipe(reader, writer):
    try:
        while not reader.at_eof():
            writer.write(await reader.read(2048))
    finally:
        writer.close()


async def handle_client(local_reader, local_writer):
    try:
        remote_reader, remote_writer = await asyncio.open_connection(
            g_proxy_addr, g_proxy_port)
        pipe1 = pipe(local_reader, remote_writer)
        pipe2 = pipe(remote_reader, local_writer)
        await asyncio.gather(pipe1, pipe2)
    finally:
        local_writer.close()


def main(listen_addr, listen_port, proxy_addr, proxy_port):
    loop = asyncio.get_event_loop()
    coro = asyncio.start_server(handle_client, listen_addr, listen_port)
    server = loop.run_until_complete(coro)
    global g_proxy_addr
    g_proxy_addr = proxy_addr
    global g_proxy_port
    g_proxy_port = proxy_port
    # Serve requests until Ctrl+C is pressed
    # print('Serving on {}'.format(server.sockets[0].getsockname()))
    try:
        loop.run_forever()
    finally:
        server.close()
        loop.run_until_complete(server.wait_closed())
        loop.close()


def proxy_server(listen_addr, listen_port, proxy_addr, proxy_port):
    """
    Прокси сервер для проброса портов, он является генератором и принимает 4 команды:
    1) pause - останавливает выполнение сервера
    2) continue - восстанавливает работу сервера
    3) status - пишет в лог статус сервера
    4) exit - завершает сервер
    пример:
    >>> gen_proxy = proxy_server('127.0.0.1', 15432, '10.79.8.49', 5432)
    >>> next(gen_proxy)
    # прокси сервер выполняет проксирование...
    >>> gen_proxy.send('pause')
    # прокси сервер становлен
    >>> gen_proxy.send('continue')
    :param listen_addr: прослушиваемый адрес
    :param listen_port:  прослушиваемый порт
    :param proxy_addr: адрес куда проксируем
    :param proxy_port:  порт куда проксируем
    :yield: сообщение с ошибкой если передана несуществующая команда либо ничего
    """
    p = Process(target=main, args=(listen_addr, listen_port, proxy_addr, proxy_port))
    p.start()
    logger.info('starting proxy server')
    while True:
        command = yield
        logger.info(f"Command:{command}")
        if command == 'pause' and p.is_alive():
            p.terminate()
            p.join()
            p = None
            logger.info('pausing proxy server')
        elif command == 'continue' and not p:
            p = Process(target=main, args=(listen_addr, listen_port, proxy_addr, proxy_port))
            p.start()
            logger.info('continuing proxy server')
        elif command == 'status':
            logger.info(f"Status:{str(p)}")
        elif command == 'exit' and p.is_alive():
            p.terminate()
            p.join()
            logger.info('terminate')
            break
        else:
            logger.info('Error: Undefined command')
            yield 'Error: Undefined command'