package org.apache.poi.xwpf;

import static org.junit.Assert.assertEquals;

import java.nio.charset.Charset;
import java.security.MessageDigest;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;
import org.apache.poi.poifs.crypt.CryptoFunctions;
import org.apache.poi.poifs.crypt.HashAlgorithm;
import org.apache.poi.util.LittleEndian;
import org.junit.Test;

public class TestDocumentProtection2 {
    static final int InitialCodeArray[] = { 
        0xE1F0, 0x1D0F, 0xCC9C, 0x84C0, 0x110C, 0x0E10, 0xF1CE, 
        0x313E, 0x1872, 0xE139, 0xD40F, 0x84F9, 0x280C, 0xA96A, 
        0x4EC3
    };
    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}
    };
    
    
    @SuppressWarnings("unused")
    @Test
    public void testHWPFEnc() throws Exception {
/**
 <w:documentProtection
  w:edit="trackedChanges"
  w:enforcement="1"
  w:cryptProviderType="rsaFull"
  w:cryptAlgorithmClass="hash"
  w:cryptAlgorithmType="typeAny"
  w:cryptAlgorithmSid="4"
  w:cryptSpinCount="100000"
  w:hash="MUHbcmpC9AnlLsd9v3lW0j30y6E="
  w:salt="2Z+i7o/0EZyUNakVeWzU/w=="/>        
 */
        String strPassword = "Example";
        
        // Generate the Salt
        byte arrSalt[] = Base64.decodeBase64("2Z+i7o/0EZyUNakVeWzU/w==");

        // Iterations specifies the number of times the hashing function shall be iteratively run (using each
        // iteration's result as the input for the next iteration).
        int iterations = 100000;
        
        //Array to hold Key Values
        byte[] generatedKey = new byte[4];

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

            // 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[strPassword.length()];
            
            for (int intLoop = 0; intLoop < strPassword.length(); intLoop++) {
                int intTemp = strPassword.charAt(intLoop);
                arrByteChars[intLoop] = (byte)(intTemp & 0x00FF);
                if (arrByteChars[intLoop] == 0)
                    arrByteChars[intLoop] = (byte)((intTemp & 0xFF00) >> 8);
            }

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

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

            // --> For each character in the strPassword:
            //      --> 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 key’s high-order word with the corresponding word from 
            //          the Encryption Matrix
            for (int intLoop = 0; intLoop < arrByteChars.length; intLoop++) {
                int tmp = intMaxPasswordLength - arrByteChars.length + intLoop;
                for (int intBit = 0; intBit < 7; intBit++) {
                    if ((arrByteChars[intLoop] & (0x0001 << intBit)) != 0) {
                        intHighOrderWord ^= EncryptionMatrix[tmp][intBit];
                    }
                }
            }
            
            // Compute the low-order word of the new key:
            
            // Initialize with 0
            int intLowOrderWord = 0;

            // For each character in the strPassword, going backwards
            for (int intLoopChar = arrByteChars.length - 1; intLoopChar >= 0; intLoopChar--) {
                // low-order word = (((low-order word SHR 14) AND 0x0001) OR (low-order word SHL 1) AND 0x7FFF)) XOR character
                intLowOrderWord = (((intLowOrderWord >> 14) & 0x0001) | ((intLowOrderWord << 1) & 0x7FFF)) ^ arrByteChars[intLoopChar];
            }

            // Lastly,low-order word = (((low-order word SHR 14) AND 0x0001) OR (low-order word SHL 1) AND 0x7FFF)) XOR strPassword length XOR 0xCE4B.
            intLowOrderWord = (((intLowOrderWord >> 14) & 0x0001) | ((intLowOrderWord << 1) & 0x7FFF)) ^ arrByteChars.length ^ 0xCE4B;

            // The byte order of the result shall be reversed [Example: 0x64CEED7E becomes 7EEDCE64. end example],
            // and that value shall be hashed as defined by the attribute values.
            
            LittleEndian.putShort(generatedKey, 0, (short)intLowOrderWord);
            LittleEndian.putShort(generatedKey, 2, (short)intHighOrderWord);
        }
        
        // Implementation Notes List:
        // --> In this third stage, the reversed byte order legacy hash from the second stage shall be converted to Unicode hex 
        // --> string representation
        // Example: If the single byte string 7EEDCE64 is converted to Unicode hex string it will be represented in memory as
        // the following byte stream: 37 00 45 00 45 00 44 00 43 00 45 00 36 00 34 00.
        String sb = Hex.encodeHexString(generatedKey).toUpperCase();
        generatedKey = sb.getBytes(Charset.forName("UTF-16LE"));

        HashAlgorithm hashAlgo = HashAlgorithm.sha1;
        
        // Implementation Notes List:
        // Word requires that the initial hash of the password with the salt not be considered in the count.
        // The initial hash of salt + key is not included in the iteration count.
        // Before calculating the initial hash, you are supposed to prepend (not append) the salt to the key
        MessageDigest sha1 = CryptoFunctions.getMessageDigest(hashAlgo);
        sha1.update(arrSalt);
        byte hash[] = sha1.digest(generatedKey);

        byte[] iterator = new byte[4];
        for (int intTmp = 0; intTmp < iterations; intTmp++) {
            //When iterating on the hash, you are supposed to append the current iteration number.
            LittleEndian.putInt(iterator, 0, intTmp);
            sha1.reset();
            sha1.update(hash);
            sha1.update(iterator);
            sha1.digest(hash, 0, hash.length);
        }

        assertEquals("MUHbcmpC9AnlLsd9v3lW0j30y6E=",Base64.encodeBase64String(hash));
    }
}