// Crypto/WzAes.cpp

 /*

 This code implements Brian Gladman's scheme

 specified in password Based File Encryption Utility.

 

 Note: you must include MyAes.cpp to project to initialize AES tables

 */

 

 #include "StdAfx.h"

 

 #include "../Common/StreamObjects.h"

 #include "../Common/StreamUtils.h"

 

 #include "Pbkdf2HmacSha1.h"

 #include "RandGen.h"

 #include "WzAes.h"

 

 // define it if you don't want to use speed-optimized version of Pbkdf2HmacSha1

 // #define _NO_WZAES_OPTIMIZATIONS

 

 namespace NCrypto {

 namespace NWzAes {

 

 const unsigned int kAesKeySizeMax = 32;

 

 static const UInt32 kNumKeyGenIterations = 1000;

 

 STDMETHODIMP CBaseCoder::CryptoSetPassword(const Byte *data, UInt32 size)

 {

   if(size > kPasswordSizeMax)

     return E_INVALIDARG;

   _key.Password.SetCapacity(size);

   memcpy(_key.Password, data, size);

   return S_OK;

 }

 

 #define SetUi32(p, d) { UInt32 x = (d); (p)[0] = (Byte)x; (p)[1] = (Byte)(x >> 8); \

     (p)[2] = (Byte)(x >> 16); (p)[3] = (Byte)(x >> 24); }

 

 void CBaseCoder::EncryptData(Byte *data, UInt32 size)

 {

   unsigned int pos = _blockPos;

   for (; size > 0; size--)

   {

     if (pos == AES_BLOCK_SIZE)

     {

       if (++_counter[0] == 0)

         _counter[1]++;

       UInt32 temp[4];

       Aes_Encode32(&Aes, temp, _counter);

       SetUi32(_buffer,      temp[0]);

       SetUi32(_buffer + 4,  temp[1]);

       SetUi32(_buffer + 8,  temp[2]);

       SetUi32(_buffer + 12, temp[3]);

       pos = 0;

     }

     *data++ ^= _buffer[pos++];

   }

   _blockPos = pos;

 }

 

 #ifndef _NO_WZAES_OPTIMIZATIONS

 

 static void BytesToBeUInt32s(const Byte *src, UInt32 *dest, int destSize)

 {

   for (int i = 0 ; i < destSize; i++)

       dest[i] =

           ((UInt32)(src[i * 4 + 0]) << 24) |

           ((UInt32)(src[i * 4 + 1]) << 16) |

           ((UInt32)(src[i * 4 + 2]) <<  8) |

           ((UInt32)(src[i * 4 + 3]));

 }

 

 #endif

 

 STDMETHODIMP CBaseCoder::Init()

 {

   UInt32 keySize = _key.GetKeySize();

   UInt32 keysTotalSize = 2 * keySize + kPwdVerifCodeSize;

   Byte buf[2 * kAesKeySizeMax + kPwdVerifCodeSize];

   

   // for (int ii = 0; ii < 1000; ii++)

   {

     #ifdef _NO_WZAES_OPTIMIZATIONS

 

     NSha1::Pbkdf2Hmac(

       _key.Password, _key.Password.GetCapacity(),

       _key.Salt, _key.GetSaltSize(),

       kNumKeyGenIterations,

       buf, keysTotalSize);

 

     #else

 

     UInt32 buf32[(2 * kAesKeySizeMax + kPwdVerifCodeSize + 3) / 4];

     UInt32 key32SizeTotal = (keysTotalSize + 3) / 4;

     UInt32 salt[kSaltSizeMax * 4];

     UInt32 saltSizeInWords = _key.GetSaltSize() / 4;

     BytesToBeUInt32s(_key.Salt, salt, saltSizeInWords);

     NSha1::Pbkdf2Hmac32(

       _key.Password, _key.Password.GetCapacity(),

       salt, saltSizeInWords,

       kNumKeyGenIterations,

       buf32, key32SizeTotal);

     for (UInt32 j = 0; j < keysTotalSize; j++)

       buf[j] = (Byte)(buf32[j / 4] >> (24 - 8 * (j & 3)));

     

     #endif

   }

 

   _hmac.SetKey(buf + keySize, keySize);

   memcpy(_key.PwdVerifComputed, buf + 2 * keySize, kPwdVerifCodeSize);

   

   _blockPos = AES_BLOCK_SIZE;

   for (int i = 0; i < 4; i++)

     _counter[i] = 0;

 

   Aes_SetKeyEncode(&Aes, buf, keySize);

   return S_OK;

 }

 

 /*

 STDMETHODIMP CEncoder::WriteCoderProperties(ISequentialOutStream *outStream)

 {

   Byte keySizeMode = 3;

   return outStream->Write(&keySizeMode, 1, NULL);

 }

 */

 

 HRESULT CEncoder::WriteHeader(ISequentialOutStream *outStream)

 {

   UInt32 saltSize = _key.GetSaltSize();

   g_RandomGenerator.Generate(_key.Salt, saltSize);

   Init();

   RINOK(WriteStream(outStream, _key.Salt, saltSize));

   return WriteStream(outStream, _key.PwdVerifComputed, kPwdVerifCodeSize);

 }

 

 HRESULT CEncoder::WriteFooter(ISequentialOutStream *outStream)

 {

   Byte mac[kMacSize];

   _hmac.Final(mac, kMacSize);

   return WriteStream(outStream, mac, kMacSize);

 }

 

 STDMETHODIMP CDecoder::SetDecoderProperties2(const Byte *data, UInt32 size)

 {

   if (size != 1)

     return E_INVALIDARG;

   _key.Init();

   Byte keySizeMode = data[0];

   if (keySizeMode < 1 || keySizeMode > 3)

     return E_INVALIDARG;

   _key.KeySizeMode = keySizeMode;

   return S_OK;

 }

 

 HRESULT CDecoder::ReadHeader(ISequentialInStream *inStream)

 {

   UInt32 saltSize = _key.GetSaltSize();

   UInt32 extraSize = saltSize + kPwdVerifCodeSize;

   Byte temp[kSaltSizeMax + kPwdVerifCodeSize];

   RINOK(ReadStream_FAIL(inStream, temp, extraSize));

   UInt32 i;

   for (i = 0; i < saltSize; i++)

     _key.Salt[i] = temp[i];

   for (i = 0; i < kPwdVerifCodeSize; i++)

     _pwdVerifFromArchive[i] = temp[saltSize + i];

   return S_OK;

 }

 

 static bool CompareArrays(const Byte *p1, const Byte *p2, UInt32 size)

 {

   for (UInt32 i = 0; i < size; i++)

     if (p1[i] != p2[i])

       return false;

   return true;

 }

 

 bool CDecoder::CheckPasswordVerifyCode()

 {

   return CompareArrays(_key.PwdVerifComputed, _pwdVerifFromArchive, kPwdVerifCodeSize);

 }

 

 HRESULT CDecoder::CheckMac(ISequentialInStream *inStream, bool &isOK)

 {

   isOK = false;

   Byte mac1[kMacSize];

   RINOK(ReadStream_FAIL(inStream, mac1, kMacSize));

   Byte mac2[kMacSize];

   _hmac.Final(mac2, kMacSize);

   isOK = CompareArrays(mac1, mac2, kMacSize);

   return S_OK;

 }

 

 STDMETHODIMP_(UInt32) CEncoder::Filter(Byte *data, UInt32 size)

 {

   EncryptData(data, size);

   _hmac.Update(data, size);

   return size;

 }

 

 STDMETHODIMP_(UInt32) CDecoder::Filter(Byte *data, UInt32 size)

 {

   _hmac.Update(data, size);

   EncryptData(data, size);

   return size;

 }

 

 }}