# -*- coding: utf-8 -*-
'''
Microsoft certificate management via the PKI Client PowerShell module.
https://technet.microsoft.com/en-us/itpro/powershell/windows/pkiclient/pkiclient

The PKI Client PowerShell module is only available on Windows 8+ and Windows
Server 2012+.
https://technet.microsoft.com/en-us/library/hh848636(v=wps.620).aspx

:platform:      Windows

:depends:
    - PowerShell 4
    - PKI Client Module (Windows 8+ / Windows Server 2012+)

.. versionadded:: 2016.11.0
'''
# Import Python libs
from __future__ import absolute_import, unicode_literals, print_function
import ast
import logging
import os

# Import salt libs
import salt.utils.json
import salt.utils.platform
import salt.utils.powershell
import salt.utils.versions
from salt.exceptions import SaltInvocationError

# Import 3rd party libs
from salt.ext import six

_DEFAULT_CONTEXT = 'LocalMachine'
_DEFAULT_FORMAT = 'cer'
_DEFAULT_STORE = 'My'
_LOG = logging.getLogger(__name__)

# Define the module's virtual name
__virtualname__ = 'win_pki'


def __virtual__():
    '''
    Requires Windows
    Requires Windows 8+ / Windows Server 2012+
    Requires PowerShell
    Requires PKI Client PowerShell module installed.
    '''
    if not salt.utils.platform.is_windows():
        return False, 'Only available on Windows Systems'

    if salt.utils.versions.version_cmp(__grains__['osversion'], '6.2.9200') == -1:
        return False, 'Only available on Windows 8+ / Windows Server 2012 +'

    if not __salt__['cmd.shell_info']('powershell')['installed']:
        return False, 'Powershell not available'

    if not salt.utils.powershell.module_exists('PKI'):
        return False, 'PowerShell PKI module not available'

    return __virtualname__


def _cmd_run(cmd, as_json=False):
    '''
    Ensure that the Pki module is loaded, and convert to and extract data from
    Json as needed.
    '''
    cmd_full = ['Import-Module -Name PKI; ']

    if as_json:
        cmd_full.append(r'ConvertTo-Json -Compress -Depth 4 -InputObject '
                        r'@({0})'.format(cmd))
    else:
        cmd_full.append(cmd)
    cmd_ret = __salt__['cmd.run_all'](
        six.text_type().join(cmd_full), shell='powershell', python_shell=True)

    if cmd_ret['retcode'] != 0:
        _LOG.error('Unable to execute command: %s\nError: %s', cmd,
                   cmd_ret['stderr'])

    if as_json:
        try:
            items = salt.utils.json.loads(cmd_ret['stdout'], strict=False)
            return items
        except ValueError:
            _LOG.error('Unable to parse return data as Json.')

    return cmd_ret['stdout']


def _validate_cert_path(name):
    '''
    Ensure that the certificate path, as determind from user input, is valid.
    '''
    cmd = r"Test-Path -Path '{0}'".format(name)

    if not ast.literal_eval(_cmd_run(cmd=cmd)):
        raise SaltInvocationError(r"Invalid path specified: {0}".format(name))


def _validate_cert_format(name):
    '''
    Ensure that the certificate format, as determind from user input, is valid.
    '''
    cert_formats = ['cer', 'pfx']

    if name not in cert_formats:
        message = ("Invalid certificate format '{0}' specified. Valid formats:"
                   ' {1}').format(name, cert_formats)
        raise SaltInvocationError(message)


def get_stores():
    '''
    Get the certificate location contexts and their corresponding stores.

    :return: A dictionary of the certificate location contexts and stores.
    :rtype: dict

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.get_stores
    '''
    ret = dict()
    cmd = r"Get-ChildItem -Path 'Cert:\' | " \
          r"Select-Object LocationName, StoreNames"

    items = _cmd_run(cmd=cmd, as_json=True)

    for item in items:
        ret[item['LocationName']] = list()

        for store in item['StoreNames']:
            ret[item['LocationName']].append(store)
    return ret


def get_certs(context=_DEFAULT_CONTEXT, store=_DEFAULT_STORE):
    '''
    Get the available certificates in the given store.

    :param str context: The name of the certificate store location context.
    :param str store: The name of the certificate store.

    :return: A dictionary of the certificate thumbprints and properties.
    :rtype: dict

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.get_certs
    '''
    ret = dict()
    cmd = list()
    blacklist_keys = ['DnsNameList']
    store_path = r'Cert:\{0}\{1}'.format(context, store)

    _validate_cert_path(name=store_path)

    cmd.append(r"Get-ChildItem -Path '{0}' | Select-Object".format(store_path))
    cmd.append(' DnsNameList, SerialNumber, Subject, Thumbprint, Version')

    items = _cmd_run(cmd=six.text_type().join(cmd), as_json=True)

    for item in items:
        cert_info = dict()
        for key in item:
            if key not in blacklist_keys:
                cert_info[key.lower()] = item[key]

        names = item.get('DnsNameList', None)
        if isinstance(names, list):
            cert_info['dnsnames'] = [name.get('Unicode') for name in names]
        else:
            cert_info['dnsnames'] = []
        ret[item['Thumbprint']] = cert_info
    return ret


def get_cert_file(name, cert_format=_DEFAULT_FORMAT, password=''):
    '''
    Get the details of the certificate file.

    :param str name: The filesystem path of the certificate file.
    :param str cert_format: The certificate format. Specify 'cer' for X.509, or
        'pfx' for PKCS #12.
    :param str password: The password of the certificate. Only applicable to pfx
        format. Note that if used interactively, the password will be seen by all minions.
        To protect the password, use a state and get the password from pillar.

    :return: A dictionary of the certificate thumbprints and properties.
    :rtype: dict

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.get_cert_file name='C:\\certs\\example.cer'
    '''
    ret = dict()
    cmd = list()
    blacklist_keys = ['DnsNameList']
    cert_format = cert_format.lower()

    _validate_cert_format(name=cert_format)

    if not name or not os.path.isfile(name):
        _LOG.error('Path is not present: %s', name)
        return ret

    if cert_format == 'pfx':
        if password:
            cmd.append('$CertObject = New-Object')
            cmd.append(' System.Security.Cryptography.X509Certificates.X509Certificate2;')
            cmd.append(r" $CertObject.Import('{0}'".format(name))
            cmd.append(",'{0}'".format(password))
            cmd.append(",'DefaultKeySet') ; $CertObject")
            cmd.append(' | Select-Object DnsNameList, SerialNumber, Subject, '
                    'Thumbprint, Version')
        else:
            cmd.append(r"Get-PfxCertificate -FilePath '{0}'".format(name))
            cmd.append(' | Select-Object DnsNameList, SerialNumber, Subject, '
                    'Thumbprint, Version')
    else:
        cmd.append('$CertObject = New-Object')
        cmd.append(' System.Security.Cryptography.X509Certificates.X509Certificate2;')
        cmd.append(r" $CertObject.Import('{0}'); $CertObject".format(name))
        cmd.append(' | Select-Object DnsNameList, SerialNumber, Subject, '
                   'Thumbprint, Version')

    items = _cmd_run(cmd=six.text_type().join(cmd), as_json=True)

    for item in items:
        for key in item:
            if key not in blacklist_keys:
                ret[key.lower()] = item[key]

        ret['dnsnames'] = [name['Unicode'] for name in item['DnsNameList']]

    if ret:
        _LOG.debug('Certificate thumbprint obtained successfully: %s', name)
    else:
        _LOG.error('Unable to obtain certificate thumbprint: %s', name)
    return ret


def import_cert(name,
                cert_format=_DEFAULT_FORMAT,
                context=_DEFAULT_CONTEXT,
                store=_DEFAULT_STORE,
                exportable=True,
                password='',
                saltenv='base'):
    '''
    Import the certificate file into the given certificate store.

    :param str name: The path of the certificate file to import.
    :param str cert_format: The certificate format. Specify 'cer' for X.509, or
        'pfx' for PKCS #12.
    :param str context: The name of the certificate store location context.
    :param str store: The name of the certificate store.
    :param bool exportable: Mark the certificate as exportable. Only applicable
        to pfx format.
    :param str password: The password of the certificate. Only applicable to pfx
        format. Note that if used interactively, the password will be seen by all minions.
        To protect the password, use a state and get the password from pillar.
    :param str saltenv: The environment the file resides in.

    :return: A boolean representing whether all changes succeeded.
    :rtype: bool

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.import_cert name='salt://cert.cer'
    '''
    cmd = list()
    thumbprint = None
    store_path = r'Cert:\{0}\{1}'.format(context, store)
    cert_format = cert_format.lower()

    _validate_cert_format(name=cert_format)

    cached_source_path = __salt__['cp.cache_file'](name, saltenv)

    if not cached_source_path:
        _LOG.error('Unable to get cached copy of file: %s', name)
        return False

    if password:
        cert_props = get_cert_file(name=cached_source_path, cert_format=cert_format, password=password)
    else:
        cert_props = get_cert_file(name=cached_source_path, cert_format=cert_format)

    current_certs = get_certs(context=context, store=store)

    if cert_props['thumbprint'] in current_certs:
        _LOG.debug("Certificate thumbprint '%s' already present in store: %s",
                   cert_props['thumbprint'], store_path)
        return True

    if cert_format == 'pfx':
        # In instances where an empty password is needed, we use a
        # System.Security.SecureString object since ConvertTo-SecureString will
        # not convert an empty string.
        if password:
            cmd.append(r"$Password = ConvertTo-SecureString "
                       r"-String '{0}'".format(password))
            cmd.append(' -AsPlainText -Force; ')
        else:
            cmd.append('$Password = New-Object System.Security.SecureString; ')

        cmd.append(r"Import-PfxCertificate "
                   r"-FilePath '{0}'".format(cached_source_path))
        cmd.append(r" -CertStoreLocation '{0}'".format(store_path))
        cmd.append(r" -Password $Password")

        if exportable:
            cmd.append(' -Exportable')
    else:
        cmd.append(r"Import-Certificate "
                   r"-FilePath '{0}'".format(cached_source_path))
        cmd.append(r" -CertStoreLocation '{0}'".format(store_path))

    _cmd_run(cmd=six.text_type().join(cmd))

    new_certs = get_certs(context=context, store=store)

    for new_cert in new_certs:
        if new_cert not in current_certs:
            thumbprint = new_cert

    if thumbprint:
        _LOG.debug('Certificate imported successfully: %s', name)
        return True
    _LOG.error('Unable to import certificate: %s', name)
    return False


def export_cert(name,
                thumbprint,
                cert_format=_DEFAULT_FORMAT,
                context=_DEFAULT_CONTEXT,
                store=_DEFAULT_STORE,
                password=''):
    '''
    Export the certificate to a file from the given certificate store.

    :param str name: The destination path for the exported certificate file.
    :param str thumbprint: The thumbprint value of the target certificate.
    :param str cert_format: The certificate format. Specify 'cer' for X.509, or
        'pfx' for PKCS #12.
    :param str context: The name of the certificate store location context.
    :param str store: The name of the certificate store.
    :param str password: The password of the certificate. Only applicable to pfx
        format. Note that if used interactively, the password will be seen by all minions.
        To protect the password, use a state and get the password from pillar.

    :return: A boolean representing whether all changes succeeded.
    :rtype: bool

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.export_cert name='C:\\certs\\example.cer' thumbprint='AAA000'
    '''
    cmd = list()
    thumbprint = thumbprint.upper()
    cert_path = r'Cert:\{0}\{1}\{2}'.format(context, store, thumbprint)
    cert_format = cert_format.lower()

    _validate_cert_path(name=cert_path)
    _validate_cert_format(name=cert_format)

    if cert_format == 'pfx':
        # In instances where an empty password is needed, we use a
        # System.Security.SecureString object since ConvertTo-SecureString will
        # not convert an empty string.
        if password:
            cmd.append(r"$Password = ConvertTo-SecureString "
                       r"-String '{0}'".format(password))
            cmd.append(' -AsPlainText -Force; ')
        else:
            cmd.append('$Password = New-Object System.Security.SecureString; ')

        cmd.append(r"Export-PfxCertificate "
                   r"-Cert '{0}' -FilePath '{1}'".format(cert_path, name))
        cmd.append(r" -Password $Password")
    else:
        cmd.append(r"Export-Certificate "
                   r"-Cert '{0}' -FilePath '{1}'".format(cert_path, name))

    cmd.append(r" | Out-Null; Test-Path -Path '{0}'".format(name))

    ret = ast.literal_eval(_cmd_run(cmd=six.text_type().join(cmd)))

    if ret:
        _LOG.debug('Certificate exported successfully: %s', name)
    else:
        _LOG.error('Unable to export certificate: %s', name)
    return ret


def test_cert(thumbprint,
              context=_DEFAULT_CONTEXT,
              store=_DEFAULT_STORE,
              untrusted_root=False,
              dns_name='',
              eku=''):
    '''
    Check the certificate for validity.

    :param str thumbprint: The thumbprint value of the target certificate.
    :param str context: The name of the certificate store location context.
    :param str store: The name of the certificate store.
    :param bool untrusted_root: Whether the root certificate is required to be
        trusted in chain building.
    :param str dns_name: The DNS name to verify as valid for the certificate.
    :param str eku: The enhanced key usage object identifiers to verify for the
        certificate chain.

    :return: A boolean representing whether the certificate was considered
        valid.
    :rtype: bool

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.test_cert thumbprint='AAA000' dns_name='example.test'
    '''
    cmd = list()
    thumbprint = thumbprint.upper()
    cert_path = r'Cert:\{0}\{1}\{2}'.format(context, store, thumbprint)
    cmd.append(r"Test-Certificate -Cert '{0}'".format(cert_path))

    _validate_cert_path(name=cert_path)

    if untrusted_root:
        cmd.append(' -AllowUntrustedRoot')
    if dns_name:
        cmd.append(" -DnsName '{0}'".format(dns_name))
    if eku:
        cmd.append(" -EKU '{0}'".format(eku))

    cmd.append(' -ErrorAction SilentlyContinue')

    return ast.literal_eval(_cmd_run(cmd=six.text_type().join(cmd)))


def remove_cert(thumbprint, context=_DEFAULT_CONTEXT, store=_DEFAULT_STORE):
    '''
    Remove the certificate from the given certificate store.

    :param str thumbprint: The thumbprint value of the target certificate.
    :param str context: The name of the certificate store location context.
    :param str store: The name of the certificate store.

    :return: A boolean representing whether all changes succeeded.
    :rtype: bool

    CLI Example:

    .. code-block:: bash

        salt '*' win_pki.remove_cert thumbprint='AAA000'
    '''
    thumbprint = thumbprint.upper()
    store_path = r'Cert:\{0}\{1}'.format(context, store)
    cert_path = r'{0}\{1}'.format(store_path, thumbprint)
    cmd = r"Remove-Item -Path '{0}'".format(cert_path)

    current_certs = get_certs(context=context, store=store)

    if thumbprint not in current_certs:
        _LOG.debug("Certificate '%s' already absent in store: %s", thumbprint,
                   store_path)
        return True

    _validate_cert_path(name=cert_path)
    _cmd_run(cmd=cmd)

    new_certs = get_certs(context=context, store=store)

    if thumbprint in new_certs:
        _LOG.error('Unable to remove certificate: %s', cert_path)
        return False
    _LOG.debug('Certificate removed successfully: %s', cert_path)
    return True