/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package br.gov.serpro.cert;

import br.gov.serpro.setup.Setup;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.security.Provider;
import java.security.ProviderException;
import java.security.Security;
import java.security.cert.CertificateFactory;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.Map;
import sun.security.pkcs11.wrapper.CK_ATTRIBUTE;
import sun.security.pkcs11.wrapper.CK_C_INITIALIZE_ARGS;
import sun.security.pkcs11.wrapper.Functions;
import sun.security.pkcs11.wrapper.PKCS11;
import sun.security.pkcs11.wrapper.PKCS11Exception;
import static sun.security.pkcs11.wrapper.PKCS11Constants.*;

//TODO: Deal with wildcards for environments variables.

/**
 *
 * @author esa
 */
class Token{

    private final Setup setup;
    private String name;
    private String libraryPath;
    private Provider tokenProvider;
    private boolean registered = false;
    private long slot;

    static long CK_OBJECT_CLASS;
    static long CK_OBJECT_HANDLE;

    private Token(final Setup setup) {
        this.setup = setup;
    }

    Token(String name, String libraryPath, final Setup setup){
        this(setup);
        this.setName(name);
        this.setLibraryPath(libraryPath);
    }

    public boolean isRegistered() {
        return this.registered;
    }

    public void setLibraryPath(String libraryPath) {
        this.libraryPath = libraryPath;
    }

    public void setName(String name) {
        this.name = name;
    }

    public String getName() {
        return this.name;
    }

    public String getProviderName(){
        return this.tokenProvider.getName();
    }

    protected void registerToken(long slot) throws IOException{

        this.slot = slot;
        String tokenConfiguration = new String("name = " + name + "_" + slot + "\n" +
            "library = " + libraryPath + "\nslot = " + slot +
            "\ndisabledMechanisms = {\n" + "CKM_SHA1_RSA_PKCS\n}" +
            "\n");

        try{
            this.registered = false;
            if (libraryExists()){
                Provider pkcs11Provider = new sun.security.pkcs11.SunPKCS11(new ByteArrayInputStream(tokenConfiguration.getBytes()));
                this.tokenProvider = pkcs11Provider;
                if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                    System.out.println("Adding provider: "+pkcs11Provider.getName());
                    System.out.println("Provider info: " + pkcs11Provider.getInfo());
                }
                Security.addProvider(pkcs11Provider);
                this.setName(this.tokenProvider.getName());
                this.registered = true;
            }
        }
        catch (ProviderException e){
            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                e.printStackTrace();
                System.out.println("Não foi possível inicializar o seguinte token: " + tokenConfiguration);
            }
        }
    }

    protected void unregisterToken(){
        Security.removeProvider(this.tokenProvider.getName());
        this.registered = false;
    }

    Map<String, String> getAliases() throws IOException{

        if (setup.getParameter("debug").equalsIgnoreCase("true")) {
            System.out.println("Getting Aliases");
        }

        Map<String, String> aliases = new HashMap<String, String>();

        CK_C_INITIALIZE_ARGS initArgs = new CK_C_INITIALIZE_ARGS();
        String functionList = "C_GetFunctionList";

        initArgs.flags = CKF_OS_LOCKING_OK;

        PKCS11 tmpPKCS11 = null;
        try {
            try {
                tmpPKCS11 = PKCS11.getInstance(libraryPath, functionList, initArgs, false);
            } catch (IOException ex) {
                if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                    Logger.getLogger(TokenCollection.class.getName()).log(Level.SEVERE, null, ex);
                }
                throw ex;
            }
        } catch (PKCS11Exception e) {
            try {
                initArgs = null;
                tmpPKCS11 = PKCS11.getInstance(libraryPath, functionList, initArgs, true);
            } catch (IOException ex) {
                if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                    Logger.getLogger(TokenCollection.class.getName()).log(Level.SEVERE, null, ex);
                }
            } catch (PKCS11Exception ex) {
                if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                    Logger.getLogger(TokenCollection.class.getName()).log(Level.SEVERE, null, ex);
                }
            }
        }

        try {
            // cria sessão pública rw. com flag CKF_SERIAL_SESSION
            long session = tmpPKCS11.C_OpenSession(this.slot, CKF_SERIAL_SESSION, null, null);

            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                System.out.println("session number:  "+session);
            }

            // TODO: Verifica se está logado, senão loga usuário. Pede pin? ou recebe pin?

            CK_ATTRIBUTE[] TEMPLATE_CERTIFICATE = {new CK_ATTRIBUTE(CKA_CLASS, CKO_CERTIFICATE)};
            CK_ATTRIBUTE[] TEMPLATE_PKEY = {new CK_ATTRIBUTE(CKA_CLASS, CKO_PRIVATE_KEY)};
            CK_ATTRIBUTE[] TEMPLATE_KEY_LABEL_ID = {new CK_ATTRIBUTE(CKA_LABEL), new CK_ATTRIBUTE(CKA_ID)};
            CK_ATTRIBUTE[] TEMPLATE_CERT_LABEL_ID = {
                new CK_ATTRIBUTE(CKA_LABEL),
                new CK_ATTRIBUTE(CKA_ID),
                new CK_ATTRIBUTE(CKA_VALUE)
            };

            tmpPKCS11.C_FindObjectsInit(session, TEMPLATE_CERTIFICATE);
            long[] certs = tmpPKCS11.C_FindObjects(session, 20);

            tmpPKCS11.C_FindObjectsFinal(session);

            tmpPKCS11.C_FindObjectsInit(session, TEMPLATE_PKEY);
            long[] keys = tmpPKCS11.C_FindObjects(session, 20);

            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                System.out.println("Private Keys: "+keys.length);
            }

            for (long key : keys){
                tmpPKCS11.C_GetAttributeValue(session, key, TEMPLATE_KEY_LABEL_ID);

                if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                    System.out.print("Private key ID: ");
                    for (byte b : (byte [])TEMPLATE_KEY_LABEL_ID[1].pValue){
                        System.out.print(b);
                    }
                    System.out.println();
                    if (TEMPLATE_KEY_LABEL_ID[0].pValue != null)
                    {
                        System.out.println("Private key LABEL: "+new String((char [])TEMPLATE_KEY_LABEL_ID[0].pValue));
                    }
                    System.out.println("\nCerts:");
                }

                for (long cert : certs){
                    tmpPKCS11.C_GetAttributeValue(session, cert, TEMPLATE_CERT_LABEL_ID);

                    if (Functions.equals((byte [])TEMPLATE_KEY_LABEL_ID[1].pValue, 
                            (byte [])TEMPLATE_CERT_LABEL_ID[1].pValue)){
                        if (TEMPLATE_CERT_LABEL_ID[0].pValue != null)
                        {
                            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                                System.out.println("Certificate LABEL: "+new String((char [])TEMPLATE_CERT_LABEL_ID[0].pValue));
                            }
                            ByteArrayInputStream in = new ByteArrayInputStream((byte [])TEMPLATE_CERT_LABEL_ID[2].pValue);
                            CertificateFactory cf = CertificateFactory.getInstance("X.509");
                            X509Certificate certObj = (X509Certificate)cf.generateCertificate(in);
                            if (certObj.getBasicConstraints() == -1 ){
                                aliases.put(new String((char [])TEMPLATE_CERT_LABEL_ID[0].pValue),
                                        certObj.getSubjectX500Principal().getName());
                            }
                        }
                    }

                }
                if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                    System.out.println();
                }
            }

            tmpPKCS11.C_CloseSession(session);

        } catch (PKCS11Exception ex) {
            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                Logger.getLogger(TokenCollection.class.getName()).log(Level.SEVERE, null, ex);
            }
        } catch (Throwable t) {
            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                Logger.getLogger(TokenCollection.class.getName()).log(Level.SEVERE, null, t);
            }
        }

        return aliases;

    }

    public boolean libraryExists(){

        File libraryFile = new File(libraryPath);
        if (libraryFile.exists()){
            if (setup.getParameter("debug").equalsIgnoreCase("true")) {
                System.out.println("Arquivo " + libraryPath + " existe.");
            }
            return true;
        }
        
        if (setup.getParameter("debug").equalsIgnoreCase("true")) {
            System.out.println("Biblioteca do Token/SmartCard " + name + " não foi encontrada: " + libraryPath);
        }
        
        return false;
    }
    
}