Python で WebSocket を使って SSH をトンネリングしてみた

WebSocket が使われるようになった今、ひょっとしてブラウザー通信風にトンネルするプログラムは簡単に書けるんじゃないかと思って、試しに書いてみました。

感想

asyncio に馴染むのに時間がかかり、動かない時はエラーも出ずに固まるので苦労しましたが、asyncio Streams や、Patterns に典型的な実装パターンが紹介されていて、少し楽ができました。

websockets パッケージのおかげで、サーバー、クライアントともに 60 行程度の小さなプログラムになりました。想像以上にコンパクトです。

サーバーで --certfile を指定しなかった場合に動作するかは確認できていません。ssl.create_default_context() ってサーバーでは効かないのかな？

こんな小さなプログラムでもわからないことはあるもので、Windows では ssh からの切断時に read で OSError 64 が発生しました。asyncio Streams にはそんなこと書かれていないんですが、何でしょうね？ とりあえず無視するようにしましたが、Windows で例外を出さずに切断する方法を知りたいです。

軽く SSH が使えることを確認しましたが、それ以上ではないので、まだバグが残っている可能性は大です。

自分で書けば人様のプログラムを動かすよりも安全だと思いますし、書いてみるのは楽しかったです。

使用例

サーバー側

$ python3 -m venv .venv
$ . .venv/bin/activate
$ pip install websockets
$ sudo .venv/bin/python3 server.py --certfile /etc/ssl/chained.crt --keyfile /etc/ssl/private/private.key

サービス化

$ sudo systemctl edit --force --full wstunnel-to-ssh.service

wstunnel-to-ssh.service

[Service]
ExecStart=/path/to/script-dir/.venv/bin/python3 \
            /path/to/script-dir/server.py \
              --certfile /etc/ssl/chained.crt \
              --keyfile /etc/ssl/private/private.key

[Install]
WantedBy=multi-user.target

クライアント側

$ python3 -m venv .venv
$ . .venv/bin/activate
$ pip install websockets
$ python3 client.py -l 10022 wss://example.com 

$ ssh -p 10022 example.com

ソース コード

サーバー側

server.py

server.py

#!/usr/bin/env python3

import asyncio
import ssl
from argparse import ArgumentParser
from websockets.asyncio.server import serve

wss_port = 443
ssh_port = 22
size = 1024

async def ssh_to_ws(reader, websocket):
    while data := await reader.read(size):
        await websocket.send(data)

async def ws_to_ssh(writer, websocket):
    async for message in websocket:
        writer.write(message)
        await writer.drain()
    writer.close()
    await writer.wait_closed()

async def ws_server(websocket):
    reader, writer = await asyncio.open_connection(port=ssh_port)
    remote_address = websocket.remote_address
    print(f"connected from: {remote_address}")
    tasks = []
    tasks.append(asyncio.create_task(ssh_to_ws(reader, websocket)))
    tasks.append(asyncio.create_task(ws_to_ssh(writer, websocket)))
    _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
    for task in pending:
        task.cancel()
    await websocket.close()
    await websocket.wait_closed()
    print(f"connection closed: {remote_address}")

async def main(args):
    if args.certfile:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(args.certfile, keyfile=args.keyfile, password=args.password)
    else:
        ssl_context = ssl.create_default_context()

    server = await serve(ws_server, port=wss_port, ssl=ssl_context)
    addrs = ", ".join(str(sock.getsockname()) for sock in server.sockets)
    print(f"Serving on {addrs}")
    await server.serve_forever()


parser = ArgumentParser()
parser.add_argument("--certfile", metavar="<PEM certificates for the server and the CAs")
parser.add_argument("--keyfile", metavar="<private key file>")
parser.add_argument("--password", metavar="<private key password>")
args = parser.parse_args()

try:
    asyncio.run(main(args))
except KeyboardInterrupt:
    print()

クライアント側

client.py

client.py

#!/usr/bin/env python3

import asyncio
import ssl
from argparse import ArgumentParser
from functools import partial
from websockets.asyncio.client import connect

size = 1024

async def ssh_client_to_ws(reader, websocket):
    try:
        while data := await reader.read(size):
            await websocket.send(data)
    except OSError as e:
        # workaround
        if not e.winerror or e.winerror != 64:
            raise e

async def ws_to_ssh_client(writer, websocket):
    async for message in websocket:
        writer.write(message)
        await writer.drain()
    writer.close()
    await writer.wait_closed()

async def tcp_server(url, cafile, reader, writer):
    ssl_context = ssl.create_default_context(cafile=cafile)
    async with connect(url, ssl=ssl_context) as websocket:
        print(f"connected: ${websocket.remote_address}")
        tasks = []
        tasks.append(asyncio.create_task(ssh_client_to_ws(reader, websocket)))
        tasks.append(asyncio.create_task(ws_to_ssh_client(writer, websocket)))
        _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for task in pending:
            task.cancel()
        await websocket.close()
        await websocket.wait_closed()
        print("connection closed")

async def main(args):
    server = await asyncio.start_server(
        partial(tcp_server, args.url, args.cafile),
        port=args.l
    )

    addrs = ", ".join(str(sock.getsockname()) for sock in server.sockets)
    print(f"Serving on {addrs}")

    async with server:
        await server.serve_forever()


parser = ArgumentParser()
parser.add_argument("-l", type=int, default=22, metavar="<listening port>")
parser.add_argument("--cafile", metavar="<CA cert file>")
parser.add_argument("url", help="a URL like wss://example.com")
args = parser.parse_args()

try:
    asyncio.run(main(args))
except KeyboardInterrupt:
    print()