Simple Python 3 Secure WebSocket Server (SSL/TLS)

wss_server.py

#!/usr/bin/env python3
# Author: Sean Pesce

# Shell command to create a self-signed TLS certificate and private key:
#    openssl req -new -newkey rsa:4096 -x509 -sha256 -days 365 -nodes -out cert.crt -keyout private.key

import asyncio
import ssl
import sys
import websockets


async def handle_client(wsock, path):
    client_str = f'{wsock.remote_address[0]}:{wsock.remote_address[1]}'
    print(f'[Client connected] {client_str} - {path}')
    
    async for msg in wsock:
        print(f'[Client message] {client_str}\n{msg}')
        resp = 'RESPONSE'
        await wsock.send(resp)
        print(f'[Server response] {client_str}\n{resp}\n')


def serve(host, port, cert_fpath, privkey_fpath):
    """
    Begin listening on the specified interface (host) and port. If file paths are provided for a
    certificate and private key, the server will automatically use TLS (HTTPS/WSS). If not, the server
    will automatically serve plaintext (HTTP).
    """
    ssl_context = None
    if cert_fpath and privkey_fpath:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(certfile=cert_fpath, keyfile=privkey_fpath, password='')

    proto = 'WS'
    if ssl_context:
        proto = 'WSS'
    print(f'Listening for {proto} connections on {host}:{port}\n')
    return websockets.serve(handle_client, host, port, ssl=ssl_context)


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(f'Usage:\n  {sys.argv[0]} <port> [PEM certificate file] [private key file]')
        sys.exit()
    
    PORT = int(sys.argv[1])
    CERT_FPATH = None
    PRIVKEY_FPATH = None
    if len(sys.argv) > 3:
        CERT_FPATH = sys.argv[2]
        PRIVKEY_FPATH = sys.argv[3]
    
    wssd = serve('0.0.0.0', PORT, CERT_FPATH, PRIVKEY_FPATH)
    asyncio.get_event_loop().run_until_complete(wssd)
    asyncio.get_event_loop().run_forever()