TLS server and client demo

tls-client.py

#!/usr/bin/env python3

import argparse, contextlib, socket, ssl

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', '-o', default='localhost')
    parser.add_argument('--port', '-p', type=int, default=4430)
    args = parser.parse_args()

    context = ssl.create_default_context()
    # This is necessary for self-signed server certificates
    context.load_verify_locations('server.crt')

    with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as client:
        client.connect((socket.gethostbyname(args.host), args.port))
        with contextlib.closing(context.wrap_socket(client, server_hostname=args.host)) as client2:
            client2.sendall(b'hello there\n')
            print(client2.recv())

tls-server.py

#!/usr/bin/env python3

import argparse, contextlib, logging, socket, ssl

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s:%(levelname)s:%(message)s', level=logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument('--port', '-p', type=int, default=4430)
    args = parser.parse_args()

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('', args.port))
    server.listen()

    context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    # openssl req -x509 -nodes -newkey rsa:4096 -keyout server.key -out server.crt -subj '/C=US/O=Acme/CN=localhost'
    context.load_cert_chain('server.crt', 'server.key')

    while True:
        client, addr = server.accept()
        with contextlib.closing(client) as client:
            logging.info('connection from {}:{}'.format(*addr))
            with contextlib.closing(context.wrap_socket(client, server_side=True)) as client2:
                client2.send(client2.recv())