certo/get_cert.py

#!/usr/bin/env python3

import hashlib
import logging
import ssl
import socket
import time

import click


logging.basicConfig(level=logging.INFO,
                    format='[%(levelname)-4s] %(message)s',
                    datefmt='%Y-%m-%d %H:%M')
logger = logging.getLogger('certo')


# The following inspired by:
# https://stackoverflow.com/questions/17667903/python-socket-receive-large-amount-of-data
# https://www.binarytides.com/receive-full-data-with-the-recv-socket-function-in-python/
def recv_msg(sock, timeout):
    # Read message length and unpack it into an integer
    data = b''
    begin = time.time()
    n = 1024
    while len(data) < n or time.time() - begin < timeout:
        try:
            data += sock.recv(n - len(data))
            logger.debug("Partial: %r" % data)
        except socket.timeout:
            pass
        if b'Ready to start TLS\r\n' in data:
            break
        time.sleep(0.1)
    if len(data) == 0:
        return None
    logger.debug("Data: %r" % data)
    return data


def establish_conn(addr, port, starttls):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(5)
    try:
        if starttls:
            logger.debug("Using STARTTLS")
            logger.debug("Connecting to %s:%s" % (addr, port))
            sock.connect((addr, port))
            sock.send(b"STARTTLS\r\n")
            data = recv_msg(sock, 5)
            if data is None:
                raise socket.error
            wrapped_socket = ssl.wrap_socket(sock)
        else:
            wrapped_socket = ssl.wrap_socket(sock)
            wrapped_socket.connect((addr, port))

        return wrapped_socket.getpeercert(True)

    finally:
        wrapped_socket.close()


def get_cert(addr, port, starttls):
    cert =  establish_conn(addr, port, starttls)
    pem_cert = ssl.DER_cert_to_PEM_cert(cert)
    logger.debug("The certificate is:\n%s" % pem_cert)

    return cert


def capitalize_and_colons(in_hash):
    in_hash = in_hash.upper()
    new_hash = in_hash[0:2]
    for i in range(2, len(in_hash), 2):
        new_hash += ":" + in_hash[i:i+2]
    return new_hash


def compute_fingerprints(cert, with_colons):
    thumb_md5 = hashlib.md5(cert).hexdigest()
    thumb_sha1 = hashlib.sha1(cert).hexdigest()
    thumb_sha256 = hashlib.sha256(cert).hexdigest()
    logger.info("MD5:    " + thumb_md5)
    if with_colons:
        logger.info("        " + capitalize_and_colons(thumb_md5))
    logger.info("SHA1:   " + thumb_sha1)
    if with_colons:
        logger.info("        " + capitalize_and_colons(thumb_sha1))
    logger.info("SHA256: " + thumb_sha256)
    if with_colons:
        logger.info("        " + capitalize_and_colons(thumb_sha256))


@click.command()
@click.argument('address')#, help="address to be used to retrieve the certificate")
@click.option('-p', '--port', default=443, type=click.IntRange(1,65535), help="The port to connect to.")
@click.option('--starttls', is_flag=True, flag_value=True, help="Whether to use starttls on connection.")
@click.option('--debug/--nodebug', is_flag=True, flag_value=False, help="Debug output.")
@click.option('-o', '--output', help="Path to save the certificate to.")
@click.option('--colons/--nocolons', is_flag=True, flag_value=False, help="Whether to output also hashed with colons")
def doit(address, port, starttls, debug, output, colons):
    if debug:
        logger.setLevel(logging.DEBUG)
    cert = get_cert(address, port, starttls)
    if output:
        with open(output, 'w') as f:
            logger.debug("Opening file %s" % output)
            f.write(ssl.DER_cert_to_PEM_cert(cert))
        logger.info("The certificate has been saved to %s" % output)
    compute_fingerprints(cert, colons)

if __name__ == '__main__':
    doit()