Python + asyncio でポートフォワーディングをしてみる

ポートフォワーディングといっても、ローカルからローカルに流すやつ。

「あるポートに来た TCP 接続をそのまま別のポートに流す」というやつをやりたくて、まぁそれだけなら nc と mkfifo を組み合わせたらできるらしいのだけど、それって Python の asyncio で書けるんじゃねえかな？と思って書いてみたらできた、という話。

port_forward.py

from asyncio import StreamReader, StreamWriter, gather, open_connection, run, start_server
from asyncio.exceptions import CancelledError
from logging import Logger
from typing import NamedTuple


class Addr(NamedTuple):
    host: str
    port: int

    @classmethod
    def parse(cls, addr: str) -> "Addr":
        host, port = addr.split(":")
        return cls(host, int(port))

    def __str__(self):
        return f"{self.host}:{self.port}"


class PortForward:
    def __init__(self, local: Addr, remote: Addr, logger: Logger):
        self.local = local
        self.remote = remote
        self.logger = logger

    async def serve(self, **kwargs):
        server = await start_server(self._handler, *self.local, **kwargs)
        async with server:
            self.logger.info(f"Serving on {self.local}, forwarding to {self.remote}")
            await server.serve_forever()

    async def _handler(self, local_reader: StreamReader, local_writer: StreamWriter):
        try:
            remote_reader, remote_writer = await open_connection(*self.remote)
        except:
            self.logger.exception(f"Failed to connect to remote {self.remote}")
            local_writer.close()
            await local_writer.wait_closed()
            return

        try:
            await gather(pipe(local_reader, remote_writer), pipe(remote_reader, local_writer))
        except CancelledError:
            self.logger.warning("Connection cancelled")
        except:
            self.logger.exception("Connection error")


async def pipe(reader: StreamReader, writer: StreamWriter):
    try:
        while True:
            data = await reader.read(4096)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    finally:
        writer.close()
        await writer.wait_closed()


if __name__ == "__main__":
    import sys
    from logging import basicConfig, getLogger

    basicConfig(level="INFO")

    if len(sys.argv) != 3:
        print(f"Usage:\n  {sys.argv[0]} <local_host>:<local_port> <remote_host>:<remote_port>")
        sys.exit(1)

    try:
        run(PortForward(Addr.parse(sys.argv[1]), Addr.parse(sys.argv[2]), getLogger()).serve())
    except KeyboardInterrupt:
        pass

例えばポート 2000 にきた通信をポート 3000 に流したい場合はこう。

$ python port_forward.py localhost:2000 localhost:3000

※ Python 3.9 以上なら動くっぽいことは確認したが、それより古いバージョンだとわからない。

はじめは Transport や Protocol みたいな asyncio の低レベル API を使わないといけないかと思っていたが、これくらいの用途であれば高レベル API だけでいい感じに書けるのだなぁ、という知見を得た。