# -----------------------------------------------------------------------------
# Utility methods for Cert Store manipulation
# -----------------------------------------------------------------------------

from java.io import FileInputStream, FileReader
from java.lang import String
from java.security import KeyFactory, KeyStore
from esapi import EntityStoreAPI
from com.vordel.security.openssl import PKCS12
from com.vordel.es.util import ShorthandKeyFinder
from org.bouncycastle.openssl import PasswordFinder, PEMReader

# -----------------------------------------------------------------------------
# Password finder class
# -----------------------------------------------------------------------------
class MyPasswordFinder(PasswordFinder):
    
    password = None
    
    def __init__(self, password):
        self.password = password
        
    def getPassword(self):
        return String(self.password).toCharArray();
        
# -----------------------------------------------------------------------------
# Class to define utility methods to add gets Certificates and keys from 
# different resources and add them to the Cert Store.
# -----------------------------------------------------------------------------
class CertStoreUtil(object):
    
    # -------------------------------------------------------------------------
    # Utility method to add a Certificate to the Cert Store
    # -------------------------------------------------------------------------
    @classmethod
    def addCertToStore(cls, esapi, alias, cert, privateKey=None):
        # Gets the Cert store using short hand key
        shkf = ShorthandKeyFinder(esapi.es)
        certStore = shkf.getEntity('/[Certificates]name=Certificate Store')
        escapedAlias = shkf.escapeFieldValue(alias)
        shk = '/[Certificates]name=Certificate Store/[Certificate]dname=' + escapedAlias
        
        # See if the certificate alias already exists in the entity store, 
        # if it does then update it thereby preserving any references to any HTTPS interfaces that are using this cert
        certEntity = shkf.getEntity(shk)
        
        if certEntity != None:
            
            # Updates the existing certificate in the certstore
            certEntity.setBinaryValue("content", cert.getEncoded())
            if privateKey != None:
                encrypted = esapi.encryptBytes(privateKey.getEncoded())
                certEntity.setStringField("key", encrypted) 
            else:
                # no private key specified - check if there was a private key in the earlier instance in the entity store - if so, remove it
                entityPKey = certEntity.getStringValue ("key")
                
                if entityPKey is not None and len (entityPKey) > 0:
                    certEntity.removeField ("key")
                
            esapi.updateEntity(certEntity)
            return
        else:
            # Adds the certificate to the certstore 
            certEntity = esapi.createEntity("Certificate")
            certEntity.setStringField("dname", alias)
            certEntity.setBinaryValue("content", cert.getEncoded())
            if privateKey != None:
                encrypted = esapi.encryptBytes(privateKey.getEncoded())
                certEntity.setStringField("key", encrypted)
            esapi.addEntity(certStore, certEntity)

    # -------------------------------------------------------------------------
    # Utility method to get Key from P12 file. 
    # -------------------------------------------------------------------------
    @classmethod 
    def getKeyFromP12(cls, file, password=None):
        io = FileInputStream(file)
        pkcs12 = PKCS12(io)
        if password != None:
            pkcs12.decrypt(String(password).toCharArray());
        key = pkcs12.getKey()
        return key

    # -------------------------------------------------------------------------
    # Utility method to get Certificate from P12 file.
    # -------------------------------------------------------------------------
    @classmethod 
    def getCertFromP12(cls, file, password):
        # Get the Key store
        ks = KeyStore.getInstance("PKCS12")
        io = FileInputStream(file)
        # Loads the KeyStore from the file input stream and uses the password 
        # to unlock the keystore, or to check the integrity of the keystore data. 
        # If the password is not given for integrity checking, then 
        # integrity checking is not performed.         
        ks.load(io, String(password).toCharArray())
        io.close()

        # Lists all the alias names of this keystore. 
        aliases = ks.aliases()

        # Returns the first Certificate that matches an alias name
        while (aliases.hasMoreElements()): 
            alias = aliases.nextElement()
            if (ks.isKeyEntry(alias)):
                return ks.getCertificate(alias)        
        return None
        
    # -------------------------------------------------------------------------
    # Utility method to get private key
    # -------------------------------------------------------------------------
    @classmethod 
    def getKeyFromPEM(cls, file, password=None):
        io = FileInputStream(file)
        # PEMReader is used for reading OpenSSL PEM encoded streams containing 
        # X509 certificates, PKCS8 encoded keys and PKCS7 objects. 
        if password == None:
            reader = PEMReader(FileReader(file))
        else:
            reader = PEMReader(FileReader(file), MyPasswordFinder(password))
        keyPair = reader.readObject();
        if keyPair != None:
            return keyPair.getPrivate(); 
        return None