1
Fork 0

Use logging instead

This commit is contained in:
Jake Howard 2019-01-01 20:26:36 +00:00
parent a6a4069158
commit 88622440e0
Signed by: jake
GPG key ID: 57AFB45680EDD477

View file

@ -1,12 +1,18 @@
import argparse import argparse
import asyncio import asyncio
import logging
from collections import namedtuple from collections import namedtuple
from functools import partial from functools import partial
logger = logging.getLogger(__name__)
Route = namedtuple("Route", ["listen_port", "destination_host", "destination_port"]) Route = namedtuple("Route", ["listen_port", "destination_host", "destination_port"])
BUFFER_SIZE = 4096 BUFFER_SIZE = 4096
FORMAT = "%(levelname)s: %(message)s"
logging.basicConfig(format=FORMAT)
def destination_host_display(route: Route): def destination_host_display(route: Route):
return "{}:{}".format(route.destination_host, route.destination_port) return "{}:{}".format(route.destination_host, route.destination_port)
@ -32,8 +38,8 @@ async def handle_client(route, local_reader, local_writer):
await asyncio.gather( await asyncio.gather(
pipe(local_reader, remote_writer), pipe(remote_reader, local_writer) pipe(local_reader, remote_writer), pipe(remote_reader, local_writer)
) )
except ConnectionRefusedError: except (ConnectionRefusedError, OSError):
print("Connection to {} refused".format(destination_host_display(route))) logger.debug("Connection to {} refused".format(destination_host_display(route)))
pass pass
finally: finally:
local_writer.close() local_writer.close()
@ -43,7 +49,7 @@ async def create_proxy_pipe(route: Route):
server = await asyncio.start_server( server = await asyncio.start_server(
partial(handle_client, route), "0.0.0.0", route.listen_port partial(handle_client, route), "0.0.0.0", route.listen_port
) )
print( logger.info(
"Routing from {} to {}".format( "Routing from {} to {}".format(
route.listen_port, destination_host_display(route) route.listen_port, destination_host_display(route)
) )
@ -56,9 +62,12 @@ async def main():
parser.add_argument( parser.add_argument(
"-R", "--route", action="append", required=True, type=parse_argument "-R", "--route", action="append", required=True, type=parse_argument
) )
parser.add_argument("--verbose", action="store_true")
args = parser.parse_args() args = parser.parse_args()
logger.setLevel(logging.DEBUG if args.verbose else logging.INFO)
servers = [create_proxy_pipe(route) for route in args.route] servers = [create_proxy_pipe(route) for route in args.route]
print("Starting servers...") logger.debug("Starting servers...")
await asyncio.gather(*servers) await asyncio.gather(*servers)
@ -66,4 +75,4 @@ if __name__ == "__main__":
try: try:
asyncio.run(main()) asyncio.run(main())
except KeyboardInterrupt: except KeyboardInterrupt:
print("Process terminated") logger.error("Process terminated")