#!/usr/bin/env python

import optparse
import os
import socket
import ssl
import sys

SERVER = ("127.0.0.1", 39898)
MAX_MSG_SIZE = 1024


def do_ssl_server(key_file, cert_file, ca_file):
    try:
        # Start listening on some port
        lis_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        lis_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        lis_sock.bind(SERVER)
        lis_sock.listen(socket.SOMAXCONN)
        print "Listening: " + str(SERVER)
        # Accept an SSL connection
        net_sock, cli_addr = lis_sock.accept()
        print "New TCP connection with %s:%d" % cli_addr
        ssl_sock = ssl.wrap_socket(net_sock,
                                   server_side=True,
                                   keyfile=key_file,
                                   certfile=cert_file,
                                   ca_certs=ca_file,
                                   cert_reqs=ssl.CERT_REQUIRED)
        print "Successful SSL connection with %s:%d" % cli_addr
        # Read some data
        data = ssl_sock.read(MAX_MSG_SIZE)
        print "Data: " + str(data)
        # Close sockets
        ssl_sock.close()
        lis_sock.close()
        print "Success"
    except socket.error as err:
        print "Socket Error: " + str(err)
        sys.exit(1)


def do_ssl_client(key_file, cert_file, ca_file):
    try:
        tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        ssl_sock = ssl.wrap_socket(tcp_sock,
                server_side=False,
                keyfile=key_file,
                certfile=cert_file,
                ca_certs=ca_file,
                cert_reqs=ssl.CERT_REQUIRED)
        ssl_sock.connect(SERVER)
        ssl_sock.send("Hello, World!\n")
        ssl_sock.close()
        print "Success"
    except socket.error as err:
        print "Socket Error: " + str(err)
        exit(1)


def usage():
    print os.path.basename(__file__) + " <client|server> " + \
            "-k <key_file> -c <cert_file> -x <ca_file>\n"


if __name__ == "__main__":
    parser = optparse.OptionParser()
    parser.add_option("-k", "--key_file", dest="key_file", help="Key File")
    parser.add_option("-c", "--cert_file", dest="cert_file", help="Cert File")
    parser.add_option("-x", "--ca_file", dest="ca_file", help="CA File")
    cmd_ops, args = parser.parse_args()
    # Sanity checks
    if len(args) != 1:
        usage()
        print "Invalid number of commands"
        sys.exit(1)
    if cmd_ops.key_file is None:
        usage()
        print "Key file is missing"
        sys.exit(1)
    if cmd_ops.cert_file is None:
        usage()
        print "Cert file is missing"
        sys.exit(1)
    if cmd_ops.ca_file is None:
        usage()
        print "CA file is missing"
        sys.exit(1)
    # Now do either server or client action
    if args[0] == "server":
        do_ssl_server(key_file=cmd_ops.key_file,
                cert_file=cmd_ops.cert_file,
                ca_file=cmd_ops.ca_file)
    elif args[0] == "client":
        do_ssl_client(key_file=cmd_ops.key_file,
                cert_file=cmd_ops.cert_file,
                ca_file=cmd_ops.ca_file)
    else:
        usage()
        print "Invalid command: " + str(args[0])
        sys.exit(1)
    # Done
    sys.exit(0)