/* Copyright 1999 Enhanced Software Technologies Inc.
 * Copyright 2002 The TapiocaStor Group;
 * All Rights Reserved
 *
 * This module is released under a BSD-style Open Source license.
 * See file LICENSE for licensing information. 
 *
 * Written December 1999 Eric Lee Green
 * Converted to use Rijndael January 2002 Eric Lee Green
 * Converted into Ruby module January 2002 Eric Lee Green
 *
 *
 */
/* Aes Module:
 * This module provides an interface to the Aes encryption algorithm
 *
 * Note: This implementation is generously cribbed from md5module.c, which
 * had some of the same criteria.
 *
 * This was based on the twofish-py module originally written at
 * Enhanced Software Technologies for BRU-Pro prior to its purchase by
 * Atipa Corporation and, later, by The Tolis Group. Release of the
 * code was authorized by the President of EST, Ted Cook, prior to the
 * purchase of EST by Atipa Corporation.  I have stripped out the
 * CFB-8 mode and the CBC mode, leaving nothing but CFB-128 and ECB
 * modes, and converted everything to the Ruby scheme.  If anybody
 * else wants more feedback modes, feel free to add them yourself :-}.
 */

#include <ruby.h>
#include "aes.h"
#include <stdio.h>

typedef unsigned char u1byte;

typedef struct {
  keyInstance encrypt_key;  /* state info for the Aes cipher... */
  keyInstance decrypt_key;  
  int key_gen;  /* a flag for whether we're keyed or not... */
  cipherInstance cipher;
  u1byte cfb_blk[16];
  u1byte cfb_crypt[16]; /* encrypted cfb_blk for CFB128 mode */
  int cfb128_idx;       /* where we are in the CFB128 cfb_blk). */ 
} AesObject;

VALUE eAES;


static VALUE s_init(VALUE self) {
  AesObject *rijnp;

  int result;

  Data_Get_Struct(self,AesObject, rijnp);

  return self;
}
  
static VALUE s_new(VALUE self) {
  AesObject *rijnp;
  VALUE rijnp_data;

  rijnp = ALLOC(AesObject);

  
  cipherInit(&rijnp->cipher,MODE_ECB,NULL);
  rijnp->cfb128_idx=-1;  /* start @ start of cfb_blk... */  
  rijnp->key_gen=0; /* have not done a key yet... */

  rijnp_data=Data_Wrap_Struct(self,0,free,rijnp);

  rb_obj_call_init(rijnp_data,0,NULL);
  return rijnp_data;
}

/* static char new_doc [] = 
"new() -> _aesobject:\n\
\n\
Return a new Aes object\n";
*/

/**************
* The DEL method:
***************/
static void
aes_dealloc(AesObject *rijnp) {
  PyMem_DEL(rijnp);  /* add it back to the free object list. */
}

/***************
 * The set_key method: We require a hex one, unlike the old twofish-py.
 ***************/
static VALUE aes_set_key(VALUE self, VALUE key) {
  unsigned char *hexkey;
  int hexkey_len;   /* sorry, this is what '#' parses :-(. */
  int i;
  AesObject *rijnp;

  Data_Get_Struct(self,AesObject, rijnp);

  Check_Type(key, T_STRING);

  hexkey_len=RSTRING(key)->len;
  hexkey=RSTRING(key)->ptr;


  if (hexkey_len != 16*2 && hexkey_len != 24*2 && hexkey_len != 32*2) {
    /* Okay, we must raise an exception and exit :-(. */
    rb_raise(rb_eArgError, "wrong key length (must be 16, 24, or 32 bytes,not %d)", hexkey_len/2);

    return self; 
  }

  /* okay, we do have it.... */

  /* sigh, double key schedule set up time. Sigh sigh sigh. */
  makeKey(&rijnp->encrypt_key,DIR_ENCRYPT,hexkey_len*4,hexkey);

  makeKey(&rijnp->decrypt_key,DIR_DECRYPT,hexkey_len*4,hexkey);
  rijnp->key_gen=1;    /* voila! */
  
  return self; /* sigh!  */
}

/* static char set_key_doc [] =
"set_key(arg): \n\
   Set the key for this Aes object. The key must be 128, 192, or 256 \n\
  bits in length (16, 24, or 32 characters). Raises aes.error with a\
  value of 'Invalid Key' if the key is any other length. \n\
  \n";
*/

static VALUE aes_encrypt(VALUE self, VALUE args) {
  AesObject *rijnp;
  unsigned char *data;
  int data_len;
  VALUE retval;

  u1byte out_blk[16];

  Check_Type(args, T_STRING); /* make sure it's a string */
  data_len=RSTRING(args)->len;
  data=RSTRING(args)->ptr;

  Data_Get_Struct(self,AesObject, rijnp);
  if (data_len != 16) {
    rb_raise(rb_eArgError, "wrong data length (must be 16 bytes, found %d bytes)", data_len);
    return self;
  }

  if (!rijnp->key_gen) {
    /* sorry, no key has been initialized! */
    rb_raise(eAES,"must set up a key before you can encrypt!");
    return self; 
  }
  /* okay, now to encrypt it: */
   blockEncrypt(&rijnp->cipher,&rijnp->encrypt_key,data,128,out_blk);

  /* encrypt(self->ctx,(u1byte *) data,out_blk); */
  
  /* Now to create a Ruby string out of the result: */
   return rb_str_new((char *)out_blk, 16);

}

/* static char encrypt_doc[] = "encrypt(data)->string\n\
Given a 128-bit data value, returns that value encrypted with the current\n\
key. If the data is other than 128 bits, raises an exception. If no key \n\
has been set up, raises an exception.\n"; 
*/

/** CFB mode helper routines.... */

static VALUE cfb_salt(VALUE self, VALUE args) {
  AesObject *rijnp;
  unsigned char *src;
  unsigned char *dest;
  int src_len;
  int i;

  Check_Type(args,T_STRING); /* make sure it's a string. */
  src=RSTRING(args)->ptr;
  src_len=RSTRING(args)->len;

  if (src_len != 16) {
    rb_raise(rb_eArgError, "wrong data length (must be 16 bytes, found %d bytes)", src_len);
    return self;
  }

   Data_Get_Struct(self,AesObject,rijnp);

   rijnp->cfb128_idx=-1;
   dest=rijnp->cfb_blk;

  for (i=0;i<16;i++) {
    *dest++ = *src++;
  }
  return self;  /* sigh! */
}

/*  static char cfb_salt_doc[] =
"cfb_salt(salt):\n\
\n\
When given a 16-byte salt value, stores it in the _aes object's\n\
internal salt buffer for use in further CFB encryptions.\n";
*/

/* Now for the actual cfb_encrypt routine: */
static VALUE cfb_encrypt(VALUE self,VALUE args) {
  AesObject *rijnp;
  unsigned char *src;
  unsigned char *destdata;
  int srclen;
  int i,ch;
 
  VALUE retvalue;

  Check_Type(args,T_STRING);  /* make sure it's a string. */
  
  src=RSTRING(args)->ptr;
  srclen=RSTRING(args)->len;

   Data_Get_Struct(self,AesObject,rijnp);
  
  /* Now to alloc our result buffer: will be same length as original data. */
  destdata=(unsigned char *) malloc(srclen);

  /* Now to encrypt each individual character in cfb mode: */
  for (i=0; i<srclen; i++) {
    if ((rijnp->cfb128_idx < 0) || (rijnp->cfb128_idx > 15)) {
       blockEncrypt(&rijnp->cipher,&rijnp->encrypt_key,rijnp->cfb_blk,128,rijnp->cfb_crypt);

      /* encrypt(rijnp->ctx,rijnp->cfb_blk,rijnp->cfb_crypt); */
      rijnp->cfb128_idx=0;
    }
    /* XOR the data with a byte from our encrypted buffer. */ 
    ch=src[i] ^ rijnp->cfb_crypt[rijnp->cfb128_idx];
    /* do output feedback: put crypted byte into next block to be crypted */
    rijnp->cfb_blk[rijnp->cfb128_idx++]=ch; 
    destdata[i]=(unsigned char) ch;
  }

  /* Now create a return value: */
  retvalue=rb_str_new((char *)destdata,srclen);
  /* free our malloced memory... */
  free(destdata);
  return retvalue; /* and return the Python string object! */
}

/* static char cfb_decrypt_doc[] =
"cfb_decrypt(data)->string\n\
Given a stream of bytes, returns a stream of bytes decrypted in CFB \n\
(Cipher FeedBack) mode. See Schneir, page 200. Note: Must first call\n\
cfb_salt(data) to set up the initial unique salt value!\n";
*/

/* static char cfb_decrypt128_doc[] =
"cfb_decrypt128(data)->string\n\
Given a stream of bytes, returns a stream of bytes decrypted in CFB128 \n\
(Cipher FeedBack) mode. See Schneir, page 200. Note: Must first call\n\
cfb_salt(data) to set up the initial unique salt value!\n"; 
*/

 
/* Now for the actual cfb_decrypt routine: */
static VALUE cfb_decrypt(VALUE *self, VALUE args)  {
   AesObject *rijnp;
  unsigned char *src;
  unsigned char *destdata;
  int srclen;
  int i;
  unsigned char ch;
  
  VALUE retvalue;

  Check_Type(args, T_STRING); /* make sure it's a string */
  srclen=RSTRING(args)->len;
  src=RSTRING(args)->ptr;

  Data_Get_Struct(self,AesObject, rijnp);

  /* Now to alloc our result buffer: will be same length as original data. */
  destdata=(unsigned char *) malloc(srclen);
 
  /* Now to encrypt each individual character in cfb mode: */
  for (i=0; i<srclen; i++) {
    if (rijnp->cfb128_idx < 0 || rijnp->cfb128_idx > 15) {

      blockEncrypt(&rijnp->cipher,&rijnp->encrypt_key,rijnp->cfb_blk,128,rijnp->cfb_crypt);
           

      /* encrypt(rijnp->ctx,rijnp->cfb_blk,rijnp->cfb_crypt); */
      rijnp->cfb128_idx=0;
    }
    ch=src[i];
    destdata[i]=ch ^ rijnp->cfb_crypt[rijnp->cfb128_idx]; 
    /* do output feedback: put crypted byte into next block to be crypted */
    rijnp->cfb_blk[rijnp->cfb128_idx++]=ch; 
  }
  /* Now create a return value: */
  retvalue=rb_str_new((char *)destdata,srclen);

  /* free our malloced memory... */
  free(destdata);
  return retvalue; /* and return the Python string object! */
}

/* Now for the decrypt: */

static VALUE aes_decrypt(VALUE self,VALUE args) {
  AesObject *rijnp;
  unsigned char *data;
  int data_len;

  u1byte out_blk[16];

  Check_Type(args, T_STRING); /* make sure it's a string */
  data_len=RSTRING(args)->len;
  data=RSTRING(args)->ptr;


  if (data_len != 16) {
    rb_raise(rb_eArgError, "wrong data length (must be 16 bytes, found %d bytes)", data_len);
    return 0;
  }

  Data_Get_Struct(self,AesObject, rijnp);

  if (!rijnp->key_gen) {
    /* sorry, no key has been initialized! */
    rb_raise(eAES,"must set up a key before you can decrypt!");
    return 0; 
  }
   /* okay, now to decrypt it: */

     blockDecrypt(&rijnp->cipher,&rijnp->decrypt_key,data,128,out_blk);
  /*   decrypt(rijnp->ctx,(u1byte *) data,out_blk); */

  
  /* Now to create a Python string out of the result: */
  return rb_str_new((char *)out_blk,16);
}

/* static char decrypt_doc[] = "decrypt(data)->string\n\
Given a 128-bit data value, returns that value decrypted with the current\n\
key. If the data is other than 128 bits, raises an exception. If no key \n\
has been set up, raises an exception.\n";
*/
/*
static char module_doc [] = 
"This module implements the interface to Dr. Brian Gladman's implementation\n\
of Counterpane's 'aes' encryption algorithm. This is the low-level\n\
interface that operates solely upon 128-bit data values. For stream-oriented\n\
variants, see the 'aes.py' module.\n\
\n\
Functions:\n\
new() -- return a new aes object.\n\
\n\
Special Objects: 
\n\
AesType -- type object for Aes objects. \n";
*/
/*
static char AesType_doc [] =
"A AesType represents an object that applies the Aes encryption and\n\
decryption algorithms to data. This low-level type operates solely upon\n\
128-bit data values and accepts 128, 192, or 256-bit key values. \n\
\n\
Methods:\n\
\n\
set_key(key) -- sets the key for further encryption or decryptions\n\
encrypt(val) -- encrypts a 128-bit value, returns a 128-bit encrypted value.\n\
decrypt(val) -- decrypts a 128-bit value, returns a 128-bit decrypted value.\n\
cfb_salt(salt) -- initializes CFB mode counter.\n\
cfb_encrypt(string) -- returns a string of same length encrypted in CFB mode.\n\
cfb_decrypt(string) -- returns a string of same length decrypted in CFB mode.\n\
";
*/
static VALUE error_error(VALUE obj)
{
    return rb_iv_get(obj, "mesg");
}

static VALUE error_errno(VALUE obj)
{
    return rb_iv_get(obj, "errno");
}

void 
Init_aes() 
{
  VALUE cAES=rb_define_class("AES",rb_cObject);
  /* create our own error class: */
   eAES =   rb_define_class("AESError", rb_eStandardError); 
   rb_define_method(eAES, "error", error_error, 0);
   rb_define_method(eAES, "errno", error_errno, 0);
   rb_define_const(eAES,"WRONG_KEY_LENGTH",INT2NUM(1));  

  rb_define_singleton_method(cAES,"new",s_new,0);
  rb_define_singleton_method(cAES,"init",s_init,0);
  rb_define_method(cAES,"initialize",s_init,0);
  
  rb_define_method(cAES,"set_key",aes_set_key,1);
  rb_define_method(cAES,"encrypt",aes_encrypt,1);
  rb_define_method(cAES,"decrypt",aes_decrypt,1);
  rb_define_method(cAES,"cfb_salt",cfb_salt,1);
  rb_define_method(cAES,"cfb_encrypt",cfb_encrypt,1);
  rb_define_method(cAES,"cfb_decrypt",cfb_decrypt,1);
}