/* camellia-rb.c
 *
 * Copyright (c) 2008-2009
 * NTT (Nippon Telegraph and Telephone Corporation) . All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer as
 *   the first lines of this file unmodified.
 * 2. Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY NTT ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL NTT BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#include <ruby.h>
#include "camellia.h"
#include <stdio.h>

// for version before Ruby 1.8.5
#ifndef RSTRING_PTR
#  define RSTRING_PTR(s) (RSTRING(s)->ptr)
#endif
#ifndef RSTRING_LEN
#  define RSTRING_LEN(s) (RSTRING(s)->len)
#endif

typedef unsigned char u1byte;

// error class
static VALUE eCamellia;

// camellia object
typedef struct {
	int				key_len;		// key lengh
	KEY_TABLE_TYPE	key;			// sub keys
	int				key_gen;		// key initialized flag
	u1byte			cfb_blk[16];	// buffer for CFB
	u1byte			cfb_crypt[16];	// buffer for CFB
	int				cfb128_idx;		// index for buffer of CFB
	u1byte			cbc_blk[16];	// buffer for CBC
	u1byte			cbc_crypt[16];	// buffer for CBC
	int				cbc128_idx;		// index for buffer of CBC
	u1byte			cbc_pchar;		// padding char(default:0x05)
} camelliaObject;

/*
# return a new camellia object
*/
static VALUE s_new(VALUE self)
{

	camelliaObject *camellia;
	VALUE camellia_data;

	camellia = ALLOC(camelliaObject);
	camellia->key_gen = 0;	// key is not initialized
	camellia->cfb128_idx = -2;	// CFB index is initialized
	camellia->cbc128_idx = -2;	// CBC index is initialized
	camellia->cbc_pchar = 0x05;
	camellia->key_gen = 0;

	camellia_data = Data_Wrap_Struct(self, 0, free, camellia);

	//rb_obj_call_init(camellia_data, 0, NULL);
	return camellia_data;
}

/*
# set_key method
*/
static VALUE set_key(VALUE self, VALUE key)
{
	unsigned char *skey;
	int skey_len;
	camelliaObject *camellia;

	Data_Get_Struct(self, camelliaObject, camellia);

	Check_Type(key, T_STRING);
	skey_len = RSTRING_LEN(key);
	skey = (unsigned char *)RSTRING_PTR(key);

	// check key length
	if (skey_len != 16 && skey_len != 24 && skey_len != 32)
	{
    	rb_raise(rb_eArgError, "wrong key length (must be 16, 24, or 32 bytes,not %d)", skey_len);
		return self; 
	}

	// key generation
	camellia->key_len = skey_len*8;
	Camellia_Ekeygen(camellia->key_len, skey, camellia->key);
	camellia->key_gen = 1;	//key is initialized

	//return self;
	return Qnil;
}

/*
# encryption method
*/
static VALUE encrypt(VALUE self, VALUE args)
{
	camelliaObject *camellia;
	unsigned char *data;
	int data_len;
	u1byte out_blk[16];

	Check_Type(args, T_STRING);
	data_len = RSTRING_LEN(args);
	data = (unsigned char *)RSTRING_PTR(args);

	Data_Get_Struct(self, camelliaObject, camellia);

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

	// check if key is initialized
	if (!camellia->key_gen)
	{
		rb_raise(eCamellia,"must set up a key before you can encrypt!");
		return self; 
	}

	// encryption
   Camellia_EncryptBlock(camellia->key_len, data, camellia->key, out_blk);

	return rb_str_new((char *)out_blk, 16);
}

/*
# decryption method
*/
static VALUE decrypt(VALUE self,VALUE args)
{
	camelliaObject *camellia;
	unsigned char *data;
	int data_len;
	u1byte out_blk[16];

	Check_Type(args, T_STRING);
	data_len = RSTRING_LEN(args);
	data = (unsigned char *)RSTRING_PTR(args);

	// check data length
	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, camelliaObject, camellia);

	// check if key is initialized
	if (!camellia->key_gen)
	{
		rb_raise(eCamellia,"must set up a key before you can decrypt!");
		return 0; 
	}

	// decryption
	Camellia_DecryptBlock(camellia->key_len, data, camellia->key, out_blk);

	return rb_str_new((char *)out_blk,16);
}

/*
# CFB IV set method
*/
static VALUE cfb_salt(VALUE self, VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	unsigned char *dest;
	int src_len;
	int i;

	Check_Type(args,T_STRING);
	src = (unsigned char *)RSTRING_PTR(args);
	src_len = RSTRING_LEN(args);

	// check IV length
	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, camelliaObject, camellia);

	// CFB initial set
	camellia->cfb128_idx = -1;
	dest = camellia->cfb_blk;
	for (i = 0; i < 16; i++)
	{
		*dest++ = *src++;
	}
	
	return self;
}

/*
# CFB encryption method
*/
static VALUE cfb_encrypt(VALUE self,VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	unsigned char *destdata;
	int srclen;
	int i,ch;
	VALUE retvalue;

	Check_Type(args,T_STRING);
	src = (unsigned char *)RSTRING_PTR(args);
	srclen = RSTRING_LEN(args);

	Data_Get_Struct(self,camelliaObject,camellia);

	// check if key is initialized
	if (!camellia->key_gen)
	{
		rb_raise(eCamellia,"must set up a key before you can cfb_encrypt!");
		return self; 
	}

	// check if IV is initialized
	if (camellia->cfb128_idx != -1)
	{
		rb_raise(eCamellia,"must set up a salt before you can cfb_encrypt!");
		return self; 
	}

	destdata = (unsigned char *)malloc(srclen);

	// CFB encryption
	for (i = 0; i < srclen; i++)
	{
	   	if ((camellia->cfb128_idx < 0) || (camellia->cfb128_idx > 15))
		{
			Camellia_EncryptBlock(camellia->key_len, camellia->cfb_blk, camellia->key, camellia->cfb_crypt);
			camellia->cfb128_idx=0;
		}
		ch = src[i] ^ camellia->cfb_crypt[camellia->cfb128_idx];
		camellia->cfb_blk[camellia->cfb128_idx++] = ch; 
		destdata[i] = (unsigned char)ch;
	}

	retvalue = rb_str_new((char *)destdata, srclen);
	free(destdata);
	return retvalue;
}

/*
# CFB decryption method
*/
static VALUE cfb_decrypt(VALUE *self, VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	unsigned char *destdata;
	int srclen;
	int i;
	unsigned char ch;
	VALUE retvalue;

	Check_Type(args, T_STRING);
	srclen = RSTRING_LEN(args);
	src = (unsigned char *)RSTRING_PTR(args);

	Data_Get_Struct(self,camelliaObject, camellia);

	// check if key is initialized
	if (!camellia->key_gen)
	{
		rb_raise(eCamellia,"must set up a key before you can cfb_decrypt!");
		return 0; 
	}

	// check if IV is initialized
	if (camellia->cfb128_idx != -1)
	{
		rb_raise(eCamellia,"must set up a salt before you can cfb_decrypt!");
		return 0; 
	}

	destdata = (unsigned char *)malloc(srclen);

	// CFB decryption
	for (i = 0; i < srclen; i++)
	{
		if (camellia->cfb128_idx < 0 || camellia->cfb128_idx > 15)
		{
			Camellia_EncryptBlock(camellia->key_len,  camellia->cfb_blk, camellia->key, camellia->cfb_crypt);
			camellia->cfb128_idx=0;
		}
		ch = src[i];
		destdata[i] = ch ^ camellia->cfb_crypt[camellia->cfb128_idx]; 
		camellia->cfb_blk[camellia->cfb128_idx++] = ch; 
	}

	retvalue=rb_str_new((char *)destdata, srclen);
	free(destdata);
	return retvalue;
}

/*
# CBC IV set
*/
static VALUE cbc_salt(VALUE self, VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	unsigned char *dest, *dest2;
	int src_len;
	int i;

	Check_Type(args,T_STRING);
	src = (unsigned char *)RSTRING_PTR(args);
	src_len = RSTRING_LEN(args);

	// check IV length
	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, camelliaObject, camellia);

	// CBC initial set
	camellia->cbc128_idx = -1;
	dest = camellia->cbc_blk;
	dest2 = camellia->cbc_crypt;
	for (i = 0; i < 16; i++)
	{
		*dest++ = *src;
		*dest2++ = *src++;
	}
	
	return self;
}

/*
# CBC padding charactor set
*/
static VALUE cbc_pchar(VALUE self, VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	int src_len;

	Check_Type(args,T_STRING);
	src = (unsigned char *)RSTRING_PTR(args);
	src_len = RSTRING_LEN(args);

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

	Data_Get_Struct(self, camelliaObject, camellia);

	// pchar set
	camellia->cbc_pchar = *src;
	
	return self;
}

/*
# CBC encryption method
*/
static VALUE cbc_encrypt(VALUE self,VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	unsigned char *destdata;
	int srclen;
	int i, j, ch, destidx;
	VALUE retvalue;

	Check_Type(args,T_STRING);
	src = (unsigned char *)RSTRING_PTR(args);
	srclen = RSTRING_LEN(args);

	Data_Get_Struct(self,camelliaObject,camellia);

	// check if key is initialized
	if (!camellia->key_gen)
	{
		rb_raise(eCamellia,"must set up a key before you can cbc_encrypt!");
		return self; 
	}

	// check if IV is initialized
	if (camellia->cbc128_idx != -1)
	{
		rb_raise(eCamellia,"must set up a salt before you can cbc_encrypt!");
		return self; 
	}

	if (srclen % 16 == 0)
		destdata = (unsigned char *)malloc(srclen);
	else
		destdata = (unsigned char *)malloc((srclen/16+1)*16);

	// CBC encryption
	destidx = 0;
	camellia->cbc128_idx = 0;
	for (i = 0; i < srclen; i++)
	{
		ch = src[i] ^ camellia->cbc_crypt[camellia->cbc128_idx];
		camellia->cbc_blk[camellia->cbc128_idx++] = ch; 

	   	if (camellia->cbc128_idx == 16)
		{
			Camellia_EncryptBlock(camellia->key_len, camellia->cbc_blk, camellia->key, camellia->cbc_crypt);
			camellia->cbc128_idx=0;
			for (j = 0; j < 16; j++)
			{
				destdata[destidx++] = camellia->cbc_crypt[j];
			}
		}
	}
	if (srclen % 16 != 0)
	{
		while (camellia->cbc128_idx < 16)
		{
			camellia->cbc_blk[camellia->cbc128_idx] = camellia->cbc_pchar ^ camellia->cbc_crypt[camellia->cbc128_idx++];
		}
	
		Camellia_EncryptBlock(camellia->key_len, camellia->cbc_blk, camellia->key, camellia->cbc_crypt);
		for (j = 0; j < 16; j++)
		{
			destdata[destidx++] = camellia->cbc_crypt[j];
		}
	}

	if (srclen % 16 == 0)
		retvalue = rb_str_new((char *)destdata, srclen);
	else
		retvalue = rb_str_new((char *)destdata, (srclen/16+1)*16);
	free(destdata);
	return retvalue;
}

/*
# CBC decryption method
*/
static VALUE cbc_decrypt(VALUE *self, VALUE args)
{
	camelliaObject *camellia;
	unsigned char *src;
	unsigned char *destdata;
	int srclen;
	int i, j, destidx;
	u1byte tmp[16];
	VALUE retvalue;

	Check_Type(args, T_STRING);
	srclen = RSTRING_LEN(args);
	src = (unsigned char *)RSTRING_PTR(args);
	
	Data_Get_Struct(self,camelliaObject, camellia);

	// check if key is initialized
	if (!camellia->key_gen)
	{
		rb_raise(eCamellia,"must set up a key before you can cbc_decrypt!");
		return 0; 
	}

	// check if IV is initialized
	if (camellia->cbc128_idx != -1)
	{
		rb_raise(eCamellia,"must set up a salt before you can cbc_decrypt!");
		return 0; 
	}

	destdata = (unsigned char *)malloc(srclen);

	// CBC decryption
	destidx = 0;
	camellia->cbc128_idx = 0;
	for (i = 0; i < srclen; i++)
	{
		tmp[camellia->cbc128_idx] = camellia->cbc_blk[camellia->cbc128_idx];
		camellia->cbc_blk[camellia->cbc128_idx++] = src[i];

		if (camellia->cbc128_idx == 16)
		{
			Camellia_DecryptBlock(camellia->key_len,  camellia->cbc_blk, camellia->key, camellia->cbc_crypt);
			camellia->cbc128_idx=0;
			
			for (j = 0; j < 16; j++)
			{
				destdata[destidx++] = camellia->cbc_crypt[j] ^ tmp[j];
			}
		}
	}

	// remove paddingchar
	while (destdata[srclen-1] == camellia->cbc_pchar)
	{
		srclen--;
	}

	retvalue=rb_str_new((char *)destdata, srclen);
	free(destdata);
	return retvalue;
}

// 
void Init_camellia() 
{
	VALUE cCamellia = rb_define_class("Camellia", rb_cObject);
	rb_define_alloc_func(cCamellia, s_new);
	
	//rb_define_singleton_method(cCamellia, "new", s_new, 0);
	rb_define_private_method(cCamellia, "initialize", set_key, 1);
	//rb_define_method(cCamellia, "set_key", set_key, 1);
	rb_define_method(cCamellia, "encrypt", encrypt, 1);
	rb_define_method(cCamellia, "decrypt", decrypt, 1);
	rb_define_method(cCamellia, "cfb_salt", cfb_salt, 1);
	rb_define_method(cCamellia, "cfb_encrypt", cfb_encrypt, 1);
	rb_define_method(cCamellia, "cfb_decrypt", cfb_decrypt, 1);
	rb_define_method(cCamellia, "cbc_salt", cbc_salt, 1);
	rb_define_method(cCamellia, "cbc_pchar", cbc_pchar, 1);
	rb_define_method(cCamellia, "cbc_encrypt", cbc_encrypt, 1);
	rb_define_method(cCamellia, "cbc_decrypt", cbc_decrypt, 1);
	
	eCamellia = rb_define_class("Error", rb_eStandardError);
	//rb_define_method(eCamellia, "error", error_error, 0);
	//rb_define_method(eCamellia, "errno", error_errno, 0);

	rb_define_const(eCamellia, "WRONG_KEY_LENGTH", INT2NUM(1));
}