# -*- coding: utf-8 -*-
import os
import sys
import subprocess
import time
import ConfigParser
import uuid
import ssl
import logging
import cgi, errno
import urlparse
from OpenSSL import crypto, SSL
import requests
import traceback


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 CertificateBase(object):

  def __init__(self):
    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 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)
    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)
      os.chmod(output_file, 0640)

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

    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)
    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 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:
      #with open(key_file) as f_key:
      #  key = f_key.read()
      return self.validateCertAndKey(cert_file, key_file)
    return True


  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 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

class CertificateAuthority(CertificateBase):

  def __init__(self, openssl_binary,
      openssl_configuration=None, certificate='cacert.pem',
      key='cakey.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()
    
    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 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])


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

    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/get/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)

    with open(self.cacertificate, 'w') as f:
      f.write(response.text)

  def signCertificateWeb(self, csr_file):

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

    #req = self.freadCertificateRequest(csr_file)
    #common_name = self.getCN(req)

    data = {'csr': csr}
    retry = 0
    sleep_time = 10
    request_url = '%s/request' % 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, key is: %s" % csr_key)
    else:
      response = self._request('post', request_url, data=data)

      while (not response or response.status_code != 200) and retry < self.max_retry:
  
        self.logger.error("%s: Failed to send 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('post', request_url, data=data)
  
      if response.status_code != 200:
        raise Exception("ERROR: failed to post CSR after % retry. Exiting..." % max_retry)
  
      self.logger.info("CSR succefully sent.")
      self.logger.debug("Server reponse with csr key is %s" % response.text)
      csr_key = response.text
      with open(csr_key_file, 'w') as fkey:
        fkey.write(response.text)

    self.logger.info("Waiting for signed certificate...")

    reply_url = '%s/get/%s.cert.pem' % (self.ca_url, csr_key)
    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)

    os.chmod(cert_temp, 0640)
    self.logger.info("Validating signed certificate...")

    if self.checkCertificateValidity(self.cacertificate, cert_temp, self.key):
      os.rename(cert_temp, self.certificate)
    else:
      raise Exception("Error: Certificate validation failed. " \
          "This signed certificate should be revoked!")
    self.logger.info("Certificate correctly saved at %s." % self.certificate)
