/* ====================================================================
   Licensed to the Apache Software Foundation (ASF) under one or more
   contributor license agreements.  See the NOTICE file distributed with
   this work for additional information regarding copyright ownership.
   The ASF licenses this file to You under the Apache License, Version 2.0
   (the "License"); you may not use this file except in compliance with
   the License.  You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
==================================================================== */
package org.apache.poi.poifs.crypt;

import java.nio.charset.Charset;
import java.security.DigestException;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.MessageDigest;
import java.security.Provider;
import java.security.Security;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.RC2ParameterSpec;

import org.apache.poi.EncryptedDocumentException;
import org.apache.poi.util.LittleEndian;
import org.apache.poi.util.LittleEndianConsts;
import org.apache.poi.util.StringUtil;

/**
 * Helper functions used for standard and agile encryption
 */
public class CryptoFunctions {
    /**
     * 2.3.4.7 ECMA-376 Document Encryption Key Generation (Standard Encryption)
     * 2.3.4.11 Encryption Key Generation (Agile Encryption)
     * 
     * The encryption key for ECMA-376 document encryption [ECMA-376] using agile encryption MUST be 
     * generated by using the following method, which is derived from PKCS #5: Password-Based
     * Cryptography Version 2.0 [RFC2898].
     * 
     * Let H() be a hashing algorithm as determined by the PasswordKeyEncryptor.hashAlgorithm
     * element, H_n be the hash data of the n-th iteration, and a plus sign (+) represent concatenation. The
     * password MUST be provided as an array of Unicode characters. Limitations on the length of the
     * password and the characters used by the password are implementation-dependent. The initial
     * password hash is generated as follows:
     * 
     * - H_0 = H(salt + password)
     * 
     * The salt used MUST be generated randomly. The salt MUST be stored in the
     * PasswordKeyEncryptor.saltValue element contained within the \EncryptionInfo stream (1) as
     * specified in section 2.3.4.10. The hash is then iterated by using the following approach:
     * 
     * - H_n = H(iterator + H_n-1)
     * 
     * where iterator is an unsigned 32-bit value that is initially set to 0x00000000 and then incremented
     * monotonically on each iteration until PasswordKey.spinCount iterations have been performed.
     * The value of iterator on the last iteration MUST be one less than PasswordKey.spinCount.
     * 
     * For POI, H_final will be calculated by {@link #generateKey(byte[],HashAlgorithm,byte[],int)}
     *
     * @param password
     * @param hashAlgorithm
     * @param salt
     * @param spinCount
     * @return the hashed password
     */
    public static byte[] hashPassword(String password, HashAlgorithm hashAlgorithm, byte salt[], int spinCount) {
        return hashPassword(password, hashAlgorithm, salt, spinCount, true);
    }
        
    /**
     * Generalized method for read and write protection hash generation.
     * The difference is, read protection uses the order iterator then hash in the hash loop, whereas write protection
     * uses first the last hash value and then the current iterator value
     *
     * @param password
     * @param hashAlgorithm
     * @param salt
     * @param spinCount
     * @param iteratorFirst if true, the iterator is hashed before the n-1 hash value,
     *        if false the n-1 hash value is applied first
     * @return the hashed password
     */
    public static byte[] hashPassword(String password, HashAlgorithm hashAlgorithm, byte salt[], int spinCount, boolean iteratorFirst) {
        // If no password was given, use the default
        if (password == null) {
            password = Decryptor.DEFAULT_PASSWORD;
        }
        
        MessageDigest hashAlg = getMessageDigest(hashAlgorithm);
        
        hashAlg.update(salt);
        byte[] hash = hashAlg.digest(StringUtil.getToUnicodeLE(password));
        byte[] iterator = new byte[LittleEndianConsts.INT_SIZE];

        byte[] first = (iteratorFirst ? iterator : hash);
        byte[] second = (iteratorFirst ? hash : iterator);
        
        try {
            for (int i = 0; i < spinCount; i++) {
                LittleEndian.putInt(iterator, 0, i);
                hashAlg.reset();
                hashAlg.update(first);
                hashAlg.update(second);
                hashAlg.digest(hash, 0, hash.length); // don't create hash buffer everytime new
            }
        } catch (DigestException e) {
            throw new EncryptedDocumentException("error in password hashing");
        }
        
        return hash;
    }    

    /**
     * 2.3.4.12 Initialization Vector Generation (Agile Encryption)
     * 
     * Initialization vectors are used in all cases for agile encryption. An initialization vector MUST be
     * generated by using the following method, where H() is a hash function that MUST be the same as
     * specified in section 2.3.4.11 and a plus sign (+) represents concatenation:
     * 1. If a blockKey is provided, let IV be a hash of the KeySalt and the following value:
     *    blockKey: IV = H(KeySalt + blockKey)
     * 2. If a blockKey is not provided, let IV be equal to the following value:
     *    KeySalt:IV = KeySalt.
     * 3. If the number of bytes in the value of IV is less than the the value of the blockSize attribute
     *    corresponding to the cipherAlgorithm attribute, pad the array of bytes by appending 0x36 until
     *    the array is blockSize bytes. If the array of bytes is larger than blockSize bytes, truncate the
     *    array to blockSize bytes. 
     **/
    public static byte[] generateIv(HashAlgorithm hashAlgorithm, byte[] salt, byte[] blockKey, int blockSize) {
        byte iv[] = salt;
        if (blockKey != null) {
            MessageDigest hashAlgo = getMessageDigest(hashAlgorithm);
            hashAlgo.update(salt);
            iv = hashAlgo.digest(blockKey);
        }
        return getBlock36(iv, blockSize);
    }

    /**
     * 2.3.4.11 Encryption Key Generation (Agile Encryption)
     * 
     * ... continued ...
     * 
     * The final hash data that is used for an encryption key is then generated by using the following
     * method:
     * 
     * - H_final = H(H_n + blockKey)
     * 
     * where blockKey represents an array of bytes used to prevent two different blocks from encrypting
     * to the same cipher text.
     * 
     * If the size of the resulting H_final is smaller than that of PasswordKeyEncryptor.keyBits, the key
     * MUST be padded by appending bytes with a value of 0x36. If the hash value is larger in size than
     * PasswordKeyEncryptor.keyBits, the key is obtained by truncating the hash value. 
     *
     * @param passwordHash
     * @param hashAlgorithm
     * @param blockKey
     * @param keySize
     * @return intermediate key
     */
    public static byte[] generateKey(byte[] passwordHash, HashAlgorithm hashAlgorithm, byte[] blockKey, int keySize) {
        MessageDigest hashAlgo = getMessageDigest(hashAlgorithm);
        hashAlgo.update(passwordHash);
        byte[] key = hashAlgo.digest(blockKey);
        return getBlock36(key, keySize);
    }

    public static Cipher getCipher(SecretKey key, CipherAlgorithm cipherAlgorithm, ChainingMode chain, byte[] vec, int cipherMode) {
        return getCipher(key, cipherAlgorithm, chain, vec, cipherMode, null);
    }

    /**
     * Initialize a new cipher object with the given cipher properties
     * If the given algorithm is not implemented in the JCE, it will try to load it from the bouncy castle
     * provider.
     *
     * @param key the secrect key
     * @param cipherAlgorithm the cipher algorithm
     * @param chain the chaining mode
     * @param vec the initialization vector (IV), can be null
     * @param cipherMode Cipher.DECRYPT_MODE or Cipher.ENCRYPT_MODE
     * @param padding
     * @return the requested cipher
     * @throws GeneralSecurityException
     * @throws EncryptedDocumentException if the initialization failed or if an algorithm was specified,
     *   which depends on a missing bouncy castle provider 
     */
    public static Cipher getCipher(Key key, CipherAlgorithm cipherAlgorithm, ChainingMode chain, byte[] vec, int cipherMode, String padding) {
        int keySizeInBytes = key.getEncoded().length;
        if (padding == null) padding = "NoPadding";
        
        try {
            // Ensure the JCE policies files allow for this sized key
            if (Cipher.getMaxAllowedKeyLength(cipherAlgorithm.jceId) < keySizeInBytes*8) {
                throw new EncryptedDocumentException("Export Restrictions in place - please install JCE Unlimited Strength Jurisdiction Policy files");
            }

            Cipher cipher;
            if (cipherAlgorithm == CipherAlgorithm.rc4) {
                cipher = Cipher.getInstance(cipherAlgorithm.jceId);
            } else if (cipherAlgorithm.needsBouncyCastle) {
                registerBouncyCastle();
                cipher = Cipher.getInstance(cipherAlgorithm.jceId + "/" + chain.jceId + "/" + padding, "BC");
            } else {
                cipher = Cipher.getInstance(cipherAlgorithm.jceId + "/" + chain.jceId + "/" + padding);
            }
            
            if (vec == null) {
                cipher.init(cipherMode, key);
            } else {
                AlgorithmParameterSpec aps;
                if (cipherAlgorithm == CipherAlgorithm.rc2) {
                    aps = new RC2ParameterSpec(key.getEncoded().length*8, vec);
                } else {
                    aps = new IvParameterSpec(vec);
                }
                cipher.init(cipherMode, key, aps);
            }
            return cipher;
        } catch (GeneralSecurityException e) {
            throw new EncryptedDocumentException(e);
        }
    }    
    
    /**
     * Returns a new byte array with a truncated to the given size. 
     * If the hash has less then size bytes, it will be filled with 0x36-bytes
     *
     * @param hash the to be truncated/filled hash byte array
     * @param size the size of the returned byte array
     * @return the padded hash
     */
    public static byte[] getBlock36(byte[] hash, int size) {
        return getBlockX(hash, size, (byte)0x36);
    }

    /**
     * Returns a new byte array with a truncated to the given size. 
     * If the hash has less then size bytes, it will be filled with 0-bytes
     *
     * @param hash the to be truncated/filled hash byte array
     * @param size the size of the returned byte array
     * @return the padded hash
     */
    public static byte[] getBlock0(byte[] hash, int size) {
        return getBlockX(hash, size, (byte)0);
    }
    
    private static byte[] getBlockX(byte[] hash, int size, byte fill) {
        if (hash.length == size) return hash;
        
        byte[] result = new byte[size];
        Arrays.fill(result, fill);
        System.arraycopy(hash, 0, result, 0, Math.min(result.length, hash.length));
        return result;
    }
    
    public static MessageDigest getMessageDigest(HashAlgorithm hashAlgorithm) {
        try {
            if (hashAlgorithm.needsBouncyCastle) {
                registerBouncyCastle();
                return MessageDigest.getInstance(hashAlgorithm.jceId, "BC");
            } else {
                return MessageDigest.getInstance(hashAlgorithm.jceId);
            }
        } catch (GeneralSecurityException e) {
            throw new EncryptedDocumentException("hash algo not supported", e);
        }
    }
    
    public static Mac getMac(HashAlgorithm hashAlgorithm) {
        try {
            if (hashAlgorithm.needsBouncyCastle) {
                registerBouncyCastle();
                return Mac.getInstance(hashAlgorithm.jceHmacId, "BC");
            } else {
                return Mac.getInstance(hashAlgorithm.jceHmacId);
            }
        } catch (GeneralSecurityException e) {
            throw new EncryptedDocumentException("hmac algo not supported", e);
        }
    }

    @SuppressWarnings("unchecked")
    public static void registerBouncyCastle() {
        if (Security.getProvider("BC") != null) return;
        try {
            ClassLoader cl = Thread.currentThread().getContextClassLoader();
            String bcProviderName = "org.bouncycastle.jce.provider.BouncyCastleProvider";
            Class<Provider> clazz = (Class<Provider>)cl.loadClass(bcProviderName);
            Security.addProvider(clazz.newInstance());
        } catch (Exception e) {
            throw new EncryptedDocumentException("Only the BouncyCastle provider supports your encryption settings - please add it to the classpath.");
        }
    }

    private static final int InitialCodeArray[] = { 
        0xE1F0, 0x1D0F, 0xCC9C, 0x84C0, 0x110C, 0x0E10, 0xF1CE, 
        0x313E, 0x1872, 0xE139, 0xD40F, 0x84F9, 0x280C, 0xA96A, 
        0x4EC3
    };

    private static final byte PadArray[] = {
        (byte)0xBB, (byte)0xFF, (byte)0xFF, (byte)0xBA, (byte)0xFF,
        (byte)0xFF, (byte)0xB9, (byte)0x80, (byte)0x00, (byte)0xBE,
        (byte)0x0F, (byte)0x00, (byte)0xBF, (byte)0x0F, (byte)0x00
    };
    
    private static final int EncryptionMatrix[][] = {
        /* char 1  */ {0xAEFC, 0x4DD9, 0x9BB2, 0x2745, 0x4E8A, 0x9D14, 0x2A09},
        /* char 2  */ {0x7B61, 0xF6C2, 0xFDA5, 0xEB6B, 0xC6F7, 0x9DCF, 0x2BBF},
        /* char 3  */ {0x4563, 0x8AC6, 0x05AD, 0x0B5A, 0x16B4, 0x2D68, 0x5AD0},
        /* char 4  */ {0x0375, 0x06EA, 0x0DD4, 0x1BA8, 0x3750, 0x6EA0, 0xDD40},
        /* char 5  */ {0xD849, 0xA0B3, 0x5147, 0xA28E, 0x553D, 0xAA7A, 0x44D5},
        /* char 6  */ {0x6F45, 0xDE8A, 0xAD35, 0x4A4B, 0x9496, 0x390D, 0x721A},
        /* char 7  */ {0xEB23, 0xC667, 0x9CEF, 0x29FF, 0x53FE, 0xA7FC, 0x5FD9},
        /* char 8  */ {0x47D3, 0x8FA6, 0x0F6D, 0x1EDA, 0x3DB4, 0x7B68, 0xF6D0},
        /* char 9  */ {0xB861, 0x60E3, 0xC1C6, 0x93AD, 0x377B, 0x6EF6, 0xDDEC},
        /* char 10 */ {0x45A0, 0x8B40, 0x06A1, 0x0D42, 0x1A84, 0x3508, 0x6A10},
        /* char 11 */ {0xAA51, 0x4483, 0x8906, 0x022D, 0x045A, 0x08B4, 0x1168},
        /* char 12 */ {0x76B4, 0xED68, 0xCAF1, 0x85C3, 0x1BA7, 0x374E, 0x6E9C},
        /* char 13 */ {0x3730, 0x6E60, 0xDCC0, 0xA9A1, 0x4363, 0x86C6, 0x1DAD},
        /* char 14 */ {0x3331, 0x6662, 0xCCC4, 0x89A9, 0x0373, 0x06E6, 0x0DCC},
        /* char 15 */ {0x1021, 0x2042, 0x4084, 0x8108, 0x1231, 0x2462, 0x48C4}
    };

    /**
     * This method generates the xor verifier for word documents &lt; 2007 (method 2).
     * Its output will be used as password input for the newer word generations which
     * utilize a real hashing algorithm like sha1.
     * 
     * @param password the password
     * @return the hashed password
     * 
     * @see <a href="http://msdn.microsoft.com/en-us/library/dd905229.aspx">2.3.7.4 Binary Document Password Verifier Derivation Method 2</a>
     * @see <a href="http://blogs.msdn.com/b/vsod/archive/2010/04/05/how-to-set-the-editing-restrictions-in-word-using-open-xml-sdk-2-0.aspx">How to set the editing restrictions in Word using Open XML SDK 2.0</a>
     * @see <a href="http://www.aspose.com/blogs/aspose-blogs/vladimir-averkin/archive/2007/08/20/funny-how-the-new-powerful-cryptography-implemented-in-word-2007-turns-it-into-a-perfect-tool-for-document-password-removal.html">Funny: How the new powerful cryptography implemented in Word 2007 turns it into a perfect tool for document password removal.</a>
     */
    public static int createXorVerifier2(String password) {
        //Array to hold Key Values
        byte[] generatedKey = new byte[4];

        //Maximum length of the password is 15 chars.
        final int maxPasswordLength = 15; 
        
        if (!"".equals(password)) {
            // Truncate the password to 15 characters
            password = password.substring(0, Math.min(password.length(), maxPasswordLength));

            // Construct a new NULL-terminated string consisting of single-byte characters:
            //  -- > Get the single-byte values by iterating through the Unicode characters of the truncated Password.
            //   --> For each character, if the low byte is not equal to 0, take it. Otherwise, take the high byte.
            byte[] arrByteChars = new byte[password.length()];
            
            for (int i = 0; i < password.length(); i++) {
                int intTemp = password.charAt(i);
                byte lowByte = (byte)(intTemp & 0x00FF);
                byte highByte = (byte)((intTemp & 0xFF00) >> 8);
                arrByteChars[i] = (lowByte != 0 ? lowByte : highByte);
            }

            // Compute the high-order word of the new key:

            // --> Initialize from the initial code array (see below), depending on the passwords length. 
            int highOrderWord = InitialCodeArray[arrByteChars.length - 1];

            // --> For each character in the password:
            //      --> For every bit in the character, starting with the least significant and progressing to (but excluding) 
            //          the most significant, if the bit is set, XOR the keys high-order word with the corresponding word from 
            //          the Encryption Matrix
            for (int i = 0; i < arrByteChars.length; i++) {
                int tmp = maxPasswordLength - arrByteChars.length + i;
                for (int intBit = 0; intBit < 7; intBit++) {
                    if ((arrByteChars[i] & (0x0001 << intBit)) != 0) {
                        highOrderWord ^= EncryptionMatrix[tmp][intBit];
                    }
                }
            }
            
            // Compute the low-order word of the new key:
            
            // SET Verifier TO 0x0000
            short verifier = 0;

            // FOR EACH PasswordByte IN PasswordArray IN REVERSE ORDER
            for (int i = arrByteChars.length-1; i >= 0; i--) {
                // SET Verifier TO Intermediate3 BITWISE XOR PasswordByte
                verifier = rotateLeftBase15Bit(verifier);
                verifier ^= arrByteChars[i];
            }

            // as we haven't prepended the password length into the input array
            // we need to do it now separately ...
            verifier = rotateLeftBase15Bit(verifier);
            verifier ^= arrByteChars.length;
            
            // RETURN Verifier BITWISE XOR 0xCE4B
            verifier ^= 0xCE4B; // (0x8000 | ('N' << 8) | 'K')

            // The byte order of the result shall be reversed [password "Example": 0x64CEED7E becomes 7EEDCE64],
            // and that value shall be hashed as defined by the attribute values.
            
            LittleEndian.putShort(generatedKey, 0, verifier);
            LittleEndian.putShort(generatedKey, 2, (short)highOrderWord);
        }
        
        return LittleEndian.getInt(generatedKey);
    }

    /**
     * This method generates the xored-hashed password for word documents &lt; 2007.
     */
    public static String xorHashPassword(String password) {
        int hashedPassword = createXorVerifier2(password);
        return String.format("%1$08X", hashedPassword);
    }
    
    /**
     * Convenience function which returns the reversed xored-hashed password for further 
     * processing in word documents 2007 and newer, which utilize a real hashing algorithm like sha1.
     */
    public static String xorHashPasswordReversed(String password) {
        int hashedPassword = createXorVerifier2(password);
        
        return String.format("%1$02X%2$02X%3$02X%4$02X"
            , ( hashedPassword >>> 0 ) & 0xFF
            , ( hashedPassword >>> 8 ) & 0xFF
            , ( hashedPassword >>> 16 ) & 0xFF
            , ( hashedPassword >>> 24 ) & 0xFF
        );
    }

    /**
     * Create the verifier for xor obfuscation (method 1)
     *
     * @see <a href="http://msdn.microsoft.com/en-us/library/dd926947.aspx">2.3.7.1 Binary Document Password Verifier Derivation Method 1</a>
     * @see <a href="http://msdn.microsoft.com/en-us/library/dd905229.aspx">2.3.7.4 Binary Document Password Verifier Derivation Method 2</a>
     * 
     * @param password the password
     * @return the verifier (actually a short value)
     */
    public static int createXorVerifier1(String password) {
        // the verifier for method 1 is part of the verifier for method 2
        // so we simply chop it from there
        return createXorVerifier2(password) & 0xFFFF;
    }
 
    /**
     * Create the xor key for xor obfuscation, which is used to create the xor array (method 1)
     *
     * @see <a href="http://msdn.microsoft.com/en-us/library/dd924704.aspx">2.3.7.2 Binary Document XOR Array Initialization Method 1</a>
     * @see <a href="http://msdn.microsoft.com/en-us/library/dd905229.aspx">2.3.7.4 Binary Document Password Verifier Derivation Method 2</a>
     * 
     * @param password the password
     * @return the xor key
     */
    public static int createXorKey1(String password) {
        // the xor key for method 1 is part of the verifier for method 2
        // so we simply chop it from there
        return createXorVerifier2(password) >>> 16;
    }

    /**
     * Creates an byte array for xor obfuscation (method 1) 
     *
     * @see <a href="http://msdn.microsoft.com/en-us/library/dd924704.aspx">2.3.7.2 Binary Document XOR Array Initialization Method 1</a>
     * @see <a href="http://docs.libreoffice.org/oox/html/binarycodec_8cxx_source.html">Libre Office implementation</a>
     *
     * @param password the password
     * @return the byte array for xor obfuscation
     */
    public static byte[] createXorArray1(String password) {
        if (password.length() > 15) password = password.substring(0, 15);
        byte passBytes[] = password.getBytes(Charset.forName("ASCII"));
        
        // this code is based on the libre office implementation.
        // The MS-OFFCRYPTO misses some infos about the various rotation sizes 
        byte obfuscationArray[] = new byte[16];
        System.arraycopy(passBytes, 0, obfuscationArray, 0, passBytes.length);
        System.arraycopy(PadArray, 0, obfuscationArray, passBytes.length, PadArray.length-passBytes.length+1);
        
        int xorKey = createXorKey1(password);
        
        // rotation of key values is application dependent
        int nRotateSize = 2; /* Excel = 2; Word = 7 */
        
        byte baseKeyLE[] = { (byte)(xorKey & 0xFF), (byte)((xorKey >>> 8) & 0xFF) };
        for (int i=0; i<obfuscationArray.length; i++) {
            obfuscationArray[i] ^= baseKeyLE[i&1];
            obfuscationArray[i] = rotateLeft(obfuscationArray[i], nRotateSize);
        }
        
        return obfuscationArray;
    }

    private static byte rotateLeft(byte bits, int shift) {
        return (byte)(((bits & 0xff) << shift) | ((bits & 0xff) >>> (8 - shift)));
    }
    
    private static short rotateLeftBase15Bit(short verifier) {
        /*
         * IF (Verifier BITWISE AND 0x4000) is 0x0000
         *    SET Intermediate1 TO 0
         * ELSE
         *    SET Intermediate1 TO 1
         * ENDIF
         */
        short intermediate1 = (short)(((verifier & 0x4000) == 0) ? 0 : 1);
        /*
         *  SET Intermediate2 TO Verifier MULTIPLED BY 2
         *  SET most significant bit of Intermediate2 TO 0
         */
        short intermediate2 = (short)((verifier<<1) & 0x7FFF);
        /*
         *  SET Intermediate3 TO Intermediate1 BITWISE OR Intermediate2
         */
        short intermediate3 = (short)(intermediate1 | intermediate2);
        return intermediate3;
    }
}
