#!/usr/bin/env python3

"""A simple Python program to test TLS connections and certificate
checks. The main goal is to show that, by default, PyOpenSSL does not
check that the hostname indicated matches the one in the certificate."""

# Can be changed on the command line
host = None
path = "/"

# Hardwired
PORT = 443

# https://www.pyopenssl.org/
# https://pyopenssl.readthedocs.io/
import OpenSSL
# We could use GNUtls <https://github.com/AGProjects/python-gnutls>
# but, at the end of 2019, it does not work with Python 3
# <https://github.com/AGProjects/python-gnutls/issues/12>

# https://github.com/drkjam/netaddr/
import netaddr

import socket
import getopt
import sys

def error(msg=None):
    if msg is None:
        msg = "Unknown error"
    print(msg,file=sys.stderr)
    sys.exit(1)
    
def usage(msg=None):
    if msg:
        print(msg,file=sys.stderr)
    print("Usage: %s --server hostname" % sys.argv[0], file=sys.stderr)

def canonicalize(hostname):
    result = hostname.lower()
    # TODO handle properly the case where it fails with UnicodeError
    # (two consecutive dots for instance) to get a custom exception
    result = result.encode('idna').decode()
    if result[len(result)-1] == '.':
        result = result[:-1]
    return result

def is_valid_ip_address(addr):
    try:
        baddr = netaddr.IPAddress(addr)
    except netaddr.core.AddrFormatError:
        return (False, None)
    return (True, baddr.version)

def get_certificate_san(x509cert):
    san = ""
    ext_count = x509cert.get_extension_count()
    for i in range(0, ext_count):
        ext = x509cert.get_extension(i)
        if "subjectAltName" in str(ext.get_short_name()):
            san = str(ext)
    return san

# Try one possible name. Names must be already canonicalized.
def match_hostname(hostname, possibleMatch):
    if possibleMatch.startswith("*."): # Wildcard
        base = possibleMatch[1:] # Skip the star
        # RFC 6125 says that we MAY accept left-most labels with
        # wildcards included (foo*bar). We don't do it here.
        try:
            (first, rest) = hostname.split(".", maxsplit=1)
        except ValueError: # One-label name
            rest = hostname
        if rest == base[1:]:
            return True
        if hostname == base[1:]:
            return True
        return False
    else:
        return hostname == possibleMatch

# Try all the names in the certificate
def validate_hostname(hostname, cert):
    # Complete specification is in RFC 6125. It is long and
    # complicated and I'm not sure we do it perfectly.
    (is_addr, family) = is_valid_ip_address(hostname)
    hostname = canonicalize(hostname)
    for alt_name in get_certificate_san(cert).split(", "):
        if alt_name.startswith("DNS:") and not is_addr:
            (start, base) = alt_name.split("DNS:")
            base = canonicalize(base)
            found = match_hostname(hostname, base)
            if found:
                return True
        elif alt_name.startswith("IP Address:") and is_addr:
            host_i = netaddr.IPAddress(hostname)
            (start, base) = alt_name.split("IP Address:")
            if base.endswith("\n"):
                base = base[:-1]
            try:
                base_i = netaddr.IPAddress(base)
            except netaddr.core.AddrFormatError:
                continue # Ignore broken IP addresses in certificates. Are we too liberal?
            if host_i == base_i: 
                return True
        else:
            pass # Ignore unknown alternative name types. 
    # According to RFC 6125, we MUST NOT try the Common Name before the Subject Alternative Names.
    cn = canonicalize(cert.get_subject().commonName)
    found = match_hostname(hostname, cn)
    if found:
        return True
    return False

# Main program
try:
    optlist, args = getopt.getopt (sys.argv[1:], "hs:p:",
                                       ["help", "server=", "path="])
    for option, value in optlist:
        if option == "--help" or option == "-h":
            usage()
            sys.exit(0)
        elif option == "--server" or option == "-s":
            host = value
        elif option == "--path" or option == "-p":
            path = value
        else:
            error("Unknown option %s" % option)
except getopt.error as reason:
    usage(reason)
    sys.exit(1)
if host is None:
    usage("No server given")
    sys.exit(1)    
        
addrinfo = socket.getaddrinfo(host, PORT, 0)

# We should loop over the IP addresses instead of taking only the first one…
sock = socket.socket(addrinfo[0][0], socket.SOCK_STREAM)
addr = addrinfo[0][4]
print("Connecting to %s ..." % str(addr))

context = OpenSSL.SSL.Context(OpenSSL.SSL.TLSv1_2_METHOD)

# Use the OS' default CAs
context.set_default_verify_paths()
context.set_verify_depth(4) # Seems ignored

# This does not check the host name, just the path from the CA, the
# expiration, may be the CRLs…
context.set_verify(OpenSSL.SSL.VERIFY_PEER | OpenSSL.SSL.VERIFY_FAIL_IF_NO_PEER_CERT | \
                   OpenSSL.SSL.VERIFY_CLIENT_ONCE,
                   lambda conn, cert, errno, depth, preverify_ok: preverify_ok)

session = OpenSSL.SSL.Connection(context, sock)
session.set_tlsext_host_name(canonicalize(host).encode()) # Server Name Indication (SNI)

# TCP
session.connect((addr))

# TLS
session.do_handshake()
cert = session.get_peer_certificate()
print("Connected, its certificate is for \"%s\", delivered by \"%s\"" % \
      (cert.get_subject().commonName,
       cert.get_issuer().commonName))
# OpenSSL has check_host but PyOpenSSL does not export it
# <https://github.com/pyca/pyopenssl/issues/795>. So, we do it
# ourselves.
valid = validate_hostname(host, cert)
if not valid:
    print("Certificate error: \"%s\" is not in the certificate" % (host))

request = """
GET %s HTTP/1.1
Host: %s
Connection: close

""" % (path, canonicalize(host))
session.write(request.replace("\n", "\r\n"))

# In a real application, we would loop to get all the data
data = session.read(256)
print(data.decode())

session.close()
sock.close()