//------------------------------------------------------------------------------
// File:          aes_reference.c
// Function:      Firmware encryption using AES reference implementation
// Supported chip(s):
//    - AT91SAM7XC128
//    - AT91SAM7XC256
// Supported toolchain(s):
//    - IAR Embedded Workbench
// Date created:  06 June 2006
// Created by:    JJo
//------------------------------------------------------------------------------
// Notes:
//
// This code is based on the AES reference implementation published by Paulo
// Barreto and Vincent Rijmen.
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// Includes
//------------------------------------------------------------------------------

#include "aes_reference.h"

#if defined(USE_ENCRYPTION) && defined(ENCRYPTION_AES_REF)

//------------------------------------------------------------------------------
// Global variables
//------------------------------------------------------------------------------

#include "aes_reference.dat"

static unsigned char shifts[3][2][4] = {
   0, 1, 2, 3,
   0, 3, 2, 1,
   
   0, 1, 2, 3,
   0, 5, 4, 3,
   
   0, 7, 5, 5,
   0, 1, 3, 4
};

__no_init static unsigned char key[KC][4];
__no_init static unsigned char expandedKey[ROUNDS+1][BC][4];
__no_init static unsigned int T0[256], T1[256], T2[256], T3[256], TF[256];

#if defined(ENCRYPTION_CBC) || defined(ENCRYPTION_CTR)
__no_init static unsigned char IV[BC][4];
#endif

//------------------------------------------------------------------------------
// Inline functions
//------------------------------------------------------------------------------

/**
 * Name:     mul
 * Purpose:  Multiplies two elements of GF(2^m)
 * Input(s):
 *  - First operand
 *  - Second operand
 * Output:   Result of multiplication
 */
static inline unsigned char mul(unsigned char a, unsigned char b) {

  if (a && b) {
    
    return Alogtable[(Logtable[a] + Logtable[b])%255];
  }
  else {
    
    return 0;
  }
}

/**
 * Name:     min
 * Purpose:  Returns the minimum between two numbers
 * Input(s):
 *  - First number
 *  - Second number
 * Ouput:    Minimum between the two operands
 */
#if defined(ENCRYPTION_CTR)
static unsigned int min(unsigned int number1, unsigned int number2) {
  
  if (number1 > number2) {
    
    return number2;
  }
  else {
    
    return number1;
  }
}
#endif

/**
 * Name:     addRoundKey
 * Purpose:  XOR text and round key together
 * Input(s):
 *  - Plain text
 *  - Round key
 */
static inline void addRoundKey(unsigned char a[BC][4], const unsigned char rk[BC][4]) {
	
  for (unsigned int i=0; i < BC; i++) {
   
    ((int *) a)[i] ^= ((int *) rk)[i];
  }
}

/**
 * Name:     keySchedule
 * Purpose:  Performs the AES key schedule
 * Input(s):
 *  - Key to use
 *  - Buffer to store expanded key schedule
 */
static inline void keySchedule(unsigned char k[KC][4], unsigned char W[ROUNDS+1][BC][4]) {
  
  // Local variables
  int t, rconpointer = 0;
  unsigned char tk[KC][4];   

  for(unsigned int j=0; j < KC; j++) {
    
    ((int *) tk)[j] = ((int *) k)[j];
  }
  
  t = 0;

  /* copy values into round key array */
  for(unsigned int j=0; (j < KC) && (t < (ROUNDS+1)*BC); j++, t++) {

    ((int *) W[t / BC])[t%BC] = ((int *) tk)[j];
  }
  
  while (t < (ROUNDS+1)*BC) { 
    
    tk[0][0] ^= S[tk[KC-1][1]] ^ rcon[rconpointer++];
    tk[0][1] ^= S[tk[KC-1][2]];
    tk[0][2] ^= S[tk[KC-1][3]];
    tk[0][3] ^= S[tk[KC-1][0]];

    if (KC != 8) {
        
      for(unsigned int j=1; j < KC; j++) {
        
        ((int *) tk)[j] ^= ((int *) tk)[j-1];
      }
    }
    else {
      for(unsigned int j=1; j < KC/2; j++) {
          
        ((int *) tk)[j] ^= ((int *) tk)[j-1];
      }
      
      tk[KC/2][0] ^= S[tk[KC/2 - 1][0]];
      tk[KC/2][1] ^= S[tk[KC/2 - 1][1]];
      tk[KC/2][2] ^= S[tk[KC/2 - 1][2]];
      tk[KC/2][3] ^= S[tk[KC/2 - 1][3]];
      
      for(unsigned int j=KC/2+1; j < KC; j++) {
        
        ((int *) tk)[j] ^= ((int *) tk)[j-1];
      }
    }
    
    /* copy values into round key array */
    for(unsigned int j=0; (j < KC) && (t < (ROUNDS+1)*BC); j++, t++) {

      ((int *) W[t/BC])[t%BC] = ((int *) tk)[j];
    }
  }
}

/**
 * Name:     invKeySchedule
 * Purpose:  Performs the AES inverse key schedule
 * Input(s):
 *  - Key to use
 *  - Buffer to store expanded key schedule
 */
#if defined(ENCRYPTION_ECB) || defined(ENCRYPTION_CBC)
static inline void invKeySchedule(unsigned char k[KC][4], unsigned char W[ROUNDS+1][BC][4]) {
  
  // Expand key normally
  keySchedule(k, W);
  
  // Apply invMixColumns to all rounds except first and last one
  for (unsigned int r=1; r < ROUNDS; r++) {
    for (unsigned int j=0; j < BC; j++) {
     
      unsigned char tmp[4];
      
      tmp[0] = mul(0x0E, W[r][j][0]) ^ mul(0x0B, W[r][j][1]) ^
               mul(0x0D, W[r][j][2]) ^ mul(0x09, W[r][j][3]);
      tmp[1] = mul(0x0E, W[r][j][1]) ^ mul(0x0B, W[r][j][2]) ^
               mul(0x0D, W[r][j][3]) ^ mul(0x09, W[r][j][0]);
      tmp[2] = mul(0x0E, W[r][j][2]) ^ mul(0x0B, W[r][j][3]) ^
               mul(0x0D, W[r][j][0]) ^ mul(0x09, W[r][j][1]);
      tmp[3] = mul(0x0E, W[r][j][3]) ^ mul(0x0B, W[r][j][0]) ^
               mul(0x0D, W[r][j][1]) ^ mul(0x09, W[r][j][2]);
      
      W[r][j][0] = tmp[0];
      W[r][j][1] = tmp[1];
      W[r][j][2] = tmp[2];
      W[r][j][3] = tmp[3];
    }
  }
}
#endif

/**
 * Name:     rotBytes
 * Purpose:  Perform the RotBytes operation needed by the AES cipher
 * Input(s):
 *  - Word32 to rotate
 * Output:   Rotated word.
 */
static inline unsigned int rotBytes(unsigned int input) {
 
  return ((input << 8) | (input >> 24));
}

/**
 * Name:     generateEncryptionLUTs
 * Purpose:  Generates the lookup tables needed for encryption
 * Input(s):
 *  - Pointer to t0
 *  - Pointer to t1
 *  - Pointer to t2
 *  - Pointer to t3
 *  - Pointer to tf
 *  - SBox
 */
#if defined(ENCRYPTION_CTR)
static inline void generateEncryptionLUTs(unsigned int * t0,
                                          unsigned int * t1,
                                          unsigned int * t2,
                                          unsigned int * t3,
                                          unsigned int * tf,
                                          unsigned char box[256]) {
 
  for (unsigned int a=0; a <= 255; a++) {
   
    // Calc t0
    t0[a] = (mul(2, box[a])) |
            (box[a] << 8) |
            (box[a] << 16) |
            (mul(3, box[a]) << 24);
    
    // Calc t1, t2, t3
    t1[a] = rotBytes(t0[a]);
    t2[a] = rotBytes(t1[a]);
    t3[a] = rotBytes(t2[a]);
    
    // Calc tf
    tf[a] = box[a] | (box[a] << 8) | (box[a] << 16) | (box[a] << 24);
  }
}
#endif

/**
 * Name:     generateDecryptionLUTs
 * Purpose:  Generates the lookup tables needed for decryption
 * Input(s):
 *  - Pointer to t0
 *  - Pointer to t1
 *  - Pointer to t2
 *  - Pointer to t3
 *  - Pointer to tf
 *  - SBox
 */
#if defined(ENCRYPTION_ECB) || defined(ENCRYPTION_CBC)
static inline void generateDecryptionLUTs(unsigned int * t0,
                                          unsigned int * t1,
                                          unsigned int * t2,
                                          unsigned int * t3,
                                          unsigned int * tf,
                                          unsigned char box[256]) {
 
  for (unsigned int a=0; a <= 255; a++) {
   
    // Calc t0
    t0[a] = (mul(0x0E, box[a])) |
            (mul(0x09, box[a]) << 8) |
            (mul(0x0D, box[a]) << 16) |
            (mul(0x0B, box[a]) << 24);
    
    // Calc t1, t2, t3
    t1[a] = rotBytes(t0[a]);
    t2[a] = rotBytes(t1[a]);
    t3[a] = rotBytes(t2[a]);
    
    // Calc tf
    tf[a] = box[a] | (box[a] << 8) | (box[a] << 16) | (box[a] << 24);
  }
}
#endif

/**
 * Name:     copyBlock
 * Purpose:  Copies a block to a buffer
 * Input(s):
 *  - Block to copy
 *   - Buffer to store copy
 */
#if defined(ENCRYPTION_CTR)
static void copyBlock(const unsigned char input[BC][4], unsigned char output[BC][4]) {
 
  for (unsigned int j=0; j < BC; j++) {
      
    ((int *) output)[j] = ((int *) input)[j];
  }
}
#endif

/**
 * Name:     encrypt
 * Purpose:  Encrypts a block of plain text using precalculated LUTs
 * Input(s):
 *  - Block of plain text to encrypt
 *  - Expanded key
 *  - Pointer to table T0
 *  - Pointer to table T1
 *  - Pointer to table T2
 *  - Pointer to table T3
 *  - Pointer to table TF
 */
#if defined(ENCRYPTION_CTR)
static inline void encrypt(unsigned char a[BC][4],
                           const unsigned char rk[ROUNDS+1][BC][4],
                           unsigned int * t0,
                           unsigned int * t1,
                           unsigned int * t2,
                           unsigned int * t3,
                           unsigned int * tf) {
  
  // Local variables
  unsigned char b[BC][4];
                            
  // First key addition	
  addRoundKey(a, rk[0]);

  // ROUNDS-1 ordinary rounds
  for(unsigned int r=1; r < ROUNDS; r++) {
    for (unsigned int j=0; j < BC; j++) {
     
      ((int *) b)[j] = t0[a[j][0]] ^
                       t1[a[(j+shifts[SC][0][1])%BC][1]] ^
                       t2[a[(j+shifts[SC][0][2])%BC][2]] ^
                       t3[a[(j+shifts[SC][0][3])%BC][3]] ^
                       ((int *) rk[r])[j];
    }
    
    if ((++r) == ROUNDS) {
      
      break;
    }
    
    for (unsigned int j=0; j < BC; j++) {
     
      ((int *) a)[j] = t0[b[j][0]] ^
                       t1[b[(j+shifts[SC][0][1])%BC][1]] ^
                       t2[b[(j+shifts[SC][0][2])%BC][2]] ^
                       t3[b[(j+shifts[SC][0][3])%BC][3]] ^
                       ((int *) rk[r])[j];
    }
  }
  
  // Last round (no MixColumns)
  for (unsigned int j=0; j < BC; j++) {
   
    ((int *) a)[j] = (t0f[b[j][0]]) ^
                     (t1f[b[(j+shifts[SC][0][1])%BC][1]]) ^
                     (t2f[b[(j+shifts[SC][0][2])%BC][2]]) ^
                     (t3f[b[(j+shifts[SC][0][3])%BC][3]]) ^
                     ((int *) rk[ROUNDS])[j];
  }
}
#endif

/**
 * Name:     decrypt
 * Purpose:  Decrypts a block of plain text using precalculated LUTs
 * Input(s):
 *  - Block of cipher text to decrypt
 *  - Expanded key
 *  - Pointer to table T0
 *  - Pointer to table T1
 *  - Pointer to table T2
 *  - Pointer to table T3
 *  - Pointer to table TF
 */
#if defined(ENCRYPTION_ECB) || defined(ENCRYPTION_CBC)
static inline void decrypt(unsigned char a[BC][4],
                           const unsigned char rk[ROUNDS+1][BC][4],
                           unsigned int * t0,
                           unsigned int * t1,
                           unsigned int * t2,
                           unsigned int * t3,
                           unsigned int * tf) {
  
  // Local variables
  unsigned char b[BC][4];
                            
  // First key addition	
  addRoundKey(a, rk[ROUNDS]);

  // ROUNDS-1 ordinary rounds
  for(unsigned int r=ROUNDS-1; r > 0; r--) {
    for (unsigned int j=0; j < BC; j++) {
     
      ((int *) b)[j] = t0[a[j][0]] ^
                       t1[a[(j+shifts[SC][1][1])%BC][1]] ^
                       t2[a[(j+shifts[SC][1][2])%BC][2]] ^
                       t3[a[(j+shifts[SC][1][3])%BC][3]] ^
                       ((int *) rk[r])[j];
    }
    
    if ((--r) == 0) {
      
      break;
    }
    
    for (unsigned int j=0; j < BC; j++) {
     
      ((int *) a)[j] = t0[b[j][0]] ^
                       t1[b[(j+shifts[SC][1][1])%BC][1]] ^
                       t2[b[(j+shifts[SC][1][2])%BC][2]] ^
                       t3[b[(j+shifts[SC][1][3])%BC][3]] ^
                       ((int *) rk[r])[j];
    }
  }
  
  // Last round (no MixColumns)
  for (unsigned int j=0; j < BC; j++) {
   
    ((int *) a)[j] = (t0f[b[j][0]]) ^
                     (t1f[b[(j+shifts[SC][1][1])%BC][1]]) ^
                     (t2f[b[(j+shifts[SC][1][2])%BC][2]]) ^
                     (t3f[b[(j+shifts[SC][1][3])%BC][3]]) ^
                     ((int *) rk[0])[j];
  }
}
#endif

/**
 * Name:     ASCII2Hex
 * Purpose:  Converts an ASCII hexadecimal representation to a raw binary one
 * Input(s):
 *  - ASCII value
 *  - Buffer to store binary value
 *  - Size of value
 */
static void ASCII2Hex(const unsigned char * ascii, unsigned char * binary, unsigned int length) {
  
  // Local variables
  unsigned char * ptr;
  
  ptr = (unsigned char *) binary;
  
  for (unsigned int i=0; i < length; i++, ptr++, ascii++) {

    if (*ascii >= 'A') {

      *ptr = *ascii - 'A' + 10;
    }
    else {

      *ptr = *ascii - '0';
    }

    *ptr <<= 4;
    ascii++;
  
    if (*ascii >= 'A') {

      *ptr += *ascii - 'A' + 10;
    }
    else {

      *ptr += *ascii - '0';
    }
  }
}

/**
 * Name:     ecb_decrypt
 * Purpose:  Decrypts a cipher text using ECB mode
 * Input(s):
 *  - Cipher text to decrypt
 *  - Buffer to store plain text
 *  - Length of cipher text
 *  - Expanded key to use
 * Output:   OK if successful, ERROR otherwise
 */
#if defined(ENCRYPTION_ECB)
static inline unsigned int ecb_decrypt(const unsigned char * cipherText,
                                       unsigned char * plainText,
                                       unsigned int length,
                                       const unsigned char expandedKey[ROUNDS+1][BC][4]) {
 
  // Local variables
  unsigned char block[BC][4];
  
  // Check input parameters
  if ((cipherText == NULL) || (plainText == NULL) || (expandedKey == NULL)) {
    
    debug_printf("AES/REF: NULL parameter(s).\n");
    return ERROR;
  }
  
  if (length%ENCRYPTION_BLOCK_LENGTH != 0) {
    
    debug_printf("AES/REF: Data length must be a multiple of the cipher block size.\n");
    return ERROR;
  }
  
  // ECB decryption
  for (unsigned int l=0; l < length;) {
    
    // Copy cipher text block, decrypt it and copy result
    for (unsigned int i=0; i < ENCRYPTION_BLOCK_LENGTH; i++) {
      
      ((char *) block)[i] = cipherText[l+i];
    }

    decrypt(block, expandedKey, T0, T1, T2, T3, TF);
    
    for (unsigned int i=0; i < ENCRYPTION_BLOCK_LENGTH; i++) {
      
      plainText[l+i] = ((char *) block)[i];
    }
    
    l += ENCRYPTION_BLOCK_LENGTH;
  }
  
  return OK;
}
#endif

/**
 * Name:     cbc_decrypt
 * Purpose:  Decrypts a cipher text using CBC mode
 * Input(s):
 *  - Cipher text to decrypt
 *  - Buffer to store plain text
 *  - Length of cipher text (in bytes)
 *  - Expanded key to use
 *  - Initialization vector to use
 * Ouput:    OK if successful, ERROR otherwise
 */
#if defined(ENCRYPTION_CBC)
static inline unsigned int cbc_decrypt(const unsigned char * cipherText,
                                       unsigned char * plainText,
                                       unsigned int length,
                                       const unsigned char expandedKey[ROUNDS+1][BC][4],
                                       unsigned char IV[BC][4]) {
               
  // Local variables
  unsigned char block[BC][4];
    
  // Check input parameters
  if ((cipherText == NULL) || (plainText == NULL)) {
    
    debug_printf("AES/REF: NULL parameter(s).\n");
    return ERROR;
  }
  
  if (length%ENCRYPTION_BLOCK_LENGTH != 0) {
    
    debug_printf("AES/REF: Cipher text length must be a multiple of the cipher block length.\n");
    return ERROR;
  }
                                         
  // Decrypt data
  for (unsigned int l=0; l < length;) {

    // Copy and decrypt a block of cipher text
    for (unsigned int i=0; i < BC; i++) {
     
      ((int *) block)[i] = ((int *) &cipherText[l])[i];
    }
   
    decrypt(block, expandedKey, T0, T1, T2, T3, TF);
    
    // Xor decrypted text & IV, copy new IV
    for (unsigned int i=0; i < BC; i++) {
     
      unsigned int tmp = ((int *) block)[i] ^ ((int *) IV)[i];
      ((int *) IV)[i] = ((int *) &cipherText[l])[i];
      ((int *) &plainText[l])[i] = tmp;
    }
    
    // Loop progression
    l += ENCRYPTION_BLOCK_LENGTH;
  }
  
  return OK;
}
#endif

/**
 * Name:     ctr_decrypt
 * Purpose:  Decrypts a cipher text using CTR mode
 * Input(s):
 *  - Cipher text to decrypt
 *  - Buffer to store plain text
 *  - Length of cipher text
 *  - Expanded key to use
 *  - Initialization vector to use
 * Output:   OK if successful, ERROR otherwise
 */
#if defined(ENCRYPTION_CTR)
static inline unsigned int ctr_decrypt(const unsigned char * cipherText,
                                       unsigned char * plainText,
                                       unsigned int length,
                                       const unsigned char expandedKey[ROUNDS+1][BC][4],
                                       unsigned char IV[BC][4]) {

  // Local variables
  unsigned char block[BC][4];
  unsigned int bytes;
  
  // Check input parameters
  if ((cipherText == NULL) || (plainText == NULL)) {
    
    return ERROR;
  }

  for (unsigned int l=0; l < length;) {

    // Copy counter and encrypt it
    copyBlock(IV, block);
    encrypt(block, expandedKey, T0, T1, T2, T3, TF);
    
    // XOR current plain text block with encrypted counter
    bytes = min(length - l, ENCRYPTION_BLOCK_LENGTH);
    
    for (unsigned int i=0; i < bytes; i++) {
     
      plainText[l+i] = cipherText[l+i] ^ ((char *) block)[i];
    }

    // Increment counter (big-endian) and number of encrypted bytes
    for (int i=ENCRYPTION_BLOCK_LENGTH-1; i >= 0; i--) {
        
      if (++((char *) IV)[i] != 0) {
        
        break;
      }
    }
    
    l += bytes;
  }
  
  return OK;
}
#endif

//------------------------------------------------------------------------------
// Functions
//------------------------------------------------------------------------------

/**
 * Name:     aes_ref_init
 * Purpose:  Initializes the AES algorithm
 */
void aes_ref_init(void) {

  debug_printf("AES/REF: Initializing ...\n");
  
  ASCII2Hex(ENCRYPTION_KEY, (unsigned char *) key, ENCRYPTION_KEY_LENGTH);
  
#if defined(ENCRYPTION_ECB) || defined(ENCRYPTION_CBC)
  
  // Initialize key schedule
  invKeySchedule(key, expandedKey);
  
  // Generate lookup tables
  generateDecryptionLUTs(T0, T1, T2, T3, TF, Si);
  
#elif defined(ENCRYPTION_CTR)
  
  // Initialize key schedule
  keySchedule(key, expandedKey);
  
  // Generate lookup tables
  generateEncryptionLUTs(T0, T1, T2, T3, TF, S);
#endif
  
#if defined(ENCRYPTION_CBC) || defined(ENCRYPTION_CTR)
  // Initialize counter
  ASCII2Hex(ENCRYPTION_IV, (unsigned char *) IV, ENCRYPTION_BLOCK_LENGTH);
#endif
  
  debug_printf("AES/REF: Initialization done.\n");
}

/**
 * Name:     aes_ref_cleanup
 * Purpose:  Cleanup the AES algorithm
 */
void aes_ref_cleanup(void) {

  debug_printf("AES/REF: Cleaning up ...\n");
  debug_printf("AES/REF: Cleanup done.\n");
}

/**
 * Name:     aes_ref_decrypt
 * Purpose:  Decrypt a cipher text of variable length
 * Input(s):
 *  - Cipher text to decrypt
 *  - Buffer to store plain text
 *  - Length of cipher text (in bytes)
 * Output:   OK if decryption was successful, ERROR otherwise.
 */
int aes_ref_decrypt(const unsigned char * cipherText,
                    unsigned char * plainText,
                    unsigned int length) {

#if defined(ENCRYPTION_ECB)
  return ecb_decrypt(cipherText, plainText, length, expandedKey);
#elif defined(ENCRYPTION_CBC)
  return cbc_decrypt(cipherText, plainText, length, expandedKey, IV);;
#elif defined(ENCRYPTION_CTR)
  return ctr_decrypt(cipherText, plainText, length, expandedKey, IV);
#endif
}

#endif // defined(USE_ENCRYPTION) && defined(ENCRYPTION_AES_REF)
