package voip;

import iaik.x509.X509Certificate;
import iaik.security.provider.IAIK;
import iaik.utils.CryptoUtils;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.Cipher;
import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

import java.security.Key;
import java.security.PublicKey;
import java.security.PrivateKey;
import java.security.Security;
import javax.crypto.spec.IvParameterSpec;
import java.security.cert.CertificateException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.InvalidAlgorithmParameterException;
import java.security.NoSuchProviderException;

public class CryptoManager
{

    public static   final int        ENC_MODE = Cipher.ENCRYPT_MODE;
    public static   final int        DEC_MODE = Cipher.DECRYPT_MODE;


    String              myDN        = null;
    String              strMyCert   = null;   // PEM encoded

    String              peerDN      = null;
    X509Certificate     peerCert    = null;

    Keystore            keystore    = null;
    SecretKey           aesKey      = null;
    KeyGenerator        aesKeyGen   = null;

    static Logger       logger      = null;

    CryptoManager()
    {
        Security.insertProviderAt(new IAIK(), 2);
        keystore = new Keystore();
        myDN = keystore.getDN();
        strMyCert = iaik.utils.Util.toPemString(keystore.getCertificate(myDN));
        try
        {
            aesKeyGen = KeyGenerator.getInstance(Constants.STR_AES);
            aesKeyGen.init(128);
        }
        catch(NoSuchAlgorithmException ex)
        {
            System.out.println("Failed to generate the shared key generator ");
        }
    }

    String generate_KE_REQ()
    {
        return Constants.STR_KEYTAG + strMyCert;
    }

    boolean verify_KE_REQ (String msg)
    {
        boolean     isValid = false;
        int         iCert = 0;
        int         iCertEnd = 0;

        if ( isValid = ((iCert = msg.indexOf(Constants.STR_CERT_BEGIN)) != -1) &&
                       ((iCertEnd = msg.indexOf(Constants.STR_CERT_END, iCert)) != -1))
        {
            String cert = msg.substring(
                            iCert,
                            iCertEnd + Constants.STR_CERT_END.length()
                            );

            try
            {
                peerCert = new X509Certificate(cert.getBytes());
                peerDN = getDN(peerCert);
            }
            catch (CertificateException e)
            {
                isValid = false;
                e.printStackTrace();
            }
         }

        return isValid;
    }

    String generateKE_REPLY ()
    {
        Assert.assertTrue(peerCert != null);

        String cipher = rsaWrap(generateAESKey(), peerCert.getPublicKey());

        String msg = Constants.STR_KEYTAG + cipher + strMyCert;

        return msg;
    }


    boolean verifyKE_REPLY (String msg)
    {
        boolean isValid = false;
        int     iCert = 0;
        int     iCertEnd = 0;
        int     iKey  = 0;

        if ( isValid = (((iKey = msg.indexOf(Constants.STR_KEYTAG)) != -1) &&
                        ((iCert = msg.indexOf(Constants.STR_CERT_BEGIN, iCert)) != -1) &&
                        ((iCertEnd = msg.indexOf(Constants.STR_CERT_END, iCert)) != -1))
           )
        {
            try
            {
                String cert = msg.substring(
                                iCert,
                                iCertEnd + Constants.STR_CERT_END.length()
                                );

                peerCert = new X509Certificate(cert.getBytes());

                peerDN = getDN(peerCert);

                String ciphered = null;
                ciphered = msg.substring(iKey + Constants.STR_KEYTAG.length(), iCert);

                aesKey = (SecretKey)rsaUnwrap(
                                        ciphered,
                                        keystore.getPrivateKey(myDN)
                                        );
            }
            catch(CertificateException ex)
            {
                isValid = false;
                logger.error (ex.getMessage());
            }
        }

        return isValid;
    }

    public void resetPeer()
    {
        peerDN = null;
        peerCert = null;

        return;
    }

    private Key generateAESKey()
    {
        Assert.assertTrue(aesKeyGen != null);

        aesKey = aesKeyGen.generateKey();
        return aesKey;
    }

    private String getDN (X509Certificate cert)
    {
        return cert.getSubjectDN().getName().toString();
    }

    public static String rsaWrap(Key aesKey, PublicKey pubKey)
    {
        Assert.assertTrue(aesKey != null);
        Assert.assertTrue(pubKey != null);

        String strWrapped = null;

        try
        {
            Cipher cipher = Cipher.getInstance(Constants.RSA_CRPT, Constants.STR_IAIK);
            cipher.init(Cipher.WRAP_MODE, pubKey);

            byte[] wrappedKey = cipher.wrap(aesKey);
            strWrapped = iaik.utils.Util.toString(wrappedKey);
        }
        catch (NoSuchAlgorithmException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (NoSuchProviderException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (NoSuchPaddingException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (IllegalBlockSizeException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (InvalidKeyException ex)
        {
            logger.error(ex.getMessage());
        }
        return strWrapped;
    }

    public static Key rsaUnwrap(String strWrappedKey, PrivateKey prvKey)
    {
        Assert.assertTrue(strWrappedKey != null);
        Assert.assertTrue(prvKey != null);

        Key unwrappedKey = null;

        try
        {
            Cipher cipher = Cipher.getInstance(Constants.RSA_CRPT, Constants.STR_IAIK);
            cipher.init(Cipher.UNWRAP_MODE, prvKey);

            byte[] wrappedKey = iaik.utils.Util.toByteArray(strWrappedKey);
            unwrappedKey = cipher.unwrap(wrappedKey, Constants.AES_CRPT, Cipher.SECRET_KEY);
        }
        catch (NoSuchAlgorithmException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (NoSuchProviderException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (NoSuchPaddingException ex)
        {
            logger.error(ex.getMessage());
        }
        catch (InvalidKeyException ex)
        {
            logger.error(ex.getMessage());
        }
        return unwrappedKey;
    }

    public byte[] aesCipher(int mode, byte[] input)
    {
        byte[] result = null;

        try
        {
            Cipher cipher = Cipher.getInstance(Constants.AES_CRPT);

            byte[] iv = iaik.utils.Util.toByteArray(
                            "8f c3 a5 36 56 b1 f7 78 c1 29 df 4e 98 48 a4 00"
                            );
                IvParameterSpec ivSpec = new IvParameterSpec(iv);
                cipher.init(mode, aesKey, ivSpec);
            result = cipher.doFinal(input);
        }
        catch (InvalidKeyException e)
        {
            e.printStackTrace();
        }
        catch (NoSuchPaddingException nspe)
        {
            nspe.printStackTrace();
        }
        catch (NoSuchAlgorithmException nsae)
        {
            nsae.printStackTrace();
        }
        catch (InvalidAlgorithmParameterException ex)
        {
            ex.printStackTrace();
        }
        catch (IllegalBlockSizeException e)
        {
            e.printStackTrace();
        }
        catch (BadPaddingException e)
        {
            e.printStackTrace();
        }

        return result;
    }
}