package com.netscape.pkisilent.common;

// --- BEGIN COPYRIGHT BLOCK ---
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; version 2 of the License.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//
// (C) 2007 Red Hat, Inc.
// All rights reserved.
// --- END COPYRIGHT BLOCK ---

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.net.URL;
import java.net.URLConnection;
import java.security.KeyPair;

import netscape.security.x509.X500Name;

import org.mozilla.jss.CryptoManager;
import org.mozilla.jss.asn1.ASN1Util;
import org.mozilla.jss.asn1.BIT_STRING;
import org.mozilla.jss.asn1.INTEGER;
import org.mozilla.jss.asn1.OBJECT_IDENTIFIER;
import org.mozilla.jss.asn1.OCTET_STRING;
import org.mozilla.jss.asn1.SEQUENCE;
import org.mozilla.jss.crypto.AlreadyInitializedException;
import org.mozilla.jss.crypto.CryptoStore;
import org.mozilla.jss.crypto.CryptoToken;
import org.mozilla.jss.crypto.IVParameterSpec;
import org.mozilla.jss.crypto.InternalCertificate;
import org.mozilla.jss.crypto.KeyGenAlgorithm;
import org.mozilla.jss.crypto.KeyGenerator;
import org.mozilla.jss.crypto.KeyPairAlgorithm;
import org.mozilla.jss.crypto.KeyPairGenerator;
import org.mozilla.jss.crypto.KeyWrapAlgorithm;
import org.mozilla.jss.crypto.KeyWrapper;
import org.mozilla.jss.crypto.PrivateKey;
import org.mozilla.jss.crypto.SymmetricKey;
import org.mozilla.jss.crypto.X509Certificate;
import org.mozilla.jss.pkcs11.PK11Token;
import org.mozilla.jss.pkix.crmf.CertReqMsg;
import org.mozilla.jss.pkix.crmf.CertRequest;
import org.mozilla.jss.pkix.crmf.CertTemplate;
import org.mozilla.jss.pkix.crmf.EncryptedKey;
import org.mozilla.jss.pkix.crmf.EncryptedValue;
import org.mozilla.jss.pkix.crmf.PKIArchiveOptions;
import org.mozilla.jss.pkix.crmf.POPOPrivKey;
import org.mozilla.jss.pkix.crmf.ProofOfPossession;
import org.mozilla.jss.pkix.primitive.AVA;
import org.mozilla.jss.pkix.primitive.AlgorithmIdentifier;
import org.mozilla.jss.pkix.primitive.Name;
import org.mozilla.jss.pkix.primitive.SubjectPublicKeyInfo;
import org.mozilla.jss.util.Password;

import com.netscape.osutil.OSUtil;

/**
 * CMS Test framework .
 * Use this class to initalize,add a certificate ,generate a certificate request from certificate database.
 */

public class ComCrypto {

    private String cdir, certnickname, keysize, keytype, tokenpwd;
    private String certpackage, pkcs10request;
    private boolean debug = true;
    private boolean DBlogin = false;
    private boolean generaterequest = false;

    private String transportcert = null;
    private boolean dualkey = false;
    public String CRMF_REQUEST = null;
    int START = 1;
    int END = START + 1;
    Password password = null;

    public static CryptoManager manager;
    public static CryptoToken token;
    private CryptoStore store;
    private Password pass1 = null, pass2 = null;

    private String bstr = "-----BEGIN NEW CERTIFICATE REQUEST-----";
    private String blob, Blob1 = null;
    private String Blob2 = null;
    private String estr = "-----END NEW CERTIFICATE REQUEST-----";

    private String certprefix = null;

    public ComCrypto() {
    };

    /**
     * Constructor . Takes the parameter certificatedbdirectory , passwordfor cert database,
     * certificatenickname,keysize, keytype(RSA/DSA)
     * 
     * @param certdbdirectory.
     * @param certdbpassword
     * @param certnickname
     * @param keysize (1024/2048/4096)
     * @param keytype (RSA/DSA)
     */

    public ComCrypto(String cd, String tpwd, String cn, String ks, String kt) {
        cdir = cd;
        tokenpwd = tpwd;
        certnickname = cn;
        keysize = ks;
        keytype = kt;
    }

    // Set and Get functions 

    public void setCertDir(String cd) {
        cdir = cd;
    }

    public void setCertnickname(String cd) {
        certnickname = cd;
    }

    public void setKeySize(String cd) {
        keysize = cd;
    }

    public void setKeyType(String cd) {
        keytype = cd;
    }

    public void setTokenPWD(String cd) {
        tokenpwd = cd;
    }

    public void setCertPackage(String cd) {
        certpackage = cd;
    }

    public void setGenerateRequest(boolean c) {
        generaterequest = c;
    }

    public void setDebug(boolean t) {
        debug = t;
    }

    public void setCertPrefix(String prefix) {
        certprefix = prefix;
    }

    /*
     * setTransportCert() should only be called when the calling profile
     * needs to do key archivals with the DRM and make sure the function
     * generateCRMFtransport() is called for the CRMF request generation
     * part.
     */
    public void setTransportCert(String tcert) {
        transportcert = tcert;
    }

    public void setDualKey(boolean dkey) {
        dualkey = dkey;
    }

    public String getPkcs10Request() {
        return pkcs10request;
    }

    /**
     * Parses the Certificate and returns SubjectDN . Takes certificate as parameter
     */

    public String getCertificateString(X509Certificate cert) {
        if (cert == null) {
            return null;
        }

        // note that it did not represent a certificate fully
        return cert.getVersion() + ";" + cert.getSerialNumber().toString() + ";"
                + cert.getIssuerDN() + ";" + cert.getSubjectDN();
    }

    /**
     * Finds and returns Certificate . Takes certificatenickname as parameter.
     */

    public X509Certificate findCert(String certname) {
        try {

            X509Certificate cert2 = manager.findCertByNickname(certname);

            return cert2;

        } catch (Exception e) {
            System.out.println("exception importing cert " + e.getMessage());
            return null;
        }

    }

    /**
     * Imports a certificate to Certificate Database. Takes certificate and nickname as parameters.
     */

    public boolean importCert(X509Certificate xcert, String nickname) {
        try {

            System.out.println(
                    "importCert x509 : importing with nickname: " + nickname);

            InternalCertificate cert2 = manager.importCertToPerm(xcert, nickname);

            cert2.setSSLTrust(2);
            return true;

        } catch (Exception e) {
            System.out.println("exception importing cert " + e.getMessage());
            return false;
        }

    }

    /**
     * Imports a certificate to Certificate Database. Takes certificate and nickname as parameters.
     */

    public boolean importCert(String cpack, String cn) {

        System.out.println("importCert string: importing with nickname: " + cn);
        try {

            String tmp = normalize(cpack);

            if (DBlogin) {
                System.out.println("Already logged into to DB");
            }

            if (manager == null) {
                System.out.println("Manager object is null");
            }

            X509Certificate cert = manager.importCertPackage(tmp.getBytes(), cn);

            return true;

        } catch (Exception e) {
            System.out.println(
                    "ERROR:exception importing cert " + e.getMessage());
            e.printStackTrace();
            return false;
        }

    }

    /* imports CA certificate
     */

    public boolean importCACert(String cpack) {

        try {
            String tmp = normalize(cpack);

            if (DBlogin) {
                System.out.println("Already logged into to DB");
            }

            if (manager == null) {
                System.out.println("Manager object is null");
            }

            X509Certificate cert = manager.importCACertPackage(tmp.getBytes());

            return true;

        } catch (Exception e) {
            System.out.println(
                    "ERROR:exception importing cert " + e.getMessage());
            return false;
        }

    }

    /**
     * Normalizes a given certificate string . Removes the extra \\ in the certificate returned by CMS server.
     */

    public String normalize(String s) {

        String val = "";

        for (int i = 0; i < s.length(); i++) {
            if ((s.charAt(i) == '\\') && (s.charAt(i + 1) == 'n')) {
                val += '\n';
                i++;
                continue;
            } else if ((s.charAt(i) == '\\') && (s.charAt(i + 1) == 'r')) {
                i++;
                continue;
            } else if (s.charAt(i) == '"') {
                continue;
            }
            val += s.charAt(i);
        }
        return val;
    }

    /**
     * Normalizes a given certificate string . Removes the extra \\ in the certificate returned by CMS server.
     */

    public String normalizeForLDAP(String s) {

        String val = "";

        for (int i = 0; i < s.length(); i++) {
            if ((s.charAt(i) == '\\') && (s.charAt(i + 1) == 'n')) {
                val += '\n' + " ";
                i++;
                continue;
            } else if ((s.charAt(i) == '\\') && (s.charAt(i + 1) == 'r')) {
                i++;
                continue;
            } else if (s.charAt(i) == '"') {
                continue;
            }
            val += s.charAt(i);
        }
        return val;
    }

    /**
     * Convert to pkcs7 format
     */

    public String pkcs7Convertcert(String s) {

        String val = "";

        int len = s.length();

        for (int i = 0; i < len; i = i + 64) {

            if (i + 64 < len) {
                val = val + s.substring(i, i + 64) + "\n";
            } else {
                val = val + s.substring(i, len);
            }

        }
        return val;
    }

    /**
     * Delete all keys frim key3.db
     **/

    public void deleteKeys() {
        try {
            int i = 0;

            store = token.getCryptoStore();
            PrivateKey[] keys = store.getPrivateKeys();

            if (debug) {
                System.out.println("Now we shall delete all the keys!");
            }

            keys = store.getPrivateKeys();
            for (i = 0; i < keys.length; i++) {
                PrivateKey key = (PrivateKey) keys[i];

                store.deletePrivateKey(key);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * Creates a new certificate database
     **/

    public boolean CreateCertDB() {
        return loginDB();

    }

    /**
     * Login to cert database
     **/

    public boolean loginDB() {
        Password pass1 = null;

        try {
            if (debug) {
                System.out.println("CRYPTO INIT WITH CERTDB:" + cdir);
            }

            // this piece of code is to create db's with certain prefix
            if (certprefix != null) {
                CryptoManager.InitializationValues vals;

                vals = new CryptoManager.InitializationValues(cdir, certprefix,
                        certprefix, "secmod.db");
                CryptoManager.initialize(vals);
            } else {
                CryptoManager.initialize(cdir);
            }

            manager = CryptoManager.getInstance();
            token = (PK11Token) manager.getInternalKeyStorageToken();
            pass1 = new Password(tokenpwd.toCharArray());
            if (token.isLoggedIn() && debug) {
                System.out.println("Already Logged in ");
            }

            if (debug) {
                System.out.println("tokenpwd:" + tokenpwd);
            }

            token.login(pass1);
            pass1.clear();

        } catch (AlreadyInitializedException e) {
            if (debug) {
                System.out.println("Crypto manager already initialized");
            }
        } catch (Exception e) {
            try {
                if (!token.isLoggedIn()) {
                    token.initPassword(pass1, pass1);
                }
                return true;
            } catch (Exception er) {
                System.err.println("some exception:" + e);
                return false;
            }
        }
        DBlogin = true;
        return true;
    }

    /**
     * Generate Certificate Request
     **/

    public synchronized boolean generateRequest() {

        System.out.println("generating pkcs10 Request");
        loginDB();

        try {
            debug = true;
            System.out.println("Generating request : keysize :" + keysize);
            System.out.println("Generating request : subject :" + certnickname);
            System.out.println("Generating request : keytype :" + keytype);

            Integer n = new Integer(keysize);

            if (generaterequest) {
                blob = token.generateCertRequest(certnickname, n.intValue(),
                        keytype, (byte[]) null, (byte[]) null, (byte[]) null);

                System.out.println("Cert Request Generated.");

                bstr = "-----BEGIN NEW CERTIFICATE REQUEST-----";
                Blob1 = blob.substring(bstr.length() + 1);
                Blob2 = Blob1.substring(0, Blob1.indexOf(estr));

                System.out.println(Blob2);
                pkcs10request = Blob2;
            }

            return true;

        } catch (Exception e) {
            System.out.println("Exception: Unable to generate request: " + e);
        }

        return false;
    }

    public String generateCRMFrequest() {
        URL url = null;
        URLConnection conn = null;
        InputStream is = null;
        BufferedReader reader = null;
        boolean success = false;
        int num = 1;
        long total_time = 0;
        KeyPair pair = null;

        System.out.println("Debug : initialize crypto Manager");
        try {

            // Step 1. initialize crypto Manager
            try {
                CryptoManager.initialize(cdir);
            } catch (Exception e) {
                // it is ok if it is already initialized 
                System.out.println("INITIALIZATION ERROR: " + e.toString());
                System.out.println("cdir = " + cdir);
            }

            // Step 2 log into database 
            try {

                System.out.println("Debug : before getInstance");

                manager = CryptoManager.getInstance();
                String token_pwd = tokenpwd;

                System.out.println("Debug : before get token");

                token = manager.getInternalKeyStorageToken();
                password = new Password(token_pwd.toCharArray());

                System.out.println("Debug : before login password");

                token.login(password);

                System.out.println("Debug : after login password");
            } catch (Exception e) {
                System.out.println("INITIALIZATION ERROR: " + e.toString());

                if (!token.isLoggedIn()) {
                    token.initPassword(password, password);
                }
            }

            // Generating CRMF request 

            KeyPairGenerator kg = token.getKeyPairGenerator(KeyPairAlgorithm.RSA);

            Integer x = new Integer(keysize);
            int key_len = x.intValue();

            kg.initialize(key_len);

            // 1st key pair
            pair = kg.genKeyPair();

            // create CRMF
            CertTemplate certTemplate = new CertTemplate();

            certTemplate.setVersion(new INTEGER(2));

            if (certnickname != null) {
                X500Name name = new X500Name(certnickname);
                ByteArrayInputStream cs = new ByteArrayInputStream(name.getEncoded());
                Name n = (Name) Name.getTemplate().decode(cs);
                certTemplate.setSubject(n);
            }

            certTemplate.setPublicKey(new SubjectPublicKeyInfo(pair.getPublic()));

            SEQUENCE seq = new SEQUENCE();
            CertRequest certReq = new CertRequest(new INTEGER(1), certTemplate,
                    seq);
            byte popdata[] = { 0x0, 0x3, 0x0 };

            ProofOfPossession pop = ProofOfPossession.createKeyEncipherment(
                    POPOPrivKey.createThisMessage(new BIT_STRING(popdata, 3)));

            CertReqMsg crmfMsg = new CertReqMsg(certReq, pop, null);

            SEQUENCE s1 = new SEQUENCE();

            // 1st : Encryption key 

            s1.addElement(crmfMsg);

            // 2nd : Signing Key

            if (dualkey) {
                System.out.println("dualkey = true");
                SEQUENCE seq1 = new SEQUENCE();
                CertRequest certReqSigning = new CertRequest(new INTEGER(1),
                        certTemplate, seq1);
                CertReqMsg signingMsg = new CertReqMsg(certReqSigning, pop, null);

                s1.addElement(signingMsg);
            }

            byte encoded[] = ASN1Util.encode(s1);

            // BASE64Encoder encoder = new BASE64Encoder(); 
            // String Req1 = encoder.encodeBuffer(encoded);
            String Req1 = OSUtil.BtoA(encoded);

            // Set CRMF_REQUEST variable 
            CRMF_REQUEST = Req1;

            System.out.println("CRMF_REQUEST = " + CRMF_REQUEST);

        } catch (Exception e) {
            System.out.println("ERROR: " + e.toString());
            e.printStackTrace();
            return null;
        }

        return CRMF_REQUEST;
    }

    /*
     * This function is used to Generated CRMF requests wrapped with the 
     * transport cert so that we can do key archival with the drm.
     * This function expects transportcert variable to be set in this class.
     * Use setTransportCert() to do the same.
     */

    public String generateCRMFtransport() {

        boolean success = false;
        int num = 1;
        long total_time = 0;
        KeyPair pair = null;

        try {
            // Step 1. initialize crypto Manager
            try {
                CryptoManager.initialize(cdir);
            } catch (Exception e) {
                // it is ok if it is already initialized 
                System.out.println("INITIALIZATION ERROR: " + e.toString());
                System.out.println("cdir = " + cdir);
            }

            // Step 2 log into database 
            try {

                System.out.println("Debug : before getInstance");

                manager = CryptoManager.getInstance();
                String token_pwd = tokenpwd;

                System.out.println("Debug : before get token");

                token = manager.getInternalKeyStorageToken();
                password = new Password(token_pwd.toCharArray());

                System.out.println("Debug : before login password");

                token.login(password);

                System.out.println("Debug : after login password");
            } catch (Exception e) {
                System.out.println("INITIALIZATION ERROR: " + e.toString());

                if (!token.isLoggedIn()) {
                    token.initPassword(password, password);
                }
            }

            // Key Pair Generation
            KeyPairGenerator kg = token.getKeyPairGenerator(KeyPairAlgorithm.RSA);
            Integer x = new Integer(keysize);
            int key_len = x.intValue();

            kg.initialize(key_len);

            pair = kg.genKeyPair();

            // wrap private key
            // BASE64Decoder decoder = new BASE64Decoder();
            // byte transport[] = decoder.decodeBuffer(transportcert);
            byte transport[] = OSUtil.AtoB(transportcert);

            X509Certificate tcert = manager.importCACertPackage(transport);

            byte iv[] = { 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1 };

            KeyGenerator kg1 = token.getKeyGenerator(KeyGenAlgorithm.DES3);
            SymmetricKey sk = kg1.generate();

            // wrap private key using session
            KeyWrapper wrapper1 = token.getKeyWrapper(
                    KeyWrapAlgorithm.DES3_CBC_PAD);

            wrapper1.initWrap(sk, new IVParameterSpec(iv));

            byte key_data[] = wrapper1.wrap((
                    org.mozilla.jss.crypto.PrivateKey) pair.getPrivate());

            // wrap session using transport
            KeyWrapper rsaWrap = token.getKeyWrapper(KeyWrapAlgorithm.RSA);

            rsaWrap.initWrap(tcert.getPublicKey(), null);

            byte session_data[] = rsaWrap.wrap(sk);

            // create CRMF
            CertTemplate certTemplate = new CertTemplate();

            certTemplate.setVersion(new INTEGER(2));

            if (certnickname != null) {
                X500Name name = new X500Name(certnickname);
                ByteArrayInputStream cs = new ByteArrayInputStream(name.getEncoded());
                Name n = (Name) Name.getTemplate().decode(cs);
                certTemplate.setSubject(n);
            }

            certTemplate.setPublicKey(new SubjectPublicKeyInfo(pair.getPublic()));

            // set extension
            AlgorithmIdentifier algS = new AlgorithmIdentifier(
                    new OBJECT_IDENTIFIER("1.2.840.113549.3.7"),
                    new OCTET_STRING(iv));

            EncryptedValue encValue = new EncryptedValue(null, algS,
                    new BIT_STRING(session_data, 0), null, null,
                    new BIT_STRING(key_data, 0));

            EncryptedKey key = new EncryptedKey(encValue);
            PKIArchiveOptions opt = new PKIArchiveOptions(key);

            SEQUENCE seq = new SEQUENCE();

            seq.addElement(
                    new AVA(new OBJECT_IDENTIFIER("1.3.6.1.5.5.7.5.1.4"), opt));

            CertRequest certReq = new CertRequest(new INTEGER(1), certTemplate,
                    seq);

            // Adding proof of possesion data
            byte popdata[] = { 0x0, 0x3, 0x0 };
            ProofOfPossession pop = ProofOfPossession.createKeyEncipherment(
                    POPOPrivKey.createThisMessage(new BIT_STRING(popdata, 3)));

            CertReqMsg crmfMsg = new CertReqMsg(certReq, pop, null);

            SEQUENCE s1 = new SEQUENCE();

            // 1st : Encryption key 
            s1.addElement(crmfMsg);

            // 2nd : Signing Key

            if (dualkey) {
                System.out.println("dualkey = true");
                SEQUENCE seq1 = new SEQUENCE();
                CertRequest certReqSigning = new CertRequest(new INTEGER(1),
                        certTemplate, seq1);
                CertReqMsg signingMsg = new CertReqMsg(certReqSigning, pop, null);

                s1.addElement(signingMsg);
            }

            byte encoded[] = ASN1Util.encode(s1);

            // BASE64Encoder encoder = new BASE64Encoder();

            // CRMF_REQUEST = encoder.encodeBuffer(encoded);
            CRMF_REQUEST = OSUtil.BtoA(encoded);

            System.out.println("Generated crmf request: ...... ");
            System.out.println("");

            System.out.println(CRMF_REQUEST);
            System.out.println("");
            System.out.println("End crmf Request:");
        } catch (Exception e) {
            System.out.println("Exception: " + e.getMessage());
        }

        return CRMF_REQUEST;
    }

} // end of class