/*

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code 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

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

import java.io.*;

import java.util.*;

import java.security.KeyStore;

import java.security.KeyStoreException;

import java.security.KeyFactory;

import java.security.KeyPairGenerator;

import java.security.KeyPair;

import java.security.SecureRandom;

import java.security.AuthProvider;

import java.security.PrivateKey;

import java.security.Provider;

import java.security.ProviderException;

import java.security.Signature;

import java.security.Security;

import java.security.cert.*;

import java.security.spec.*;

import java.security.interfaces.*;

import javax.crypto.SecretKey;

import javax.security.auth.Subject;

import javax.security.auth.login.LoginException;

import com.sun.security.auth.module.*;

import com.sun.security.auth.callback.*;

public class Basic extends PKCS11Test {

    private static final char SEP = File.separatorChar;

    private static String DIR = System.getProperty("DIR");

    private static char[] tokenPwd;

    private static final char[] ibuttonPwd =

                        new char[0];

    private static final char[] activcardPwd =

                        new char[] { '1', '1', '2', '2', '3', '3' };

    private static final char[] nssPwd =

                        new char[] { 't', 'e', 's', 't', '1', '2' };

    private static final char[] solarisPwd =

                        new char[] { 'p', 'i', 'n' };

    private static final char[] sca1000Pwd =

                        new char[] { 'p', 'a', 's', 's', 'w', 'o', 'r', 'd' };

    private static final char[] sPwd = { 'f', 'o', 'o' };

    private static SecretKey sk1;

    private static SecretKey sk2;

    private static SecretKey sk3;

    private static SecretKey sk4;

    private static RSAPrivateCrtKey pk1;

    private static PrivateKey pk2;

    private static PrivateKey pk3;

    private static Certificate[] chain1;

    private static Certificate[] chain2;

    private static Certificate[] chain3;

    private static Certificate[] chain4;

    private static X509Certificate randomCert;

    private static KeyStore ks;

    private static final String KS_TYPE = "PKCS11";

    private static Provider provider;

    private static class FooEntry implements KeyStore.Entry { }

    private static class P11SecretKey implements SecretKey {

        String alg;

        int length;

        public P11SecretKey(String alg, int length) {

            this.alg = alg;

            this.length = length;

        }

        public String getAlgorithm() { return alg; }

        public String getFormat() { return "raw"; }

        public byte[] getEncoded() { return new byte[length/8]; }

    }

    public static void main(String[] args) throws Exception {

        main(new Basic(), args);

    }

    public void main(Provider p) throws Exception {

        this.provider = p;

        // get private keys

        KeyFactory dsaKf = KeyFactory.getInstance("DSA", "SUN");

        ObjectInputStream ois1 = new ObjectInputStream

                        (new FileInputStream(new File(DIR, "pk1.key")));

        byte[] keyBytes = (byte[])ois1.readObject();

        ois1.close();

        PrivateKey tmpKey =

                kf.generatePrivate(new PKCS8EncodedKeySpec(keyBytes));

        pk1 = (RSAPrivateCrtKey)tmpKey;

        ObjectInputStream ois2 = new ObjectInputStream

                        (new FileInputStream(new File(DIR, "pk2.key")));

        keyBytes = (byte[])ois2.readObject();

        ois2.close();

        pk2 = kf.generatePrivate(new PKCS8EncodedKeySpec(keyBytes));

        ObjectInputStream ois3 = new ObjectInputStream

                        (new FileInputStream(new File(DIR, "pk3.key")));

        keyBytes = (byte[])ois3.readObject();

        pk3 = kf.generatePrivate(new PKCS8EncodedKeySpec(keyBytes));

        ois3.close();

        // get cert chains for private keys

        CertificateFactory cf = CertificateFactory.getInstance("X.509", "SUN");

        Certificate caCert = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "ca.cert")));

        Certificate ca2Cert = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "ca2.cert")));

        Certificate pk1cert = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "pk1.cert")));

        Certificate pk1cert2 = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "pk1.cert2")));

        Certificate pk2cert = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "pk2.cert")));

        Certificate pk3cert = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "pk3.cert")));

        chain1 = new Certificate[] { pk1cert, caCert };

        chain2 = new Certificate[] { pk2cert, caCert };

        chain3 = new Certificate[] { pk3cert, caCert };

        chain4 = new Certificate[] { pk1cert2, ca2Cert };

        // create secret keys

        sk1 = new P11SecretKey("DES", 64);

        sk2 = new P11SecretKey("DESede", 192);

        sk3 = new P11SecretKey("AES", 128);

        sk4 = new P11SecretKey("RC4", 128);

        // read randomCert

        randomCert = (X509Certificate)cf.generateCertificate

                        (new FileInputStream(new File(DIR, "random.cert")));

        doTest();

    }

    private static void doTest() throws Exception {

        String token = System.getProperty("TOKEN");

        String test = System.getProperty("TEST");

        if (token == null || token.length() == 0) {

            throw new Exception("token arg required");

        }

        if (test == null || test.length() == 0) {

            throw new Exception("test arg required");

        }

        if ("ibutton".equals(token)) {

            tokenPwd = ibuttonPwd;

        } else if ("activcard".equals(token)) {

            tokenPwd = activcardPwd;

        } else if ("nss".equals(token)) {

            tokenPwd = nssPwd;

        } else if ("sca1000".equals(token)) {

            tokenPwd = sca1000Pwd;

        } else if ("solaris".equals(token)) {

            tokenPwd = solarisPwd;

        }

        if ("list".equals(test)) {

            Basic.list();

        } else if ("basic".equals(test)) {

            int testnum = 1;

            if ("ibutton".equals(token)) {

                // pkey and setAttribute

                testnum = Basic.pkey(testnum);

                testnum = Basic.setAttribute(testnum);

            } else if ("activcard".equals(token)) {

                // sign

                testnum = Basic.signAlias(testnum, null);

            } else if ("nss".equals(token)) {

                // setAttribute, pkey, sign

                testnum = Basic.setAttribute(testnum);

                testnum = Basic.pkey(testnum);

                testnum = Basic.sign(testnum);

                testnum = Basic.copy(testnum);

            } else if ("solaris".equals(token)) {

                testnum = Basic.setAttribute(testnum);

                testnum = Basic.pkey(testnum);

                testnum = Basic.sign(testnum);

                testnum = Basic.skey(testnum);

                testnum = Basic.copy(testnum);

            } else if ("sca1000".equals(token)) {

                // setAttribute, pkey, sign, skey, copy

                testnum = Basic.setAttribute(testnum);

                testnum = Basic.pkey(testnum);

                testnum = Basic.sign(testnum);

                testnum = Basic.skey(testnum);

                testnum = Basic.copy(testnum);

            }

        } else if ("pkey".equals(test)) {

            Basic.pkey(1);

        } else if ("skey".equals(test)) {

            Basic.skey(1);

        } else if ("setAttribute".equals(test)) {

            Basic.setAttribute(1);

        } else if ("copy".equals(test)) {

            Basic.copy(1);

        } else if ("sign".equals(test)) {

            Basic.sign(1);

        } else if ("module".equals(test)) {

            Basic.module();

        } else if ("nss-extended".equals(test)) {

            // this only works if NSS_TEST is set to true in P11KeyStore.java

            int testnum = 1;

            testnum = Basic.setAttribute(testnum);

            testnum = Basic.pkey(testnum);

            testnum = Basic.sign(testnum);

            testnum = Basic.extended(testnum);

        } else {

            System.out.println("unrecognized command");

        }

    }

    private static int sign(int testnum) throws Exception {

        if (ks == null) {

            ks = KeyStore.getInstance(KS_TYPE, provider);

            ks.load(null, tokenPwd);

        }

        if (!ks.containsAlias("pk1")) {

            ks.setKeyEntry("pk1", pk1, null, chain1);

        }

        System.out.println("test " + testnum++ + " passed");

        return signAlias(testnum, "pk1");

    }

    private static int signAlias(int testnum, String alias) throws Exception {

        if (ks == null) {

            ks = KeyStore.getInstance(KS_TYPE, provider);

            ks.load(null, tokenPwd);

        }

        if (alias == null) {

            Enumeration enu = ks.aliases();

            if (enu.hasMoreElements()) {

                alias = (String)enu.nextElement();

            }

        }

        PrivateKey pkey = (PrivateKey)ks.getKey(alias, null);

        if ("RSA".equals(pkey.getAlgorithm())) {

            System.out.println("got [" + alias + "] signing key: " + pkey);

        } else {

            throw new SecurityException

                ("expected RSA, got " + pkey.getAlgorithm());

        }

        Signature s = Signature.getInstance("MD5WithRSA", ks.getProvider());

        s.initSign(pkey);

        System.out.println("initialized signature object with key");

        s.update("hello".getBytes());

        System.out.println("signature object updated with [hello] bytes");

        byte[] signed = s.sign();

        System.out.println("received signature " + signed.length +

                        " bytes in length");

        Signature v = Signature.getInstance("MD5WithRSA", ks.getProvider());

        v.initVerify(ks.getCertificate(alias));

        v.update("hello".getBytes());

        v.verify(signed);

        System.out.println("signature verified");

        System.out.println("test " + testnum++ + " passed");

        return testnum;

    }

    private static int copy(int testnum) throws Exception {

        if (ks == null) {

            ks = KeyStore.getInstance(KS_TYPE, provider);

            ks.load(null, tokenPwd);

        }

        KeyFactory kf = KeyFactory.getInstance("RSA", provider);

        PrivateKey pkSession = (PrivateKey)kf.translateKey(pk3);

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

        ks.setKeyEntry("pkSession", pkSession, null, chain3);

        KeyStore.PrivateKeyEntry pke =

                (KeyStore.PrivateKeyEntry)ks.getEntry("pkSession", null);

        System.out.println("pkSession = " + pke.getPrivateKey());

        Certificate[] chain = pke.getCertificateChain();

        if (chain.length != chain3.length) {

            throw new SecurityException("received chain not correct length");

        }

        for (int i = 0; i < chain.length; i++) {

            if (!chain[i].equals(chain3[i])) {

                throw new SecurityException("received chain not equal");

            }

        }

        System.out.println("test " + testnum++ + " passed");

        return testnum;

    }

    private static void list() throws Exception {

        int testnum = 1;

        ks = KeyStore.getInstance(KS_TYPE, provider);

        // check instance

        if (ks.getProvider() instanceof java.security.AuthProvider) {

            System.out.println("keystore provider instance of AuthProvider");

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException("did not get AuthProvider KeyStore");

        }

        // load

        ks.load(null, tokenPwd);

        System.out.println("test " + testnum++ + " passed");

        // aliases

        Enumeration enu = ks.aliases();

        int count = 0;

        while (enu.hasMoreElements()) {

            count++;

            System.out.println("alias " +

                                count +

                                " = " +

                                (String)enu.nextElement());

        }

    }

    private static void module() throws Exception {

        // perform Security.addProvider of P11 provider

        Security.addProvider(getSunPKCS11(System.getProperty("CUSTOM_P11_CONFIG")));

        String KS_PROVIDER = "SunPKCS11-" + System.getProperty("TOKEN");

        KeyStoreLoginModule m = new KeyStoreLoginModule();

        Subject s = new Subject();

        Map options = new HashMap();

        options.put("keyStoreURL", "NONE");

        options.put("keyStoreType", KS_TYPE);

        options.put("keyStoreProvider", KS_PROVIDER);

        options.put("debug", "true");

        m.initialize(s, new TextCallbackHandler(), new HashMap(), options);

        m.login();

        m.commit();

        System.out.println("authenticated subject = " + s);

        m.logout();

        System.out.println("authenticated subject = " + s);

    }

    /**

     * SCA1000 does not handle extended secret key tests

     * . Blowfish (CKR_TEMPLATE_INCOMPLETE)

     * . AES (CKR_TEMPLATE_INCOMPLETE)

     * . RC4 (CKR_ATTRIBUTE_TYPE_INVALID)

     * so do this instead

     */

    private static int skey(int testnum) throws Exception {

        if (ks == null) {

            ks = KeyStore.getInstance(KS_TYPE, provider);

            ks.load(null, tokenPwd);

        }

        // delete all old aliases

        Enumeration enu = ks.aliases();

        int count = 0;

        while (enu.hasMoreElements()) {

            String next = (String)enu.nextElement();

            ks.deleteEntry(next);

            System.out.println("deleted entry for: " + next);

        }

        // set good ske 1

        ks.setKeyEntry("sk1", sk1, null, null);

        System.out.println("test " + testnum++ + " passed");

        // set good ske 2

        ks.setKeyEntry("sk2", sk2, null, null);

        System.out.println("test " + testnum++ + " passed");

        // getEntry good ske 1

        KeyStore.SecretKeyEntry ske =

                (KeyStore.SecretKeyEntry)ks.getEntry("sk1", null);

        if ("DES".equals(ske.getSecretKey().getAlgorithm())) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException

                ("expected DES, got " + ske.getSecretKey().getAlgorithm());

        }

        // getEntry good ske 2

        ske = (KeyStore.SecretKeyEntry)ks.getEntry("sk2", null);

        if ("DESede".equals(ske.getSecretKey().getAlgorithm())) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException

                ("expected DESede, got " + ske.getSecretKey().getAlgorithm());

        }

        // getKey good ske 1

        SecretKey skey = (SecretKey)ks.getKey("sk1", null);

        if ("DES".equals(skey.getAlgorithm())) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException

                ("expected DES, got " + skey.getAlgorithm());

        }

        // getKey good ske 2

        skey = (SecretKey)ks.getKey("sk2", null);

        if ("DESede".equals(skey.getAlgorithm())) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException

                ("expected DESede, got " + skey.getAlgorithm());

        }

        // aliases

        enu = ks.aliases();

        count = 0;

        while (enu.hasMoreElements()) {

            count++;

            System.out.println("alias " +

                                count +

                                " = " +

                                (String)enu.nextElement());

        }

        if (count == 2) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException("expected 2 aliases");

        }

        // size

        if (ks.size() == 2) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException("expected size 2");

        }

        // isCertificateEntry sk1

        if (!ks.isCertificateEntry("sk1")) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException("expected ske");

        }

        // isKeyEntry sk1

        if (ks.isKeyEntry("sk1")) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException("expected ske");

        }

        // entryInstanceOf sk2

        if (ks.entryInstanceOf("sk2", KeyStore.SecretKeyEntry.class)) {

            System.out.println("test " + testnum++ + " passed");

        } else {

            throw new SecurityException("expected ske");

        }

        return testnum;

    }

    private static int setAttribute(int testnum) throws Exception {

        if (ks == null) {

            ks = KeyStore.getInstance(KS_TYPE, provider);

            ks.load(null, tokenPwd);

        }

        if (!ks.containsAlias("pk1")) {

            // set good pke 1

            ks.setKeyEntry("pk1", pk1, null, chain1);

            System.out.println("test " + testnum++ + " passed");

        }

        // delete all old aliases except pk1

        Enumeration enu = ks.aliases();

        int count = 0;

        while (enu.hasMoreElements()) {

            String next = (String)enu.nextElement();

            if (!"pk1".equals(next)) {

                ks.deleteEntry(next);

                System.out.println("deleted entry for: " + next);

            }

        }

        KeyStore.PrivateKeyEntry pke =

                (KeyStore.PrivateKeyEntry)ks.getEntry("pk1", null);