# -*- coding: utf-8 -*-
import os
import sys
import subprocess
import time
import ConfigParser
import uuid
import logging
import errno
# import urlparse
from OpenSSL import crypto, SSL
import requests
import traceback
from pyasn1.codec.der import encoder as der_encoder
from pyasn1.type import tag
from pyasn1_modules import rfc2459


def popenCommunicate(command_list, input=None):
  subprocess_kw = dict(stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
  if input is not None:
    subprocess_kw.update(stdin=subprocess.PIPE)
  popen = subprocess.Popen(command_list, **subprocess_kw)
  result = popen.communicate(input)[0]
  if popen.returncode is None:
    popen.kill()
  if popen.returncode != 0:
    raise ValueError('Issue during calling %r, result was:\n%s' % (
      command_list, result))
  return result

class BusyCertificateAuthority(Exception):
  """Exception raised when certificate authority is busy"""
  pass

class GeneralNames(rfc2459.GeneralNames):
    """
    rfc2459 has wrong tagset.
    """
    tagSet = tag.TagSet(
        (),
        tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0),
        )

class DistributionPointName(rfc2459.DistributionPointName):
    """
    rfc2459 has wrong tagset.
    """
    tagSet = tag.TagSet(
        (),
        tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0),
        )

class ExtensionType():

  CRL_DIST_POINTS = "crlDistributionPoints"
  BASIC_CONSTRAINTS = "basicConstraints"
  KEY_USAGE = "keyUsage"
  NS_CERT_TYPE = "nsCertType"
  NS_COMMENT = "nsComment"
  SUBJECT_KEY_ID = "subjectKeyIdentifier"
  AUTH_KEY_ID = "authorityKeyIdentifier"


class X509Extension(object):

  known_extension_list = [name for (attr, name) in vars(ExtensionType).items()
                      if attr.isupper()]

  def setX509Extension(self, ext_type, critical, value, subject=None, issuer=None):
    if not ext_type in self.known_extension_list:
      raise ValueError('Extension type is not known from ExtensionType class')

    if ext_type == ExtensionType.CRL_DIST_POINTS:
      cdp = self._getCrlDistPointExt(value)
      return crypto.X509Extension(
        b'%s' % ext_type,
        critical,
        'DER:' + cdp.encode('hex'),
        subject=subject,
        issuer=issuer,
      )
    else:
      return crypto.X509Extension(
        ext_type,
        critical,
        value,
        subject=subject,
        issuer=issuer,
      )

  def _getCrlDistPointExt(self, cdp_list):
    cdp = rfc2459.CRLDistPointsSyntax()
    position = 0
    for cdp_type, cdp_value in cdp_list:
      cdp_entry = rfc2459.DistributionPoint()
      general_name = rfc2459.GeneralName()

      if not cdp_type in ['dNSName', 'directoryName', 'uniformResourceIdentifier']:
        raise ValueError("crlDistributionPoints GeneralName '%s' is not valid" % cdp_type)

      general_name.setComponentByName(cdp_type, cdp_value)

      general_names = GeneralNames()
      general_names.setComponentByPosition(0, general_name)

      name = DistributionPointName()
      name.setComponentByName('fullName', general_names)
      cdp_entry.setComponentByName('distributionPoint', name)

      cdp.setComponentByPosition(position, cdp_entry)
      position += 1

    return der_encoder.encode(cdp)

  def setCaExtensions(self, cert_obj, issuer=None):
    """
      extensions for default certificate
    """
    cert_obj.add_extensions([
      self.setX509Extension(ExtensionType.BASIC_CONSTRAINTS, True,
                                                          "CA:TRUE, pathlen:0"),
      self.setX509Extension(ExtensionType.NS_COMMENT,
                                              False, "OpenSSL CA Certificate"),
      self.setX509Extension(ExtensionType.KEY_USAGE,
                                                 True, "keyCertSign, cRLSign"),
      self.setX509Extension(ExtensionType.SUBJECT_KEY_ID,
                                                False, "hash", subject=cert_obj),
    ])
    cert_obj.add_extensions([
      self.setX509Extension(ExtensionType.AUTH_KEY_ID,
                                  False, "keyid:always,issuer", issuer=cert_obj)
    ])

  def setDefaultExtensions(self, cert_obj, subject=None, issuer=None, crl_url=None):
    """
      extensions for default certificate
    """
    cert_obj.add_extensions([
      self.setX509Extension(ExtensionType.BASIC_CONSTRAINTS, False, "CA:FALSE"),
      self.setX509Extension(ExtensionType.NS_COMMENT,
                                        False, "OpenSSL Generated Certificate"),
      self.setX509Extension(ExtensionType.SUBJECT_KEY_ID,
                                                False, "hash", subject=subject),
    ])
    cert_obj.add_extensions([
      self.setX509Extension(ExtensionType.AUTH_KEY_ID,
                                          False, "keyid,issuer", issuer=issuer)
    ])
    if crl_url:
      cert_obj.add_extensions([
        self.setX509Extension(ExtensionType.CRL_DIST_POINTS,
                          False, [("uniformResourceIdentifier", crl_url)])
      ])

  def setDefaultCsrExtensions(self, cert_obj, subject=None, issuer=None):
    """
      extensions for certificate signature request
    """
    cert_obj.add_extensions([
      self.setX509Extension(ExtensionType.BASIC_CONSTRAINTS, False, "CA:FALSE"),
      self.setX509Extension(ExtensionType.KEY_USAGE,
                    False, "nonRepudiation, digitalSignature, keyEncipherment"),
    ])

class CertificateBase(object):

  def __init__(self):
    self.X509Extension = X509Extension()
    pass

  def validateCertAndKey(self, cert_file, key_file):
    ctx = SSL.Context(SSL.TLSv1_METHOD)
    ctx.use_privatekey(self.freadPkey(key_file))
    ctx.use_certificate(self.freadX509(cert_file))
    try:
      ctx.check_privatekey()
    except SSL.Error:
      return False
    else:
      return True

  def verifyCertificateChain(self, cert_pem, trusted_cert_list):

    cert = crypto.load_certificate(crypto.FILETYPE_PEM, cert_pem)

    # Create and fill a X509Sore with trusted certs
    store = crypto.X509Store()
    for trusted_cert in trusted_cert_list:
      store.add_cert(crypto.load_certificate(crypto.FILETYPE_PEM, trusted_cert))

    store_ctx = crypto.X509StoreContext(store, cert)
    # Returns None if certificate can be validated
    result = store_ctx.verify_certificate()

    if result is None:
      return True
    else:
      return False

  def checkCertificateValidity(self, ca_cert_file, cert_file, key_file=None):
    with open(ca_cert_file) as f_ca:
      ca_cert = f_ca.read()
    with open(cert_file) as f_cert:
      cert = f_cert.read()

    # XXX Considering only one trusted certificate here
    if not self.verifyCertificateChain(cert, [ca_cert]):
      return False
    if key_file:
      return self.validateCertAndKey(cert_file, key_file)
    return True

  def generatePrivatekey(self, output_file, size=2048):
    key = crypto.PKey()
    key.generate_key(crypto.TYPE_RSA, size)

    try:
      key_fd = os.open(output_file,
                       os.O_CREAT|os.O_WRONLY|os.O_EXCL|os.O_TRUNC,
                       0600)
    except OSError, e:
      if e.errno != errno.EEXIST:
        raise
    else:
      os.write(key_fd, crypto.dump_privatekey(crypto.FILETYPE_PEM, key))
      os.close(key_fd)

  def generateCertificateRequest(self, key_file, output_file, cn,
      country, state, locality='', email='', organization='',
      organization_unit='', digest="sha256"):

    with open(key_file) as fkey:
      key = crypto.load_privatekey(crypto.FILETYPE_PEM, fkey.read())

    req = crypto.X509Req()
    subject = req.get_subject()
    subject.CN = cn
    subject.C = country
    subject.ST = state
    subject.L = locality
    subject.O = organization
    subject.OU = organization_unit
    subject.emailAddress = email
    req.set_pubkey(key)
    self.X509Extension.setDefaultCsrExtensions(req)
    req.sign(key, digest)

    with open(output_file, 'w') as req_file:
      req_file.write(crypto.dump_certificate_request(crypto.FILETYPE_PEM, req))

    os.chmod(output_file, 0644)

  def signData(self, key_file, data, digest="sha256", output_file=None):
    """
      Sign a data using digest and return signature. If output_file is provided
      the signature will be written into the file.
    """
    with open(key_file) as fkey:
      pkey = crypto.load_privatekey(crypto.FILETYPE_PEM, fkey.read())

    sign = crypto.sign(pkey, data, digest)
    # data_base64 = base64.b64encode(sign)
    if output_file is None:
      return sign
      fd = os.open(output_file,
                   os.O_CREAT|os.O_WRONLY|os.O_EXCL|os.O_TRUNC,
                   0644)
      try:
        os.write(fd, sign)
      finally:
        os.close(fd)
    return sign

  def verifyData(self, cert_string, signature, data, digest="sha256"):
    """
      Verify the signature for a data string.

      cert_string: is the certificate content as string
      signature: is generate using 'signData' from the data to verify
      data: content to verify
      digest: by default is sha256, set the correct value
    """
    x509 = crypto.load_certificate(crypto.FILETYPE_PEM, cert)
    return crypto.verify(x509, signature, data, digest)

  def readCertificateRequest(self, csr):
    req = crypto.load_certificate_request(crypto.FILETYPE_PEM, csr)

    return req

  def freadCertificateRequest(self, csr_file):
    with open(csr_file) as f_csr:
      return self.readCertificateRequest(f_csr.read())

  def freadX509(self, cert_file):
    with open(cert_file) as f_cert:
      return self.readX509(f_cert.read())

  def readX509(self, cert):
    x509 = crypto.load_certificate(crypto.FILETYPE_PEM, cert)
    return x509

  def freadCrl(self, crl_file):
    with open(crl_file) as f_crl:
      return self.readCrl(f_crl.read())

  def readCrl(self, crl):
    return crypto.load_crl(crypto.FILETYPE_PEM, crl)

  def readPkey(self, key):
    return crypto.load_privatekey(crypto.FILETYPE_PEM, key)

  def freadPkey(self, key_file):
    with open(key_file) as f_key:
      return self.readPkey(f_key.read())

  def getSubject(self, cert):
    subject = cert.get_subject()
    subject_dict = dict(subject.get_components())
    return subject_dict

  def getX509Text(self, openssl_binary, cert_file):
    # XXX - using openssl_bin to get text information from certificate
    return popenCommunicate([openssl_binary, 'x509', '-in', cert_file, '-text'])

  def getCsrText(self, openssl_binary, csr_file):
    # XXX - using openssl_bin to get text information from certificate
    return popenCommunicate([openssl_binary, 'req', '-in', csr_file, '-text'])

  def getCrlText(self, openssl_binary, crl_file):
    # XXX - using openssl_bin to get text information from certificate
    return popenCommunicate([openssl_binary, 'crl', '-in', crl_file, '-text'])


class CertificateAuthority(CertificateBase):

  def __init__(self, openssl_binary,
      openssl_configuration=None, certificate='cacert.pem',
      key='cakey.pem', crl='cacrl.pem', ca_directory=None):
    self.key = key
    self.certificate = certificate
    self.openssl_binary = openssl_binary
    self.openssl_configuration = openssl_configuration
    self.ca_directory = ca_directory
    if self.ca_directory is None:
      self.ca_directory = os.getcwd()

    self.ca_crl = crl
    self.lock_file = os.path.join(self.ca_directory, 'ca_locked')

    for file in ['crlnumber', 'serial']:
      if not os.path.exists(os.path.join(self.ca_directory, file)):
        with open(os.path.join(self.ca_directory, file), 'w') as f:
          f.write('01')

    if not os.path.exists(os.path.join(self.ca_directory, 'index.txt')):
      with open(os.path.join(self.ca_directory, 'index.txt'), 'w') as f:
        f.write('')

  def _lock(self):
    """Checks lock and locks Certificate Authority tool

       Raises BusyCertificateAuthority"""
    if os.path.exists(self.lock_file):
      raise BusyCertificateAuthority
    open(self.lock_file, 'w').write('locked')

  def _unlock(self):
    """Checks lock and locks Certificate Authority tool"""
    if os.path.exists(self.lock_file):
      os.unlink(self.lock_file)
    else:
      pass

  def checkAuthority(self):
    if self.openssl_configuration is None:
      raise ValueError('Openssl configuration file not found!')

    file_list = [ self.key, self.certificate ]
    ca_ready = True
    for f in file_list:
      if not os.path.exists(f):
        ca_ready = False
        break
    if ca_ready:
      return
    for f in file_list:
      if os.path.exists(f):
        os.unlink(f)
    try:
      # no CA, let us create new one
      popenCommunicate([self.openssl_binary, 'req', '-nodes', '-config',
          self.openssl_configuration, '-new', '-x509', '-extensions',
          'v3_ca', '-keyout', self.key, '-out', self.certificate,
          '-days', '10950'], 'Certificate Authority %s\n' % uuid.uuid1())
    except:
      try:
        for f in file_list:
          if os.path.exists(f):
            os.unlink(f)
      except:
        # do not raise during cleanup
        pass
      raise

  def signCertificateRequest(self, csr_file, output_file):
    if self.openssl_configuration is None:
      raise ValueError('Openssl configuration file not found!')

    popenCommunicate([self.openssl_binary, 'ca', '-batch', '-config',
          self.openssl_configuration, '-out', output_file,
          '-infiles', csr_file])

  def genCertificateRevocationList(self):
    if self.openssl_configuration is None:
      raise ValueError('Openssl configuration file not found!')

    popenCommunicate([self.openssl_binary, 'ca', '-batch', '-config',
          self.openssl_configuration, '-gencrl', '-keyfile', self.key,
          '-cert', self.certificate, '-out', self.ca_crl])

  def revokeCertificate(self, cert_file, update_crl=False):
    if self.openssl_configuration is None:
      raise ValueError('Openssl configuration file not found!')

    popenCommunicate([self.openssl_binary, 'ca', '-batch', '-config',
          self.openssl_configuration, '-revoke', cert_file])

    if update_crl:
      self.genCertificateRevocationList()

  def getNextSerial(self):
    # XXX - serial is managed by openssl
    return open(os.path.join(self.ca_directory, 'serial')).read().strip()


class CertificateAuthorityRequest(CertificateBase):

  def __init__(self, key, certificate, cacertificate, ca_url,
               max_retry=10, logger=None):

    self.key = key
    self.certificate = certificate
    self.cacertificate = cacertificate
    self.ca_url = ca_url
    self.logger = logger
    self.max_retry = max_retry
    self.X509Extension = X509Extension()

    while self.ca_url.endswith('/'):
      # remove all / at end or ca_url
      self.ca_url = self.ca_url[:-1]

    if self.logger is None:
      self.logger = logging.getLogger('Certificate Request')
      self.logger.setLevel(logging.DEBUG)
      handler = logging.StreamHandler()
      handler.setLevel(logging.DEBUG)
      formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
      handler.setFormatter(formatter)

      self.logger.addHandler(handler)

    self.generatePrivatekey(self.key)

  def _request(self, method, url, data={}):
    try:
      req = getattr(requests, method)
      if data:
        return req(url, data=data)
      else:
        return req(url)
    except requests.ConnectionError, e:
      self.logger.error("Got ConnectionError while sending request to CA. Url is %s\n%s" % (
          url, str(e)))
      return None

  def getCACertificateWeb(self):

    # If cert file exists exist
    if os.path.exists(self.cacertificate) and os.stat(self.cacertificate).st_size > 0:
      return

    ca_cert_url = '%s/crt/cacert.pem' % self.ca_url
    self.logger.info("getting CA certificate file %s" % ca_cert_url)
    response = None
    while not response or response.status_code != 200:
      response = self._request('get', ca_cert_url)
      if response is not None:
        try:
          x509 = crypto.load_certificate(crypto.FILETYPE_PEM, response.text)
        except crypto.Error, e:
          traceback.print_exc()
          response = None
        else:
          break
      # sleep a bit then try again until  ca cert is ready
      time.sleep(10)

    fd = os.open(self.cacertificate,
                  os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC, 0644)
    try:
      os.write(fd, response.text)
    finally:
      os.close(fd)

  def signCertificateWeb(self, csr_file, auto_revoke=False):

    if os.path.exists(self.certificate) and os.stat(self.certificate).st_size > 0:
      return
    with open(csr_file) as fcsr:
      csr = fcsr.read()

    data = {'csr': csr}
    retry = 0
    sleep_time = 10
    request_url = '%s/csr' % self.ca_url
    # Save Cert in tmp to check later
    cert_temp = '%s.tmp' % self.certificate
    csr_key_file = '%s.key' % csr_file
    csr_key = ""

    self.logger.info("Request signed certificate from CA...")
    if os.path.exists(csr_key_file):
      with open(csr_key_file) as fkey:
        csr_key = fkey.read()

    if csr_key:
      self.logger.info("Csr was already sent to CA, using csr : %s" % csr_key)
    else:
      response = self._request('put', request_url, data=data)

      while (not response or response.status_code != 201) and retry < self.max_retry:

        self.logger.error("%s: Failed to sent CSR. \n%s" % (
            response.status_code, response.text))
        self.logger.info("will retry in %s seconds..." % sleep_time)
        time.sleep(sleep_time)
        retry += 1
        response = self._request('put', request_url, data=data)

      if response.status_code != 201:
        raise Exception("ERROR: failed to put CSR after % retry. Exiting..." % retry)

      self.logger.info("CSR succefully sent.")
      # Get csr Location from request header: http://xxx.com/csr/key
      self.logger.debug("Csr location is: %s" % response.headers['Location'])

      csr_key = response.headers['Location'].split('/')[-1]
      with open(csr_key_file, 'w') as fkey:
        fkey.write(response.text)

      # csr is xxx.csr.pem so cert is xxx.cert.pem
    self.logger.info("Waiting for signed certificate...")

    reply_url = '%s/crt/%s.cert.pem' % (self.ca_url, csr_key[:-8])
    response = self._request('get', reply_url)

    while not response or response.status_code != 200:
      time.sleep(sleep_time)
      response = self._request('get', reply_url)

    with open(cert_temp, 'w') as cf:
      cf.write(response.text)

    self.logger.info("Validating signed certificate...")
    if self.checkCertificateValidity(self.cacertificate, cert_temp, self.key):
      fd = os.open(self.certificate,
                    os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC, 0644)
      try:
        os.write(fd, response.text)
      finally:
        os.close(fd)
        os.unlink(cert_temp)
    else:
      """if auto_revoke:
        self.logger.error("Certificate validation failed. " \
            "The signed certificate is going to be revoked...")
        self.revokeCertificateRequest(cert_temp,
                                      csr_key,
                                      message="Certificate validation failed.")
        # Cleanup
        try:
          os.remove(csr_key_file)
        except OSError, e:
          if e.errno != errno.ENOENT:
            # raise
            pass"""
      raise Exception("Error: Certificate validation failed. " \
          "This signed certificate should be revoked!")
    self.logger.info("Certificate correctly saved at %s." % self.certificate)

  def revokeCertificateRequest(self, cert_file, message=""):
    """
    Send a revocation request for the givent certificate to the master.
    """
    sleep_time = 10
    retry = 0

    with open(cert_file) as f:
      cert_string = f.read()

    cert = self.readX509(cert_string)
    digest = "sha256"
    payload = json.dumps(dict(
      reason=message,
      cert=cert_string))
    signature = self.signData(self.key, payload, digest)
    request_url = '%s/crt/revoke' % self.ca_url
    data = {'digest': digest, 'payload': payload, 'signature': signature}

    self.logger.info("Sent Certificate revocation request for CN: %s." % (
                      cert.get_subject().CN))

    response = self._request('put', request_url, data=data)

    while (not response or response.status_code != 201) and retry < self.max_retry:

      self.logger.error("%s: Failed to send Rovocation request. \n%s" % (
          response.status_code, response.text))
      self.logger.info("will retry in %s seconds..." % sleep_time)
      time.sleep(sleep_time)
      retry += 1
      response = self._request('put', request_url, data=data)

    if response.status_code != 201:
      raise Exception("ERROR: failed to post revoke request after %s retry. Exiting..." % retry)

    self.logger.info("Certificate revocation request for %s successfully sent." % (
                      cert_file))
