#!/usr/bin/python3
#
#

from __future__ import print_function

import sys
import argparse
import logging

from socket import socket, AF_INET6, AF_INET, create_connection
from ssl import SSLError, CertificateError, SSLContext
from ssl import PROTOCOL_TLSv1_2, CERT_REQUIRED
from unbound import ub_ctx

from check_dane.tlsa import verify_tlsa_record
from check_dane.cert import verify_certificate, add_certificate_options

def init_connection(sslcontext, args, family):
    host = args.Host

    if args.ssl:
        port = 465 if args.port == 0 else args.port
        connection = sslcontext.wrap_socket(socket(family),
                                            server_hostname=host)
        connection.connect((host, port))
        answer = connection.recv(512)
        logging.debug(answer)

        connection.send(b"EHLO localhost\r\n")
        answer = connection.recv(512)
        logging.debug(answer)

    else:
        port = 25 if args.port == 0 else args.port

        connection = socket(family=family)
        connection.connect((host, port))
        answer = connection.recv(512)
        logging.debug(answer)

        connection.send(b"EHLO localhost\r\n")
        answer = connection.recv(512)
        logging.debug(answer)

        connection.send(b"STARTTLS\r\n")
        answer = connection.recv(512)
        logging.debug(answer)

        connection = sslcontext.wrap_socket(connection, server_hostname=host)
        connection.do_handshake()

        connection.send(b"EHLO localhost\r\n")
        answer = connection.recv(512)
        logging.debug(answer)

    return connection


def close_connection(connection):
    connection.send(b"QUIT\r\n")
    answer = connection.recv(512)
    logging.debug(answer)


def init(args):
    sslcontext = SSLContext(PROTOCOL_TLSv1_2)
    sslcontext.verify_mode = CERT_REQUIRED
    sslcontext.load_verify_locations(args.castore)

    resolver = ub_ctx()
    resolver.add_ta_file(args.ancor)

    return sslcontext, resolver


def main():
    logging.basicConfig(format='%(levelname)5s %(message)s')
    parser = argparse.ArgumentParser()
    parser.add_argument("Host")

    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    parser.add_argument("-p", "--port",
                        action="store", type=int, default=0,
                        help="SMTP port")
    parser.add_argument("--ssl",
                        action="store_true",
                        help="Use direct TLS connection instead of starttls (default: disabled)")
    parser.add_argument("--check-dane",
                        action="store_false",
                        help="Verify presented certificate via DANE (default: enabled)")
    parser.add_argument("--check-ca",
                        action="store_false",
                        help="Verify presented certificate via the CA system (default: enabled)")
    parser.add_argument("--check-expire",
                        action="store_false",
                        help="Verify presented certificate for expiration (default: enabled)")

    parser.add_argument("-a", "--ancor",
                        action="store", type=str, default="/etc/unbound/root.key",
                        help="DNSSEC root ancor")
    parser.add_argument("--castore", action="store", type=str,
                        default="/etc/ssl/certs/ca-certificates.crt",
                        help="ca certificate bundle")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("-6", "--6", action="store_true", dest="use6", help="check via IPv6 only")
    group.add_argument("-4", "--4", action="store_true", dest="use4", help="check via IPv4 only")
    group.add_argument("--64", action="store_false", dest="use64", help="check via IPv4 and IPv6 (default)")

    add_certificate_options(parser)

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    elif args.quiet:
        logging.getLogger().setLevel(logging.WARNING)
    else:
        logging.getLogger().setLevel(logging.INFO)

    port = args.port
    if port == 0:
        port = 465 if args.ssl else 25
    host = args.Host.encode('idna').decode()

    sslcontext, resolver = init(args)

    if args.use6:
        afamilies = [AF_INET6]
    elif args.use4:
        afamilies = [AF_INET6]
    else:
        afamilies = [AF_INET, AF_INET6]

    retval = 0
    for afamily in afamilies:
        try:
            connection = init_connection(sslcontext, args, afamily)
        except ConnectionRefusedError:
            logging.error("Connection refused")
            return 2

        nretval = verify_certificate(connection.getpeercert(), args)
        retval = max(retval, nretval)
        nretval = verify_tlsa_record(resolver, "_%d._tcp.%s" % (port, host),
                                     connection.getpeercert(binary_form=True))
        retval = max(retval, nretval)

        close_connection(connection)

    return retval


if __name__ == '__main__':
    sys.exit(main())