import os, sys
from time import strptime
from datetime import datetime

from BaDCA.Utils import getSHA1

import cert

from BaDCA import CSRs, Keys

class Certificate:
    filename = None
    cert = None
    serial = None
    key = None
    keyPaths = []
    notBefore = None
    notAfter = None
    modulus = None

    def __init__(self, filename = None, obj = None, keyPaths = None):
        if keyPaths is not None:
            self.addKeyPaths(keyPaths)
        if filename is not None:
            self.readFromFile(filename)
        if obj is not None:
            self.cert = obj
            self.processCertificate()

    def readFromFile(self, filename = None):
        if filename is None:
            return 0
        self.cert = cert.read(filename);
        if self.cert:
            self.processCertificate()
            return 1
        return 0

    def isValid(self):
        if self.cert is None or self.serial is None or \
           self.notBefore is None or self.notAfter is None or \
           self.modulus is None:
            return False
        return True

    # Match serials
    # TODO - should we also match the modulus?
    def __eq__(self, x):
        if self.getSerial() == x.getSerial():
            return 1
        return 0

    def getSerial(self):
        if self.serial is None:
            self.serial = cert.getSerial(self.cert)
        return self.serial

    def getKey(self):
        if self.cert is None or self.key is None:
            return None
        return self.key

    def addKeyPaths(self, paths):
        for p in paths:
            self.addKeyPath(p)

    # should we store the path in the certificate object?
    def addKeyPath(self, path):
        if not path in self.keyPaths:
            self.keyPaths.append(path)
        if self.key is not None:
            self.key.addSearchDirectory(path)

    # convenience function
    def getFingerprint(self):
        return self.getInformation('fingerprint', 'value')

    # which is the purpose to check, ca whether it is a CA purpose
    def checkPurpose(self, which, ca = 0):
        v = None
        if ca == 0:
            try:
                v = self.info['purpose']['NonCA'][which]
            except:
                print "unable to find '%s'" % which
                return 0
        else:
            try:
                v = self.info['purpose']['CA'][which]
            except:
                print "unable to find '%s'" % which
                return 0
        if v == 'Yes':
            return 1
        return 0

    def getInformation(self, section, part = None):
        if self.cert is None or len(self.info) == 0:
            return None
        if part is None:
            try:
                return self.info[section]
            except:
                return None
        try:
            return self.info[section][part]
        except:
            return None

    def asTime(self, which):
        try:
            dt = self.info[which]
            try:
                return datetime.strptime(dt, "%b %d %H:%M:%S %Y %Z")
            except AttributeError:
                try:
                    import time
                    return datetime(*time.strptime(dt, \
                                           "%b %d %H:%M:%S %Y %Z")[0:6])
                except:
                    return None
            except ValueError:
                return None        
        except KeyError:
            return None

    # Compare against supplied time (which should be in UTC)
    def checkUTC(self, dt):
        if self.notBefore is None or self.notAfter is None:
            return False
        if self.notBefore <= dt and self.notAfter >= dt:
            return True
        return False

    # Compare against UTC now
    def checkNow(self):
        dt = datetime.utcnow()
        if self.notBefore is None or self.notAfter is None:
            return False
        if self.notBefore <= dt and self.notAfter >= dt:
            return True
        return False

    # Collect all information from the certificate we need.
    # Try and find the key that was used to create the certificate.
    def processCertificate(self):
        if self.cert is None:
            return
        self.info = cert.parse(self.cert)
        self.notBefore = self.asTime("notBefore")
        self.notAfter = self.asTime("notAfter")
        self.modulus = cert.getModulus(self.cert)
        self.serial = cert.getSerial(self.cert)
        key = Keys.RSAKey(searchPath = self.keyPaths, \
                          certificate = self.cert)
        if key.isValid() and key.modulus == self.modulus:
            self.key = key

    def signRequest(self, csr = None, options = None):
        if csr is None or self.cert is None or not self.key.hasPrivate():
            return None
        if options is None:
            nCert = cert.signRequest(self.cert, csr.csr, self.key.privRSA)
        else:
            nCert = cert.signRequest(self.cert, csr.csr, self.key.privRSA, \
                                     options)
        thecert = Certificate(obj = nCert)
        if thecert is not None:
            return thecert
        return None

    def createSelfSignedFromDict(self, subject, key = None, \
                                 options = None):
        if key is None:
            if self.key is None:
                return None
            thekey = self.key
        else:
            thekey = key
        if not thekey.hasPrivate():
            return None
        self.cert = cert.createCertificateFromDict(subject, \
                                                thekey.privRSA, options)
        if self.cert is not None:
            self.key = thekey
            self.processCertificate()

    def writeToFile(self, filename = None):
        if self.cert is None:
            return 0
        if filename is None:
            filename = self.getSerial() + '.pem'
        rv = cert.write(self.cert, filename)
        if rv == 1 and self.filename is not None:
            self.filename = filename
        return rv

    def writeToDirectory(self, dirname):
        if self.cert is None:
            return 0
        if not os.path.isdir(dirname):
            return 0
        self.filename = os.path.abspath( \
                       os.path.join(dirname, self.getSerial() + '.pem'))
        return cert.write(self.cert, self.filename)

    def checkKey(self, key):
        if self.key is not None:
            if key == self.key:
                return True
            if self.modulus == self.key.modulus:
                self.key = key
                return True            
        return False