/******************************************************************
 *  The Advanced Encryption Standard (aka AES and Rijndael)       *
 *                                                                *
 * There are three modes of operation:                            *
 *   Default --- The user is prompted for a string to encrypt     *
 *   -DTEST --- Run an encryption and decryption for each of the  *
 *      three key sizes.                                          *
 *   -DITERATE --- Iteratively encrypt a block multiple times for *
 *      benchmarking purposes.                                    *
 *      For SAC SL in this mode set a preprocessor                *
 *      variable named SACSL                                      *
 *                                                                * 
 *****************************************************************/
use Array: all;
use String: {to_string,strlen,tochar};
#ifdef SACSL
use MutcIO: all;
#else
use StdIO: all;
use RuntimeError: {error};
#endif

#define MAX_LEN 1024
#define BLOCK_SIZE 16 /* This is fixed by the standard */

#ifdef SACSL
int[16] id( int in)
{
  if( in == 0) {
    retVal = genarray( [16], 0);
  }
  else {
    retVal = [ 0, 17, 34, 51, 68, 85, 102, 119, 136, 153, 170, 187, 204, 221, 238, 255];
  }
  return( retVal);
}
#endif


int main()
{

#ifdef TEST
  if( tests()) {
    returnVal = 0;
  }
  else {
    error(2, "Test Failed");
    returnVal = 1;
  }
#elif defined( ITERATE)
#ifndef SACSL
  iter = (:int) FibreScanIntArray();
  key = (:int[32]) FibreScanIntArray();
  ip = (:int[16]) FibreScanIntArray();
#else
  iter = 1;
  key = [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31];
  ip = id( 1);
#endif
  key = reshape( [8,4], key);
  ip = transpose( arrayTo4Matrix( ip));
  /*words = calcWords( key);*/
  keys = expandKey( key, 60); 
  for ( i=0 ; i<iter ; i++){
    ip = encrypt( ip, keys);
  }
  FibrePrint( ip);
  returnVal = 0;

#else
  test256Key = reshape( [8,4], [ 0,  1,  2,  3, 
                                 4,  5,  6,  7, 
                                 8,  9,  10, 11,
                                 12, 13, 14, 15, 
                                 16, 17, 18, 19,
                                 20, 21, 22, 23, 
                                 24, 25, 26, 27,
                                 28, 29, 30, 31]);
  words = calcWords( test256Key);
  keys = expandKey( test256Key, words); 

  printf( "Please enter text to be enciphered:\n");
  input = scans( MAX_LEN);
  input_bits = string2Blocks( input);
  enc = encrypt( input_bits, keys);
  printf( "enciphered data:\n");
  print( enc);
  dec = decrypt( enc, keys);
  printf( "deciphering data:\n");
  print( dec);

  printf( "The deciphered text is >%s<\n", blocks2string( dec));

  returnVal = 0;

#endif
  return( returnVal);
}


/*
 * These functions are crudely writen and slow.
 * They are here so that this code can be classified 
 * as SacSL until bit primitive support is added to Sac2C.
 */
inline
int BitXOR( int a, int b) 
{
  res = 0;
  val = 1;
  for( i=0 ; i<32 ; i++) {
    bita = ( a / val) % 2;
    bitb = ( b / val) % 2;
    if( bita + bitb == 1) {
      res += val;
    }
    val = val * 2;
  }
  return( res);
}

inline
int BitAND( int a, int b) 
{
  res = 0;
  val = 1;
  for( i=0 ; i<32 ; i++) {
    bita = ( a / val) % 2;
    bitb = ( b / val) % 2;
    if( bita + bitb == 2) {
      res += val;
    }
    val = val * 2;
  }
  return( res);
}

inline
int BitShiftLeft( int shift, int bits)
{
  for( i=0 ; i<shift ; i++) {
    bits = bits * 2;
  }
  return( bits);
}
 

inline
int calcWords( int[.,.] key)
{
  keyHeight = shape(key)[0];
  if( keyHeight == 4) {
    words = 44;
  }
  else if( keyHeight == 6) {
    words = 52;
  }
  else if( keyHeight == 8) {
    words = 60;
  }
  else {
#ifndef SACSL
    error(1, "Invalid key size. "
             "AES supports only 128 bit, 192 bit and 256 bit keys.");
#endif
    words = 0;
  }
  return( words);
}

/*
 * Reshapes a vector with a length 4N for some integer N 
 * to a 4 by N matrix
 */
int[.,.] arrayTo4Matrix( int[.] array)
{
  return( reshape( [shape(array)[0]/4, 4], 
		   array));
}

int[4,4] encrypt( int[4,4] state, int[.,.,.] keys )
{
  rounds = shape(keys)[0]-2;
  state = keyXor( state, keys[0]);
  for( rnd = 1 ; rnd <= rounds ; rnd++) {
    state = round( state, keys[rnd]);
  }

  return( keyXor( shiftRows( subBytes( state)), keys[rounds+1]));
}

int[4,4] decrypt( int[4,4] state, int[.,.,.] keys)
{
  rounds = shape(keys)[0]-2;
  state = keyXor( state, keys[rounds+1]);
  for( rnd = rounds ; rnd > 0 ; rnd--) {
      state = invRound( state, keys[rnd]);
  }

  return( keyXor( invSubBytes( invShiftRows( state)), keys[0]));
}

int[4,4] keyXor( int[4,4] state, int[.,.] key )
{
  return( xor( state, transp( key)));
}

int[4,4] round( int[4,4] state, int[.,.] key)
{
  return( keyXor( mixColumns( shiftRows( subBytes( state))), 
                  key ));
}

int[4,4] invRound( int[4,4] state, int[.,.] key)
{
  return( invMixColumns( keyXor( invSubBytes( invShiftRows( state)), 
			                   key )));
}

int[.,.] getRoundKey(int[.,.] keySpace, int round)
{
  return( tile( [4,4], [round,0], keySpace));
}

/******************************************************
 * Key expansion                                      *
 ******************************************************/

/*
 * This is annoyingly sequential and I'mnot sure right 
 * now how to break it up. Ideas welcome.
 */
/*specialize int[.,.,.] expandKey( int[4,8] key, int totalWords);*/
inline
int[.,.,.] expandKey( int[.,.] key, int totalWords)
{
  roundConstant = [ 1,   2,   4,   8, 
		                16,  32,  64,  128,
		                27,  54,  108, 216,
		                171, 77,  154, 47,
		                94,  188, 99,  198,
                    151, 53,  106, 212];

  keyHeight = shape(key)[0];
  keySpace = genarray( [totalWords,4], 0);

  keySpace = with {
    ( [0] <= iv < [keyHeight] ) : key[iv];
  }: modarray(keySpace);

  for( iv = keyHeight ; iv<totalWords ; iv++) {
    keySpace[iv] = keySpace[iv-1];
    if( iv % keyHeight == 0) {
      keySpace[iv] = rotate( [-1], keySpace[iv]);
      keySpace[iv] = subBytes( keySpace[iv]);
      keySpace[[iv,0]] = BitXOR( keySpace[iv][0], 
                                 roundConstant[(iv/keyHeight)-1]);
    }
    else if( ( totalWords > 52) && /* For 256 bit keys only */
	           ( iv % 4 == 0)) { 
      keySpace[iv] = subBytes(keySpace[iv]);
    }
    keySpace[iv] = xor( keySpace[iv], keySpace[iv-keyHeight]);
  }
  return( reshape( [(totalWords/4), 4, 4], keySpace));
}



/****************************************************
 * Cipher Operation Boxes.                          *
 * All operations have the same name as stated in   *
 * the official AES specification.                  *
 ****************************************************/

specialize int[4,4] subBytes( int[4,4] state);
specialize int[4] subBytes( int[4] state);
int[*] subBytes( int[*] state)
{
  encSbox = [
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5,
    0x30, 0x1,  0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0,
    0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC,
    0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x4,  0xC7, 0x23, 0xC3, 0x18, 0x96, 0x5,  0x9A,
    0x7,  0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x9,  0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0,
    0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x0,  0xED, 0x20, 0xFC, 0xB1, 0x5B,
    0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85,
    0x45, 0xF9, 0x2,  0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5,
    0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0xC,  0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17,
    0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88,
    0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0xB,  0xDB,
    0xE0, 0x32, 0x3A, 0xA,  0x49, 0x6,  0x24, 0x5C,
    0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9,
    0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x8,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6,
    0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x3,  0xF6, 0xE,
    0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94,
    0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0xD,  0xBF, 0xE6, 0x42, 0x68,
    0x41, 0x99, 0x2D, 0xF,  0xB0, 0x54, 0xBB, 0x16];
  return( { iv -> encSbox[state[iv]]});
}

specialize int[4,4] invSubBytes( int[4,4] state);
specialize int[4] invSubBytes( int[4] state);
int[*] invSubBytes( int[*] state)
{
  decSBox = [
    0x52,  0x9, 0x6A, 0xD5, 0x30, 0x36, 0xA5, 0x38,
    0xBF, 0x40, 0xA3, 0x9E, 0x81, 0xF3, 0xD7, 0xFB,
    0x7C, 0xE3, 0x39, 0x82, 0x9B, 0x2F, 0xFF, 0x87,
    0x34, 0x8E, 0x43, 0x44, 0xC4, 0xDE, 0xE9, 0xCB,
    0x54, 0x7B, 0x94, 0x32, 0xA6, 0xC2, 0x23, 0x3D,
    0xEE, 0x4C, 0x95, 0xB,  0x42, 0xFA, 0xC3, 0x4E,
    0x8,  0x2E, 0xA1, 0x66, 0x28, 0xD9, 0x24, 0xB2,
    0x76, 0x5B, 0xA2, 0x49, 0x6D, 0x8B, 0xD1, 0x25,
    0x72, 0xF8, 0xF6, 0x64, 0x86, 0x68, 0x98, 0x16,
    0xD4, 0xA4, 0x5C, 0xCC, 0x5D, 0x65, 0xB6, 0x92,
    0x6C, 0x70, 0x48, 0x50, 0xFD, 0xED, 0xB9, 0xDA,
    0x5E, 0x15, 0x46, 0x57, 0xA7, 0x8D, 0x9D, 0x84,
    0x90, 0xD8, 0xAB, 0x0,  0x8C, 0xBC, 0xD3, 0xA,
    0xF7, 0xE4, 0x58, 0x5,  0xB8, 0xB3, 0x45, 0x6,
    0xD0, 0x2C, 0x1E, 0x8F, 0xCA, 0x3F, 0xF,  0x2,
    0xC1, 0xAF, 0xBD, 0x3,  0x1,  0x13, 0x8A, 0x6B,
    0x3A, 0x91, 0x11, 0x41, 0x4F, 0x67, 0xDC, 0xEA,
    0x97, 0xF2, 0xCF, 0xCE, 0xF0, 0xB4, 0xE6, 0x73,
    0x96, 0xAC, 0x74, 0x22, 0xE7, 0xAD, 0x35, 0x85,
    0xE2, 0xF9, 0x37, 0xE8, 0x1C, 0x75, 0xDF, 0x6E,
    0x47, 0xF1, 0x1A, 0x71, 0x1D, 0x29, 0xC5, 0x89,
    0x6F, 0xB7, 0x62, 0xE,  0xAA, 0x18, 0xBE, 0x1B,
    0xFC, 0x56, 0x3E, 0x4B, 0xC6, 0xD2, 0x79, 0x20,
    0x9A, 0xDB, 0xC0, 0xFE, 0x78, 0xCD, 0x5A, 0xF4,
    0x1F, 0xDD, 0xA8, 0x33, 0x88, 0x7,  0xC7, 0x31,
    0xB1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xEC, 0x5F,
    0x60, 0x51, 0x7F, 0xA9, 0x19, 0xB5, 0x4A, 0xD,
    0x2D, 0xE5, 0x7A, 0x9F, 0x93, 0xC9, 0x9C, 0xEF,
    0xA0, 0xE0, 0x3B, 0x4D, 0xAE, 0x2A, 0xF5, 0xB0,
    0xC8, 0xEB, 0xBB, 0x3C, 0x83, 0x53, 0x99, 0x61,
    0x17, 0x2B, 0x4,  0x7E, 0xBA, 0x77, 0xD6, 0x26,
    0xE1, 0x69, 0x14, 0x63, 0x55, 0x21, 0xC,  0x7D];

  return( { iv -> decSBox[state[iv]]});
}

inline int[4,4] shiftRows( int[4,4] state)
{
  return( shiftRows_( state, -1));
}

inline int[4,4] invShiftRows( int[4,4] state)
{
  return( shiftRows_( state, 1));
}

int[4,4] shiftRows_( int[4,4] state, int shiftBy)
{
  return( with { 
    ([0] <= iv < [shape(state)[0]]) : rotate( iv*shiftBy, state[iv]);
  }: modarray( state)); 
}

int[.,.] addRoundKey( int[.,.] state, int[.,.] exKey, int round)
{
  rKey = tile( [4,4], [round,0], exKey);
  return ( { iv -> BitXOR( state[iv], rKey[iv])});
}

int[4,4] mixColumns(int[4,4] state)
{
  mixFwdMatrix = reshape( [4,4],
			  [ 2, 3, 1, 1,
			    1, 2, 3, 1,
			    1, 1, 2, 3,
			    3, 1, 1, 2 ]);

  return( mix( state, mixFwdMatrix));
}

int[4,4] invMixColumns(int[4,4] state)
{
  mixBackMatrix = reshape( [4,4],
			   [ 14, 11, 13, 9,
			     9,  14, 11, 13,
			     13, 9 , 14, 11, 
			     11, 13, 9,  14]);
  return( mix( state, mixBackMatrix));
}

inline int[.,.] mix( int[.,.] state, int[.,.] mixMatrix)
{
  return( { [.,iv] -> matrixMul(mixMatrix, state[.,iv])} ); 
}



/**************************************************
 *  Utility functions Used by the above functions * 
 **************************************************/

/* 
 * Elementwise Exclusive OR 
 */
inline int[*] xor( int[*] a, int[*] b)
{
  return( { iv -> BitXOR( a[iv], b[iv])});
}

/*
 * Matrix product of a matrix and a vector
 */
inline int[.] matrixMul( int[.,.] matrix, int[.] vector)
{ 
  return( { [iv] -> polyMul( matrix[[iv]], vector)});
}

/* 
 * Polynomial multiplication of a row and a column from a matrix
 */ 
int polyMul( int[*] row, int[*] column)
{
  return( with {
    ([0] <= iv < [shape(row)[0]]) : 
      reduce( polyMulScalers( column[iv], row[iv]));
  }: fold( xor, 0));
}  

/* 
 * Multiplication of two integers in the galois field. 
 */
int polyMulScalers( int a, int b) /* where b is a nibble */
{
  mask = 1;
  result = 0;

  for( i = 0 ; i<4 ; i++) { 
    if ( BitAND( mask, b) != 0) {
      result = BitXOR( result, BitShiftLeft( i, a));
    }
    mask = BitShiftLeft( 1, mask);
  }

  return( result);
}

/* 
 * Polynomial reduction in Galois field. 
 * Accepts a number up to 16 bits in size and returns a reduced 
 * 8 bit equivenlent in the Galois field.
 */
inline int reduce( int n)
{
  cancelBits = 283;
  while (n > 255) {
    sb = BitShiftLeft( calcBitsToShiftBy( n), cancelBits );
    n = BitXOR( n, sb);
  }
  return( n);
}

/* 
 * Used by reduce to calculate how many bits to shift the cancel bits by. 
 */
inline int calcBitsToShiftBy( int x)
{
  t = 16;
  while(( pow(2,t) > x) && ( t>=0)) {
    t--; 
  }
  return( t-8);
}

/* 
 * Matrix transposition 
 */
inline int[.,.] transp( int[.,.] matrix)
{
  return( {[i,j] -> matrix[[j,i]]});
}

/*
 * Calculates x raised to the power y 
 */
inline int pow( int x, int y)
{
  res = x;
  for( i=0 ; i< y-1 ; i++) {
   res *= x;
  }
  return( res);
}




#ifndef ITERATE
/*
 * The following to functions were adapted from des.sac
 * and provide the means for a simple user interface.
 */

int[4,4] string2Blocks( String::string s)
{
  len = strlen( s);
  printf( "len %d\n", len);
  pad = ( len % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE - len % BLOCK_SIZE);
  printf( "pad %d\n", pad);
  ints = toi( tochar( s) ++ genarray( [pad], ' '));
  res = reshape( [4,4], ints);
  return( res);
}

String::string blocks2string( int[4,4] bits)
{
  bytes = reshape( [BLOCK_SIZE], bits);
  res = to_string( tochar( bytes));
  return( res);
}
#endif

#ifdef TEST
/*********************************************************** 
 * Test functions                                          * 
 ***********************************************************/

bool tests()
{
  test128ip = transp( arrayTo4Matrix( [50, 67, 246, 168, 136, 90, 48, 141, 49, 49, 152, 162, 224, 55, 7, 52]));
  test128Key = reshape( [4,4], [43, 126, 21, 22, 40, 174, 210, 166, 171, 247, 21, 136, 9, 207, 79, 60]);
  res128 = reshape( [4,4], [57, 2, 220, 25, 37, 220, 17, 106, 132, 9, 133, 11, 29, 251, 151, 50]);


  test192ip = transp( arrayTo4Matrix( [ 0, 17, 34, 51, 68, 85, 102, 119, 136, 153, 170, 187, 204, 221, 238, 255]));
  test192Key = reshape( [6,4], [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]);
  res192 = transp( reshape( [4,4], [0xdd, 0xa9, 0x7c, 0xa4, 0x86, 0x4c, 0xdf, 0xe0, 
				       0x6e, 0xaf, 0x70, 0xa0, 0xec, 0xd, 0x71, 0x91]));


  test256ip = transp( arrayTo4Matrix( [ 0, 17, 34, 51, 68, 85, 102, 119, 136, 153, 170, 187, 204, 221, 238, 255]));
  test256Key = reshape( [8,4], [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]);
  res256 = transp( reshape( [4,4], [0x8e, 0xa2, 0xb7, 0xca, 0x51, 0x67, 0x45, 0xbf,
                            0xea, 0xfc, 0x49, 0x90, 0x4b, 0x49, 0x60, 0x89 ]));
  
  return( 
	  test( test192ip, test192Key, res192) &&
	  test( test256ip, test256Key, res256) &&
    test( test128ip, test128Key, res128) );
}


/*specialize bool test( int[4,4] ip, int[4,4] key, int[4,4] expectedOutput );
specialize bool test( int[4,4] ip, int[6,4] key, int[4,4] expectedOutput );
specialize bool test( int[4,4] ip, int[8,4] key, int[4,4] expectedOutput );*/
inline
bool test( int[4,4] ip, int[.,.] key, int[4,4] expectedOutput )
{
  words = calcWords( key);
  keys = expandKey( key, words); 
  
  printf("Now follows input decryption and encryption and expected result\n");
  print( ip);
  enc = encrypt( ip, keys);
  printf("Now decrypting\n");
  print( decrypt( enc, keys));
  result = encrypt( ip, keys);
  print( result);
  print(expectedOutput);
  return( all(result == expectedOutput));
}
#endif