/*
*   Byte-oriented AES-256 implementation.
*   All lookup tables replaced with 'on the fly' calculations.
*
*   Copyright (c) 2007 Ilya O. Levin, http://www.literatecode.com
*
*   Permission to use, copy, modify, and distribute this software for any
*   purpose with or without fee is hereby granted, provided that the above
*   copyright notice and this permission notice appear in all copies.
*
*   Modified: Georgiy Pruss 2015
*/
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <time.h>
#include "aes256.h"

#define __ {
#define _  }

typedef unsigned int U32;
typedef unsigned long long U64;

#define DUMP(s,b,n) {int i;printf(s);for(i=0;i<(n);++i)printf("%02x",b[i]);printf("\n");}

#define NB(i) b[i]=(BYTE)(n>>(i*8))
void u2b( U32 n, BYTE* b ) { NB(0); NB(1); NB(2); NB(3); }
void l2b( U64 n, BYTE* b ) { NB(0); NB(1); NB(2); NB(3); NB(4); NB(5); NB(6); NB(7); }

U64 aes_process_file( int enc, BYTE key[KSZ], FILE* f_i, FILE* f_o ) __
  U64 ctr, sz=enc?0:8; int i,n,m;
  BYTE aesbuf[BSZ], buf[BSZ], header[2*sizeof(U32)];
  aes256_context ctx; aes256_init(&ctx, key);
  if(enc) __ /* encrypt - make and write header */
    time_t t; U32 r; srand(time(&t));
    r = ((rand()/19%1621)*1621 + rand()/17%1621)*1621 + rand()/13%1621;
    u2b(t,header); u2b(r,&header[4]);
    m = fwrite( header, 1, sizeof(header), f_o ); if(m!=sizeof(header)) return -1; _
  else __ /* decrypt - read header */
    m = fread( header, 1, sizeof(header), f_i ); if(m!=sizeof(header)) return -2; _
  for(ctr=1;;++ctr) __
    n = fread( buf, 1, BSZ, f_i ); sz += n;
    if( n==0 ) break;
    memcpy(aesbuf,header,sizeof(header));
    l2b(ctr,aesbuf+sizeof(header));
    aes256_encrypt(&ctx, aesbuf);
    for(i=0;i<BSZ;++i) aesbuf[i]^=buf[i];
    m = fwrite( aesbuf, 1, n, f_o ); if(m!=n) return -3; _
  aes256_done(&ctx);
  return sz; _

void s2v( const char* s, BYTE* v, int n ) __
  int i,k; for(i=0;i<n;++i) { sscanf(s+2*i,"%2x",&k); v[i]=(BYTE)k; } _

int main( int ac, char* av[] ) __
  aes256_context ctx; FILE* fi; FILE* fo;
  BYTE key[KSZ], buf[BSZ], pt[BSZ], ct[BSZ]; int i;
  /* Official test samples: http://csrc.nist.gov/groups/STM/cavp/documents/aes/KAT_AES.zip */
  fi = fopen("vectors.txt","rt");
  if(fi) __ char s[200];
    for(i=0;;++i) __
      #define read(s) if(!fgets(s,sizeof(s),fi))break;if(s[strlen(s)-1]=='\x0A')s[strlen(s)-1]='\0'
      read(s); s2v( s, key, KSZ ); read(s); s2v( s, pt, BSZ ); read(s); s2v( s, ct, BSZ );
      memcpy(buf,pt,BSZ); aes256_init(&ctx, key); aes256_encrypt(&ctx, buf);
      if( memcmp(buf,ct,BSZ)!=0 ) printf("*** Error encoding!\n");
      memcpy(buf,ct,BSZ); aes256_init(&ctx, key); aes256_decrypt(&ctx, buf);
      if( memcmp(buf,pt,BSZ)!=0 ) printf("*** Error decoding!\n"); _
      aes256_done(&ctx);
    fclose(fi);
    printf("%d",i); _
  /* Time memory encrypt/decrypt */
  if(ac>2) __ time_t t0,t1;
    memset(key,123,sizeof(key));
    aes256_init(&ctx, key);
    memset(buf,234,BSZ);
    time(&t0);
    for( i=0; i<30000000; ++i ) aes256_encrypt(&ctx, buf);
    for( i=0; i<30000000; ++i ) aes256_decrypt(&ctx, buf);
    time(&t1);
    if( memcmp(buf,pt,BSZ)!=0 ) printf("Error decoding!\n");
    printf( "30M blks e/d: %d s, %.1f MB/s\n", t1-t0, 30e6*2*16/(1e6*(t1-t0)) );
    /* cygwin-64 18.8 file: 15.4, msvc2010-64 9.9 file: 11.6 MB/s */
    aes256_done(&ctx); _
  /* Time file encrypt/decrypt */
  if(ac>1) __ U64 sz; char* K="abcdefghijklmnopqrstuvwxyz123456"; time_t t0,t1;
    time(&t0);
    fi = fopen(av[1],"rb"); if(!fi) return 1;
    fo = fopen("enc","wb");
    memcpy(key,K,KSZ);
    sz = aes_process_file(1,key,fi,fo);
    fclose(fi); fclose(fo);
    time(&t1);
    if(t1>t0) printf( "Enc: %llu, %d s, %.1f MB/s\n", sz, t1-t0, (double)sz/(1e6*(t1-t0)) );
    t0 = t1;
    fi = fopen("enc","rb");
    fo = fopen("dec","wb");
    memcpy(key,K,KSZ);
    sz = aes_process_file(0,key,fi,fo);
    fclose(fi); fclose(fo);
    time(&t1);
    if(t1>t0) printf( "Dec: %llu, %d s, %.1f MB/s\n", sz, t1-t0, (double)sz/(1e6*(t1-t0)) ); _
  return 0; _
