#!/usr/bin/python
# -*- coding: UTF-8 -*-
# noinspection SpellCheckingInspection
"""
@since: 10/14/2020
@authors: Michael Mendenhall
@version: 1.6
@description: Tool to assist with creating/uploading SSL certificates to ECS
@change:
12/12/2020     1.0 - Fix outputting when password not configured in config.ini during view_certs operation

02/12/2021     1.1 - Support different hostnames for data/management interfaces #3
                   - Rewrote view_certs so it works if no certs have been uploaded yet. #2
                   - backup original certifiate before uploading new one. #1

04/07/2021     1.2 - nuke certs #10
                   - fix urllib3 warnings
                   - fix logging
                   - output additional info when viewing certs #9
07/06/2021     1.3 - Support 1024/2048/4096 private key sizes #14

09/24/2021     1.4 - #18 - Fix bug in get_issuer
                   - #19 - Remove sudo requirement and force admin user
                   - #23 - Handle Credentials with ?{}|&~![()^"

10/08/2021     1.5 - #25 - admin userid 1001

               1.6 - Set log file permissions to 0755 and chown os.geteuid/os.getegid
                   - fix userid check
"""

import sys
import os
import argparse
import logging
import time
import json
import atexit
import datetime
import traceback
import warnings
import subprocess
import getpass
import pickle

sys.path.insert(0, '{}/libs'.format(os.path.dirname(os.path.abspath(__file__))))
with warnings.catch_warnings():
    warnings.simplefilter('ignore')
    from cryptography.hazmat.primitives import serialization, hashes
    from cryptography.hazmat.primitives.asymmetric import rsa, ec

import urllib3

urllib3.disable_warnings()
import requests
# noinspection PyCompatibility
import configparser
import dns.reversename
import dns.resolver
import idna
from cryptography import x509
from pprint import pformat
# noinspection PyCompatibility
from ipaddress import IPv4Address
from dns.exception import DNSException
from OpenSSL import SSL
from socket import socket
from logging.handlers import RotatingFileHandler

__version__ = '1.6'


class EcsCertificateTool:
    def __init__(self, args):
        self.script_version = __version__
        self.script_name = 'ecs_certificate_tool'
        self.args = args
        self.time_string = time.strftime("%Y-%m-%d-%H-%M-%S")
        # file paths
        self.psnt_path = '/opt/emc/caspian/fabric/agent/data/topology.json'
        self.serial = self._get_psnt()
        self.script_dir = os.path.dirname(os.path.abspath(__file__))
        # backup paths
        self.cert_backup_dir = os.path.join(self.script_dir, 'certificate_backups')
        self.management_backup_path = os.path.join(self.cert_backup_dir,
                                                   '{}-management_{}.crt.backup'.format(self.serial, self.time_string))
        self.data_backup_path = os.path.join(self.cert_backup_dir,
                                             '{}-data_{}.crt.backup'.format(self.serial, self.time_string))
        self.data_backup_path = '{}/{}-data_{}.crt.backup'.format(self.cert_backup_dir,
                                                                  self.serial, self.time_string)
        # Generated file paths
        self.generated_files_dir = os.path.join(self.script_dir, 'generated_files')
        self.data_private_key_path = os.path.join(self.generated_files_dir, '{}-data_private.key'.format(self.serial))
        self.management_private_key_path = os.path.join(self.generated_files_dir, '{}-management_private.key'.format(self.serial))
        self.data_csr_path = os.path.join(self.generated_files_dir, '{}-data.csr'.format(self.serial))
        self.management_csr_path = os.path.join(self.generated_files_dir, '{}-management.csr'.format(self.serial))
        self.data_ssc_path = os.path.join(self.generated_files_dir, '{}-data-ssc.crt'.format(self.serial))
        self.management_ssc_path = os.path.join(self.generated_files_dir, '{}-management-ssc.crt'.format(self.serial))
        self.object_main_json_path = '/opt/emc/caspian/fabric/agent/data/object-main_network.json'
        self.log_dir = os.path.join(self.script_dir, 'logs')
        if not os.path.exists(self.log_dir):
            os.mkdir(self.log_dir)
        self.logfile = os.path.join(self.log_dir, 'certificate_tool.log')

        # Config stuff
        self.config_file_path = os.path.join(self.script_dir, 'config.ini')
        self.cred_cache_path = os.path.join(self.script_dir, '.0')
        self.config = configparser.ConfigParser(allow_no_value=True)
        self.config.optionxform = str
        self.GENERAL_CONF = 'GENERAL'
        self.SELF_SIGNED_CONF = 'SELF_SIGNED'
        self.DATA_SAN_CONF = 'DATA_SUBJECT_ALTERNATIVE_NAME'
        self.MANAGEMENT_SAN_CONF = 'MANAGEMENT_SUBJECT_ALTERNATIVE_NAME'
        self.ADVANCED_CONF = 'ADVANCED'
        self.TARGET_SAN_CONF = None
        self.DATA = 'DATA'
        self.MANAGEMENT = 'MANAGEMENT'
        # Net map
        self.network_info = {'data': [],
                             'management': [],
                             'management_fqdns': [],
                             'data_fqdns': [],
                             'private.4': [],
                             '9206': [],
                             '2181': []}
        #
        self.private_key = None
        # REST
        self.management_user = 'root'
        self.management_pass = 'ChangeMe'
        self.management_ip = 'localhost'
        self.management_port = 4443
        self.data_port = 9021
        self.AUTH_TOKEN = None
        # args
        if self.args.get('data'):
            self.TARGET_SAN_CONF = self.DATA_SAN_CONF
        else:
            self.TARGET_SAN_CONF = self.MANAGEMENT_SAN_CONF
        self.debug = self.args['debug']
        self.private_key_size = self.args.get('key_size')

    def _init_logging(self):
        handler = RotatingFileHandler(filename=self.logfile,
                                      maxBytes=3000000,
                                      mode='a',
                                      backupCount=5)
        rootLogger = logging.getLogger()
        log_formatter = logging.Formatter(
            '%(asctime)-15s [%(levelname)7s] [%(process)6s] [%(funcName)16s (%(threadName)s)] %(message).20000s')
        handler.setFormatter(log_formatter)
        rootLogger.addHandler(handler)
        consoleHandler = logging.StreamHandler(sys.stdout)
        if self.debug:
            consoleHandler.setFormatter(log_formatter)
            rootLogger.addHandler(consoleHandler)
            logging.getLogger().setLevel(logging.DEBUG)
        else:
            rootLogger.addHandler(consoleHandler)
            logging.getLogger().setLevel(logging.WARNING)

        try:
            os.chmod(self.logfile, 0755)
            os.chown(self.logfile, os.geteuid(), os.getegid())
        except Exception:
            pass  # TODO - need more handling?  Will fail if some other user owns the file

    def _logout_token(self):
        if self.AUTH_TOKEN:
            logging.debug('logging out auth token')
            url = 'https://{}:4443/logout'.format(self.management_ip)
            headers = {'X-SDS-AUTH-TOKEN': self.AUTH_TOKEN}
            response = requests.get(url=url, auth=(self.management_user, self.management_pass), verify=False,
                                    headers=headers)
            logging.debug('log out auth token response_code: {}'.format(response.status_code))

    def _get_network_info(self):
        logging.debug('Reading File: {}'.format(self.object_main_json_path))
        with open(self.object_main_json_path, 'r') as f:
            data = json.loads(f.read())
        for node in data['cluster_info']:
            data_ip = str(node['network']['data_ip'])
            data2_ip = str(node['network']['data2_ip'])
            management_ip = str(node['network']['mgmt_ip'])
            private_ip = str(node['network']['private_ip'])
            self.network_info['data'].append(data_ip)
            self.network_info['data'].append(data2_ip)
            self.network_info['management'].append(management_ip)
            self.network_info['private.4'].append(private_ip)

        self.network_info['data'] = list(set(self.network_info['data']))
        for key, value in self.network_info.iteritems():
            self.network_info[key].sort()
            logging.debug('{} : {}'.format(key, value))

    def _discover_listeners(self):
        for ip in self.network_info['private.4']:
            for port in ['2181', '9206']:
                logging.debug('checking for listener on node {} port {}'.format(ip, port))
                listener = self.get_listener_ip(private_4=ip, port=port)
                # TODO
                if listener is not None:
                    logging.debug('found listener on node {} for port {}'.format(ip, port))
                    self.network_info[port].append(listener)
                else:
                    logging.debug('no listener found for port {} on node {}'.format(port, ip))

    def _get_psnt(self):
        psnt = 'Unknown'
        try:
            with open(self.psnt_path, 'r') as f:
                data = json.loads(f.read())
                psnt = data['psnt']
                return psnt
        except:
            logging.warning('warning', exc_info=True)
            pass
        return psnt

    def _backup_certificate(self, interface):
        old_cert = None
        port = self.data_port if interface == 'data' else self.management_port
        backup_path = self.data_backup_path if interface == 'data' else self.management_backup_path
        for ip in self.network_info[interface]:
            try:
                old_cert = self.get_certificate(hostname=ip, port=port)
                break
            except:
                pass
        if old_cert:
            issuer = self.get_issuer(old_cert)
            if str(issuer) != 'localhost':
                if not os.path.exists(self.cert_backup_dir):
                    os.mkdir(self.cert_backup_dir)
                with open(backup_path, 'wb') as f:
                    f.write(old_cert.public_bytes(serialization.Encoding.PEM))
                self.msg_output('Backed up existing certificate to : {}'.format(backup_path))
            else:
                self.msg_output('Skipping existing certificate backup as certificate is issuer is localhost..')

    @staticmethod
    def msg_output(msg='', newline=True):
        logging.info(msg)
        if msg in ['PASS', 'DONE']:
            msg = '\033[92m' + msg + '\033[0m'
        if newline:
            sys.stdout.write(msg + '\n')
        else:
            sys.stdout.write(msg)

    @staticmethod
    def get_listener_ip(private_4, port):
        logging.info('Getting listening IP for node {}, port {}'.format(private_4, port))
        cmd = 'sudo ss -lstn | grep :%s  | grep -oP "(?:[0-9]{1,3}\.){3}[0-9]{1,3}" | head -1' % port
        cmd = 'ssh -n -q -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o ConnectTimeout=10 admin@{} \'{}\''.format(
            private_4, cmd)
        listener = ''
        try:
            output = subprocess.check_output(cmd, shell=True)
            listener = output.strip()
            if output:
                logging.debug('Found listener {} for port {}'.format(listener, port))
        except subprocess.CalledProcessError:
            logging.debug('cmd: {} hit exception:'.format(cmd), exc_info=True)

        return listener

    @staticmethod
    def get_hostname(ip):
        logging.debug('Getting hostname for IPv4Address: {}'.format(ip))
        try:
            name = dns.resolver.query(dns.reversename.from_address(ip), 'PTR')
            name = str(name[0])[:-1]
            logging.debug('IPv4Address: {}, hostname: {}'.format(ip, name))
            return name
        except DNSException:
            logging.warning('Failed to resolve hostname for IPv4Address: {}'.format(ip), exc_info=True)
        except:
            logging.warning('Unknown exception while resolving hostname for IPv4Address: {}'.format(ip), exc_info=True)

    def get_auth_token(self):
        self.msg_output('Authenticating using configured credentials..', newline=False)
        logging.info('Logging in using configured credentials')
        logging.debug('management_ips: {}'.format(self.network_info['management']))
        for ip in self.network_info['management']:
            try:
                url = 'https://{}:4443/login'.format(str(ip))
                response = requests.get(url, auth=(self.management_user, self.management_pass), verify=False)
                logging.debug(response.status_code)
                if response.status_code == 200:
                    token = response.headers['X-SDS-AUTH-TOKEN']
                    logging.debug('token: {}'.format(token))
                    self.AUTH_TOKEN = token
                    self.management_ip = ip
                    atexit.register(self._logout_token)
                    break
                else:
                    logging.debug('{} failed with status_code: {}'.format(url, response.status_code))
            except:
                logging.warning('Exception: ', exc_info=True)
                continue
        if not self.AUTH_TOKEN:
            raise Exception('Rest API authentication failed, please run configure_credentials')

        self.msg_output('PASS')
        self.msg_output()

    def create_rsa_private_key(self):
        self.msg_output('Creating RSA private key..', newline=False)
        # noinspection PyArgumentList
        key = rsa.generate_private_key(public_exponent=65537, key_size=self.private_key_size)
        self.private_key = key
        # noinspection PyTypeChecker
        key_pem = key.private_bytes(encoding=serialization.Encoding.PEM,
                                    format=serialization.PrivateFormat.TraditionalOpenSSL,
                                    encryption_algorithm=serialization.NoEncryption())
        logging.debug(key_pem)
        self.msg_output('DONE')
        logging.debug('Finished creating private key.')
        if self.args.get('data'):
            with open(self.data_private_key_path, 'wb') as f:
                f.write(key_pem)
            self.msg_output('Wrote private key to {}'.format(self.data_private_key_path))

        else:
            with open(self.management_private_key_path, 'wb') as f:
                f.write(key_pem)
            self.msg_output('Wrote private key to {}'.format(self.management_private_key_path))

    def create_elliptic_curve_private_key(self):
        self.msg_output('Creating Elliptic Curve private key..', newline=False)
        # noinspection PyArgumentList
        key = ec.generate_private_key(curve=ec.SECP521R1())
        self.private_key = key
        # noinspection PyTypeChecker
        key_pem = key.private_bytes(encoding=serialization.Encoding.PEM,
                                    format=serialization.PrivateFormat.TraditionalOpenSSL,
                                    encryption_algorithm=serialization.NoEncryption())
        logging.debug(key_pem)
        self.msg_output('DONE')
        logging.debug('Finished creating private key.')
        if self.args.get('data'):
            with open(self.data_private_key_path, 'wb') as f:
                f.write(key_pem)
            self.msg_output('Wrote private key to {}'.format(self.data_private_key_path))

        else:
            with open(self.management_private_key_path, 'wb') as f:
                f.write(key_pem)
            self.msg_output('Wrote private key to {}'.format(self.management_private_key_path))

    @staticmethod
    def get_certificate(hostname, port):
        hostname_idna = idna.encode(hostname)
        sock = socket()

        sock.connect((hostname, port))
        ctx = SSL.Context(SSL.SSLv23_METHOD)  # most compatible
        ctx.check_hostname = False
        ctx.verify_mode = SSL.VERIFY_NONE

        sock_ssl = SSL.Connection(ctx, sock)
        sock_ssl.set_connect_state()
        sock_ssl.set_tlsext_host_name(hostname_idna)
        sock_ssl.do_handshake()
        cert = sock_ssl.get_peer_certificate()
        crypto_cert = cert.to_cryptography()
        sock_ssl.close()
        sock.close()
        return crypto_cert

    @staticmethod
    def get_san_names(cert):
        try:
            ext = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName)
            return ext.value.get_values_for_type(x509.DNSName)
        except x509.ExtensionNotFound:
            return None

    @staticmethod
    def get_common_name(cert):
        try:
            names = cert.subject.get_attributes_for_oid(x509.NameOID.COMMON_NAME)
            return names[0].value
        except x509.AttributeNotFound:
            return None

    @staticmethod
    def get_issuer(cert):
        try:
            return cert.issuer
        except:
            return None

    @staticmethod
    def get_valid_dates(cert):
        try:
            not_valid_before = cert.not_valid_before
            not_valid_after = cert.not_valid_after
            return not_valid_before, not_valid_after
        except:
            return None, None

    def _validate_config(self):
        logging.info('Validating config...')
        self.msg_output('-' * 70)
        self.msg_output('Validating {} configuration'.format(self.GENERAL_CONF))
        self.msg_output('-' * 70)
        self.msg_output()
        general_items = self.config.items(self.GENERAL_CONF)
        for key, value in general_items:
            if not value:
                continue
            self.msg_output('Validating {} = {}..'.format(key, value), newline=False)
            x509_OID = 'x509.NameOID.{}'.format(key)
            x509.NameAttribute(eval(x509_OID), value)
            logging.debug('Successfully validated {}!'.format(key))
            self.msg_output('PASS')

        dns_names = self.config.get(self.TARGET_SAN_CONF, 'DNS_NAMES')
        ip_addresses = self.config.get(self.TARGET_SAN_CONF, 'IP_ADDRESSES')
        dns_names = [x.strip() for x in dns_names.split()]
        ip_addresses = [x.strip() for x in ip_addresses.split()]
        logging.debug('dns_names: {}'.format(dns_names))
        logging.debug('ip_addresses: {}'.format(ip_addresses))
        if dns_names:
            self.msg_output('-' * 70)
            self.msg_output('Validating DNS_NAMES configuration')
            self.msg_output('-' * 70)
            self.msg_output()
            for dns_name in dns_names:
                self.msg_output('Validating DNSName: {}..'.format(dns_name), newline=False)
                x509.DNSName(dns_name)
                # socket.gethostbyname(dns_name)
                self.msg_output('PASS')
            self.msg_output()
            logging.info('Successfully validated DNS_NAMES')

        if ip_addresses:
            self.msg_output('-' * 70)
            self.msg_output('Validating IP_ADDRESSES configuration')
            self.msg_output('-' * 70)
            self.msg_output()
            for ip in ip_addresses:
                self.msg_output('Validating IPv4Address: {}..'.format(ip), newline=False)
                ip = IPv4Address(unicode(ip))
                x509.IPAddress(ip)
                self.msg_output('PASS')
                logging.debug('Successfully validated IP: {}'.format(ip))
            self.msg_output()
            logging.info('Successfully validated IP_ADDRESSES')

        valid_days = self.config.get(self.SELF_SIGNED_CONF, 'VALID_DAYS', fallback=None)
        if valid_days:
            self.msg_output('Validating {}..'.format(self.SELF_SIGNED_CONF), newline=False)
            try:
                int(valid_days)
            except:
                raise ValueError('Valid days must be integer')
            self.msg_output('PASS')
            self.msg_output()
        # advanced
        advanced_items = self.config.items(self.ADVANCED_CONF)
        advanced_items = [x for x in advanced_items if x[1]]
        if advanced_items:
            self.msg_output('-' * 70)
            self.msg_output('Validating {} configuration'.format(self.ADVANCED_CONF))
            self.msg_output('-' * 70)
            for key, value in advanced_items:
                if not value:
                    continue
                self.msg_output('Validating {} = {}..'.format(key, value), newline=False)
                x509_OID = 'x509.NameOID.{}'.format(key)
                x509.NameAttribute(eval(x509_OID), value)
                logging.debug('Successfully validated {}!'.format(key))
                self.msg_output('PASS')
            self.msg_output()
        self.msg_output('All configurations items validated successfully!')
        self.msg_output()

    def _enable_vnest_diagnostics(self, ip):
        url = 'http://{}:9101/nest/diagnosticEnabler?enabled=true'.format(ip)
        response = requests.put(url)
        if response.status_code != 200:
            raise Exception('Failed to enable vnest diagnostics for node: {}'.format(ip))
        atexit.register(self._disable_vnest_diagnostics, ip)

    @staticmethod
    def _disable_vnest_diagnostics(ip):
        url = 'http://{}:9101/nest/diagnosticEnabler?enabled=false'.format(ip)
        requests.put(url)
        # if not response.ok:
        #    raise Exception('Failed to disable vnest diagnostics for {}'.format(ip))

    @staticmethod
    def _verify_diagnostics_enabled(ip):
        url = 'http://{}:9206/configuration/membership'.format(ip)
        response = requests.get(url)
        if response.status_code != 200:
            raise Exception('Failed to validate vnest diagnostics API is enabled for node: {}'.format(ip))

    def _key_pair_vnest_request(self, ip, request_type='GET'):
        if self.args.get('management'):
            url = 'http://{}:9206/d/?path=/config/ecs_keyCertificatePairConfig/ecs_keyCertificatePairId'.format(ip)
        elif self.args.get('data'):
            url = 'http://{}:9206/d/?path=/config/ecs_objectSvcKeyCertificatePairConfig/ecs_keyCertificatePairId'.format(
                ip)
        else:
            raise RuntimeError('management or data interface should be specified')
        if request_type == 'GET':
            url += '&scope=local&showvalue=gpb'
            response = requests.get(url)
            return response
        elif request_type == 'DELETE':
            url += '&scope=paxos'
            response = requests.delete(url)
            return response

    def do_nuke_cert(self):
        if self.args.get('data'):
            target = 'data'
        else:
            target = 'management'
        self._discover_listeners()
        ip_map = zip(self.network_info['private.4'], self.network_info['2181'], self.network_info['9206'])
        for t in ip_map:
            try:
                private_4_ip = t[0]
                object_zk_ip = t[1]
                vnest_diag_ip = t[2]
                if any([private_4_ip, object_zk_ip, vnest_diag_ip]) == '' or any(
                        [private_4_ip, object_zk_ip, vnest_diag_ip]) is None:
                    continue
                self.msg_output('Enabling vnest diagnostics for {}..'.format(private_4_ip), newline=False)
                self._enable_vnest_diagnostics(private_4_ip)
                self.msg_output('PASS')
                self.msg_output('Verifying vnest diagnostics is enabled for {}..'.format(private_4_ip), newline=False)
                self._verify_diagnostics_enabled(vnest_diag_ip)
                self.msg_output('PASS')
                self.msg_output('Backing up existing {} certificate..'.format(target))
                self._backup_certificate(interface=target)
                self.msg_output('Nuking {} certificate..'.format(target), newline=False)

                response = self._key_pair_vnest_request(vnest_diag_ip, request_type='GET')
                exists = True
                if response.status_code == 404:
                    exists = False
                if 'DELETED' in response.text:
                    exists = False
                if not exists:
                    self.msg_output('No certificates exist in vnest, certificate is already cleared')
                    return

                response = self._key_pair_vnest_request(vnest_diag_ip, request_type='DELETE')
                if response.status_code not in [404, 500, 200]:
                    raise Exception('Request to delete certificate from vnest failed.')

                self.msg_output('Verifying certificate has been deleted..')
                response = self._key_pair_vnest_request(vnest_diag_ip, request_type='GET')
                deleted = False
                if response.status_code == 404:
                    deleted = True
                if 'DELETED' in response.text:
                    deleted = True
                if not deleted:
                    raise Exception('Certificate was not deleted from vnest.')

                if target == 'data':
                    self.msg_output('Certificates successfully cleared. Please restart dataheadsvc')
                elif target == 'management':
                    self.msg_output('Certificates successfully cleared. Please restart objcontrolsvc and nginx')
                self.msg_output()
                return
            except:
                logging.debug('exception', exc_info=True)
                continue

    # noinspection PyTypeChecker
    # noinspection PyArgumentList
    def generate_csr(self):
        self.msg_output('-' * 70)
        self.msg_output('Certificate Signing Request')
        self.msg_output('-' * 70)
        self.msg_output()
        self.msg_output('Creating Certificate Signing Request..', newline=False)
        # CSR builder
        builder = x509.CertificateSigningRequestBuilder()
        # Set name attributes
        subject_list = []
        if self.config.items(self.ADVANCED_CONF):
            subject_list.extend([x509.NameAttribute(eval('x509.NameOID.{}'.format(key)), value) for key, value in
                                 self.config.items(self.ADVANCED_CONF) if value])
        subject_list.extend([x509.NameAttribute(eval('x509.NameOID.{}'.format(key)), value) for key, value in
                             self.config.items(self.GENERAL_CONF) if value])
        subject = x509.Name(subject_list)
        logging.debug('{}'.format(subject))
        builder = builder.subject_name(subject)
        dns_names = [x509.DNSName(x.strip()) for x in self.config.get(self.TARGET_SAN_CONF, 'DNS_NAMES').split()]
        ip_addresses = [x509.IPAddress(IPv4Address(unicode(ip.strip()))) for ip in
                        self.config.get(self.TARGET_SAN_CONF, 'IP_ADDRESSES').split()]
        logging.debug('{}'.format(dns_names))
        logging.debug('{}'.format(ip_addresses))
        # Set SAN extension
        san_list = []
        san_list.extend(dns_names)
        san_list.extend(ip_addresses)
        builder = builder.add_extension(
            x509.SubjectAlternativeName(san_list),
            critical=False,

        )
        # Set basic constraints
        builder = builder.add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        # Set key usage
        builder = builder.add_extension(x509.KeyUsage(digital_signature=True,
                                                      content_commitment=True,  # nonRepudiation
                                                      key_encipherment=True,
                                                      data_encipherment=False,
                                                      key_agreement=False,
                                                      key_cert_sign=False,
                                                      crl_sign=False,
                                                      encipher_only=False,
                                                      decipher_only=False), critical=True)
        # Set server auth extension
        builder = builder.add_extension(x509.ExtendedKeyUsage([x509.OID_SERVER_AUTH]), critical=True)
        # Set authorityKeyIdentifier
        # builder = builder.add_extension(x509.AuthorityKeyIdentifier(key_identifier=None,
        #                                                             authority_cert_issuer=None,
        #                                                             authority_cert_serial_number=None), critical=False)

        # Sign the CSR with our private key.
        CSR = builder.sign(private_key=self.private_key, algorithm=hashes.SHA256())
        self.msg_output('DONE')
        # Write to file
        if self.args.get('data'):
            with open(self.data_csr_path, 'wb') as f:
                f.write(CSR.public_bytes(serialization.Encoding.PEM))
            self.msg_output('Wrote certificate signing request to {}'.format(self.data_csr_path))
            output = subprocess.check_output('sudo openssl req -in {} -noout -text'.format(self.data_csr_path),
                                             shell=True)
            logging.debug('{}'.format(output))
        else:
            with open(self.management_csr_path, 'wb') as f:
                f.write(CSR.public_bytes(serialization.Encoding.PEM))
            self.msg_output('Wrote certificate signing request to {}'.format(self.management_csr_path))
            output = subprocess.check_output('sudo openssl req -in {} -noout -text'.format(self.management_csr_path),
                                             shell=True)
            logging.debug('{}'.format(output))

    # noinspection PyTypeChecker
    # noinspection PyArgumentList
    def generate_ssc(self):
        self.msg_output('-' * 70)
        self.msg_output('Self-signed certificate')
        self.msg_output('-' * 70)
        self.msg_output()
        self.msg_output('Creating self-signed certificate..', newline=False)
        valid_days = self.config.get(self.SELF_SIGNED_CONF, 'VALID_DAYS')
        dns_names = [x509.DNSName(x.strip()) for x in self.config.get(self.TARGET_SAN_CONF, 'DNS_NAMES').split()]
        ip_addresses = [x509.IPAddress(IPv4Address(unicode(ip.strip()))) for ip in
                        self.config.get(self.TARGET_SAN_CONF, 'IP_ADDRESSES').split()]
        logging.debug('{}'.format(dns_names))
        logging.debug('{}'.format(ip_addresses))
        logging.debug('Valid days: {}'.format(valid_days))
        # Set SAN extension
        san_list = []
        san_list.extend(dns_names)
        san_list.extend(ip_addresses)
        subject_list = []
        if self.config.items(self.ADVANCED_CONF):
            subject_list.extend([x509.NameAttribute(eval('x509.NameOID.{}'.format(key)), value) for key, value in
                                 self.config.items(self.ADVANCED_CONF) if value])
        subject_list.extend([x509.NameAttribute(eval('x509.NameOID.{}'.format(key)), value) for key, value in
                             self.config.items(self.GENERAL_CONF) if value])
        subject = issuer = x509.Name(subject_list)
        builder = x509.CertificateBuilder()
        builder = builder.subject_name(subject).issuer_name(issuer).public_key(self.private_key.public_key())
        builder = builder.serial_number(x509.random_serial_number())
        builder = builder.not_valid_before(datetime.datetime.utcnow())
        builder = builder.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=int(valid_days)))

        builder = builder.add_extension(
            x509.SubjectAlternativeName(san_list),
            critical=False,

        )
        # Set basic constraints
        builder = builder.add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        # Set key usage
        builder = builder.add_extension(x509.KeyUsage(digital_signature=True,
                                                      content_commitment=True,  # nonRepudiation
                                                      key_encipherment=True,
                                                      data_encipherment=False,
                                                      key_agreement=False,
                                                      key_cert_sign=False,
                                                      crl_sign=False,
                                                      encipher_only=False,
                                                      decipher_only=False), critical=True)
        # Set server auth extension
        builder = builder.add_extension(x509.ExtendedKeyUsage([x509.OID_SERVER_AUTH]), critical=True)
        # Set authorityKeyIdentifier
        builder = builder.add_extension(x509.AuthorityKeyIdentifier(key_identifier=None,
                                                                    authority_cert_issuer=None,
                                                                    authority_cert_serial_number=None), critical=False)

        # Sign the CSR with our private key.
        CERT = builder.sign(private_key=self.private_key, algorithm=hashes.SHA256())
        self.msg_output('DONE')

        if self.args.get('data'):
            with open(self.data_ssc_path, "wb") as f:
                f.write(CERT.public_bytes(serialization.Encoding.PEM))
            self.msg_output('Wrote Certificate to: {}'.format(self.data_ssc_path))
            output = subprocess.check_output('sudo openssl x509 -in {} -text -noout'.format(self.data_ssc_path),
                                             shell=True)
            logging.debug('{}'.format(output))
        else:
            with open(self.management_ssc_path, "wb") as f:
                f.write(CERT.public_bytes(serialization.Encoding.PEM))
            self.msg_output('Wrote Certificate to: {}'.format(self.management_ssc_path))
            output = subprocess.check_output('sudo openssl x509 -in {} -text -noout'.format(self.management_ssc_path),
                                             shell=True)
            logging.debug('{}'.format(output))

    # noinspection PyTypeChecker
    # noinspection PyArgumentList
    def upload_certificate(self, certificate_path, private_key_path, interface):
        self.msg_output('-' * 70)
        self.msg_output('Upload Certificate')
        self.msg_output('-' * 70)
        self.msg_output()
        # Read cert / key
        self.get_auth_token()
        url = None
        self.msg_output('Reading certificate from: {}..'.format(certificate_path), newline=False)
        with open(certificate_path, 'rb') as f:
            data = f.read()
            # certificate_object = x509.load_pem_x509_certificate(data)
            certificate_string = data.encode('utf-8')
            logging.debug('certificate_string:\n{}'.format(certificate_string))
        self.msg_output('DONE')
        self.msg_output('Reading private key from: {}..'.format(private_key_path), newline=False)
        with open(private_key_path, 'rb') as f:
            data = f.read()
            # private_key_object = serialization.load_pem_private_key(data, None)
            private_key_string = data.encode('utf-8')
            logging.debug('private_key_string:\n{}'.format(private_key_string))

        self.msg_output('DONE')
        self.msg_output('Backing up existing certificate if needed..')
        self._backup_certificate(interface=interface)

        self.msg_output('Uploading the certificate to ECS..', newline=False)
        headers = {'X-SDS-AUTH-TOKEN': self.AUTH_TOKEN,
                   'Content-type': 'application/xml',
                   'X-EMC-RESTCLIENT': 'TRUE'}
        logging.debug('headers: {}'.format(headers))
        data = '<rotate_keycertchain><key_and_certificate>'
        data = data + '<private_key>{}</private_key><certificate_chain>{}</certificate_chain>'.format(
            private_key_string,
            certificate_string)
        data = data + '</key_and_certificate></rotate_keycertchain>'
        if interface == 'data':
            url = 'https://{}:4443/object-cert/keystore'.format(self.management_ip)
        elif interface == 'management':
            url = 'https://{}:4443/vdc/keystore'.format(self.management_ip)
        if not url:
            raise Exception('Unexpected Error')
        logging.debug('upload_certificate_payload:\n{}'.format(data))
        response = requests.put(url=url, data=data, verify=False, headers=headers, allow_redirects=True)
        logging.debug('{}'.format(vars(response)))
        if response.status_code == 200:
            self.msg_output('DONE')
            logging.info('Successfully uploaded certificate!')
        else:
            logging.error(
                'Failed to upload certificate. \n response: {}\n headers: {}'.format(response.text, response.headers))
        self.msg_output()

    def do_generate_san_config(self):
        try:
            self.msg_output('=' * 70)
            self.msg_output('Generating SAN (subject alternative name) config.')
            self.msg_output('=' * 70)
            self.msg_output()
            # Get data hostnames from DNS
            for ip in self.network_info['data']:
                hostname = self.get_hostname(ip)
                self.network_info['data_fqdns'].append(hostname)
            # Get management hostnames from dns
            for ip in self.network_info['management']:
                hostname = self.get_hostname(ip)
                self.network_info['management_fqdns'].append(hostname)

            logging.debug('Reading config file: {}'.format(self.config_file_path))
            self.config.read(self.config_file_path)
            # data
            logging.debug('START: {}'.format(self.DATA_SAN_CONF))
            self.config.set(self.DATA_SAN_CONF, 'IP_ADDRESSES', ' '.join(self.network_info['data']))
            self.config.set(self.DATA_SAN_CONF, 'DNS_NAMES', ' '.join(self.network_info['data_fqdns']))
            logging.debug('END: {}'.format(self.DATA_SAN_CONF))
            self.msg_output('-' * 70)
            self.msg_output('Setting {} config'.format(self.DATA_SAN_CONF))
            self.msg_output('-' * 70)
            self.msg_output('Set DNS_NAMES to :')
            self.msg_output(pformat(self.network_info['data_fqdns']))
            self.msg_output()
            self.msg_output('Set IP_ADDRESSES to :')
            self.msg_output(pformat(self.network_info['data']))
            self.msg_output()
            # management
            logging.debug('START: {}'.format(self.MANAGEMENT_SAN_CONF))
            self.config.set(self.MANAGEMENT_SAN_CONF, 'IP_ADDRESSES', ' '.join(self.network_info['management']))
            self.config.set(self.MANAGEMENT_SAN_CONF, 'DNS_NAMES', ' '.join(self.network_info['management_fqdns']))
            logging.debug('END: {}'.format(self.MANAGEMENT_SAN_CONF))
            self.msg_output('-' * 70)
            self.msg_output('Setting {} config'.format(self.MANAGEMENT_SAN_CONF))
            self.msg_output('-' * 70)
            self.msg_output('Set DNS_NAMES to :')
            self.msg_output(pformat(self.network_info['management_fqdns']))
            self.msg_output()
            self.msg_output('Set IP_ADDRESSES to :')
            self.msg_output(pformat(self.network_info['management']))
            self.msg_output()
            with open(self.config_file_path, 'wb') as f:
                self.config.write(f)
            self.msg_output('Wrote changes to: {}'.format(self.config_file_path))
            self.msg_output('DONE')
        except:
            self.msg_output('Exception while generating SAN config:\n{}'.format(traceback.format_exc()))

    def do_view_certs(self):
        try:
            self.msg_output('=' * 70)
            self.msg_output('Data Certificates')
            self.msg_output('=' * 70)
            self.msg_output()
            for data_ip in self.network_info['data']:
                try:
                    cert = self.get_certificate(data_ip, self.data_port)
                    san = self.get_san_names(cert)
                    common_name = self.get_common_name(cert)
                    issuer = self.get_issuer(cert)
                    not_valid_before, not_valid_after = self.get_valid_dates(cert)
                    self.msg_output('{}'.format(data_ip))
                    self.msg_output('-' * 70)
                    self.msg_output(
                        'common_name: {}\nissuer: {}\nsubject_alternative_name: {}'.format(common_name, issuer, san))
                    self.msg_output(
                        'not_valid_before: {}\nnot_valid_after: {}\n'.format(not_valid_before, not_valid_after))

                except:
                    self.msg_output('{}'.format(data_ip))
                    self.msg_output('-' * 70)
                    self.msg_output('Failed to view certificate')
                    logging.warning('traceback', exc_info=True)
                    self.msg_output()
                    continue

            self.msg_output('=' * 70)
            self.msg_output('Management Certificates')
            self.msg_output('=' * 70)
            self.msg_output()
            for management_ip in self.network_info['management']:
                try:
                    cert = self.get_certificate(management_ip, self.management_port)
                    san = self.get_san_names(cert)
                    common_name = self.get_common_name(cert)
                    issuer = self.get_issuer(cert)
                    not_valid_before, not_valid_after = self.get_valid_dates(cert)
                    self.msg_output('{}'.format(management_ip))
                    self.msg_output('-' * 70)
                    self.msg_output(
                        'common_name: {}\nissuer: {}\nsubject_alternative_name: {}'.format(common_name, issuer, san))
                    self.msg_output(
                        'not_valid_before: {}\nnot_valid_after: {}\n'.format(not_valid_before, not_valid_after))

                except:
                    self.msg_output('{}'.format(management_ip))
                    self.msg_output('-' * 70)
                    self.msg_output('Failed to view certificate')
                    logging.warning('traceback', exc_info=True)
                    self.msg_output()
                    continue

        except:
            self.msg_output('Exception while viewing certificates:\n{}'.format(traceback.format_exc()))

        self.msg_output('DONE')

    def do_create_csr(self):
        try:
            if not os.path.exists(self.generated_files_dir):
                os.mkdir(self.generated_files_dir)
            logging.debug('Reading config file: {}'.format(self.config_file_path))
            self.config.read(self.config_file_path, encoding='utf-8')
            self._validate_config()
            self.create_rsa_private_key()
            self.generate_csr()
        except Exception:
            self.msg_output('Exception while creating certificate signing request:\n{}'.format(traceback.format_exc()))

    def do_create_ssc(self):
        try:
            if not os.path.exists(self.generated_files_dir):
                os.mkdir(self.generated_files_dir)
            logging.debug('Reading config file: {}'.format(self.config_file_path))
            self.config.read(self.config_file_path, encoding='utf-8')
            self._validate_config()
            self.create_rsa_private_key()
            self.generate_ssc()
        except Exception:
            self.msg_output('Exception while creating self signed certificate:\n{}'.format(traceback.format_exc()))

    def do_configure_credentials(self):
        try:
            self.msg_output('\n=======> Configuring Credentials\n')
            pw = getpass.unix_getpass(prompt='Please enter the password for the root management user: ')
            self.management_pass = pw
            self.get_auth_token()
            with open(self.cred_cache_path, 'wb') as f:
                pickle.dump(pw, f)
            self.msg_output('Successfully configured credentials!')
        except:
            logging.error('Failed to configure credentials: ', exc_info=True)

    def main(self):
        """
        main program logic
        :return:
        """
        self._init_logging()
        self.config.read(self.config_file_path)
        self._get_network_info()
        self._discover_listeners()
        if os.path.exists(self.cred_cache_path):
            with open(self.cred_cache_path, 'rb') as f:
                self.management_pass = pickle.load(f)
        if self.args['action'] == 'view_certs':
            self.do_view_certs()
        if self.args['action'] == 'generate_san':
            self.do_generate_san_config()
        if self.args['action'] == 'create_csr':
            self.do_create_csr()
        if self.args['action'] == 'create_ssc':
            self.do_create_ssc()
        if self.args['action'] == 'upload_certificate':
            certificate_path = self.args['certificate']
            private_key_path = self.args['private_key']
            try:
                if self.args.get('data'):
                    self.upload_certificate(certificate_path, private_key_path, 'data')
                if self.args.get('management'):
                    self.upload_certificate(certificate_path, private_key_path, 'management')
            except:
                self.msg_output('Exception while uploading certificate:\n{}'.format(traceback.format_exc()))
        if self.args['action'] == 'nuke_cert':
            self.do_nuke_cert()
        if self.args['action'] == 'configure_credentials':
            self.do_configure_credentials()


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--debug', action="store_true", help='debug output')
    subparsers = parser.add_subparsers(metavar='{command}', help='sub-command help', dest="action")
    # Generate SAN info
    subparsers.add_parser('view_certs', help='Shows the current certificates on the data and management interfaces')
    subparsers.add_parser('generate_san', help='Generates the subject alternative name IP addresses and domain '
                                               'names from fabric and adds them to the ini config file')

    # Certificate signing request
    csr = subparsers.add_parser('create_csr', help='Create certificate signing request')
    csr_group = csr.add_mutually_exclusive_group(required=True)
    csr.add_argument('-k', '--key_size', help='Private key size for RSA private key generation (default=2048)',
                     type=int, choices=[1024, 2048, 4096], default=2048)
    csr_group.add_argument('-d', '--data',
                           help='Create certificate signing request for data interface (ports 9020, 9021,'
                                ' 3218) (S3, CAS, NFS, etc)',
                           action='store_true')
    csr_group.add_argument('-m', '--management',
                           help='Create certificate signing request for management interface (WEB UI)',
                           action='store_true')
    # Self signed certificate
    ssc = subparsers.add_parser('create_ssc', help='Create self-signed certificate')
    ssc_group = ssc.add_mutually_exclusive_group(required=True)
    ssc.add_argument('-k', '--key_size', help='Private key size for RSA private key generation (default=2048)',
                     type=int, choices=[1024, 2048, 4096], default=2048)
    ssc_group.add_argument('-d', '--data',
                           help='Create self-signed certificate for data interface (ports 9020, 9021, 3218) (S3, CAS, NFS, etc)',
                           action='store_true')
    ssc_group.add_argument('-m', '--management',
                           help='Create self-signed certificate for management interface (WEB UI)',
                           action='store_true')
    # upload data
    upload_certificate = subparsers.add_parser('upload_certificate', help='Upload certificate to the data interface')
    upload_certificate.add_argument('-c', '--certificate', help='Filepath to the data certificate', required=True)
    upload_certificate.add_argument('-p', '--private_key', help='Filepath to private key with no password',
                                    required=True)
    upload_group = upload_certificate.add_mutually_exclusive_group(required=True)
    upload_group.add_argument('-d', '--data',
                              help='Upload certificate to the data interface',
                              action='store_true')
    upload_group.add_argument('-m', '--management',
                              help='Upload certificate to the management interface',
                              action='store_true')
    nuke_cert = subparsers.add_parser('nuke_cert')
    nuke_cert_group = nuke_cert.add_mutually_exclusive_group(required=True)
    nuke_cert_group.add_argument('-d', '--data', help='Nuke certificate for data interface', action='store_true')
    nuke_cert_group.add_argument('-m', '--management', help='Nuke certificate for management interface',
                                 action='store_true')
    subparsers.add_parser('configure_credentials', help='Configure the UI credentials for the tool to use')

    return vars(parser.parse_args())


if __name__ == '__main__':
    sys.stdout.write('ecs_certificate_tool v{}\n'.format(__version__))
    #  #25
    if os.geteuid() == 0:
        logging.warning('Please run this tool as admin user without sudo')
        # exit(100)
    file_permission_commands = ['sudo chown -R admin {}/*'.format(os.path.dirname(os.path.abspath(__file__))),
                                'sudo chgrp -R users {}/*'.format(os.path.dirname(os.path.abspath(__file__))),
                                'sudo chmod +x -R {}/*'.format(os.path.dirname(os.path.abspath(__file__)))]
    for permission_command in file_permission_commands:
        subprocess.check_call(permission_command, shell=True)
    params = get_args()
    tool_inst = EcsCertificateTool(params)
    tool_inst.main()

