/*
 * Worldvisions Tunnel Vision Software:
 *   Copyright (C) 1998 Worldvisions Computer Technology, Inc.
 * 
 * Streams with built-in cryptography on read/write.  See wvcrypto.h.
 */
#include "wvcrypto.h"
#include "strutils.h"
#include <blowfish.h>
#include <rsa.h>



////////////////////////// WvCryptoStream



WvCryptoStream::WvCryptoStream(WvStream *_slave) : WvStreamClone(&slave)
{
    slave = _slave;
}


size_t WvCryptoStream::uread(void *buf, size_t size)
{
    size_t len;
    void *ebuf = new char[size];

    len = WvStreamClone::uread(ebuf, size);
    decrypt(ebuf, buf, len);
    
    delete ebuf;
    return len;
}


size_t WvCryptoStream::uwrite(const void *buf, size_t size)
{
    size_t len;
    void *ebuf = new char[size];
    
    encrypt(buf, ebuf, size);
    len = WvStreamClone::uwrite(ebuf, size);

    delete ebuf;
    return len;
}



/////////////////////// WvXORStream



WvXORStream::WvXORStream(WvStream *_slave, unsigned char _xorvalue)
		: WvCryptoStream(_slave)
{
    xorvalue = _xorvalue;
}


void WvXORStream::encrypt(const void *inbuf, void *outbuf, size_t size)
{
    const unsigned char *i = (const unsigned char *)inbuf;
    unsigned char *o = (unsigned char *)outbuf;
    
    for (; size > 0; size--)
	*o++ = (*i++) ^ xorvalue;
}


void WvXORStream::decrypt(const void *inbuf, void *outbuf, size_t size)
{
    encrypt(inbuf, outbuf, size);
}



////////////////////////// WvBlowfishStream



WvBlowfishStream::WvBlowfishStream(WvStream *_slave, const void *_key,
				   size_t keysize)
		: WvCryptoStream(_slave)
{
    key = new BF_KEY;
    BF_set_key(key, keysize, (unsigned char *)_key);
    ennum = denum = 0;
    memset(envec, 0, sizeof(envec));
    memset(devec, 0, sizeof(devec));
}


void WvBlowfishStream::encrypt(const void *inbuf, void *outbuf, size_t size)
{
    BF_cfb64_encrypt((unsigned char *)inbuf,
		     (unsigned char *)outbuf, size, key,
		     envec, &ennum, BF_ENCRYPT);
}


void WvBlowfishStream::decrypt(const void *inbuf, void *outbuf, size_t size)
{
    BF_cfb64_encrypt((unsigned char *)inbuf,
		     (unsigned char *)outbuf, size, key,
		     devec, &denum, BF_DECRYPT);
}


WvBlowfishStream::~WvBlowfishStream()
{
    delete key;
}



////////////////////////////// WvRSAKey



WvRSAKey::WvRSAKey(char *_keystr, bool priv)
{
    unsigned char *keybuf = new unsigned char[strlen(_keystr)/2], *bufp;
    char *keystr;
    RSA *rp;
    
    keystr = strdup(_keystr);
    
    unhexify(keybuf, keystr);
    bufp = keybuf;
    rp = rsa = RSA_new();
    
    if (priv)
    {
	rsa = d2i_RSAPrivateKey(&rp, &bufp, strlen(keystr)/2);
	prv = keystr;
	
	size_t size;
	unsigned char *iend = keybuf;
	size = i2d_RSAPublicKey(rsa, &iend);
	pub = (char *)malloc(size * 2 + 1);
	hexify(pub, keybuf, size);
    }
    else
    {
	rsa = d2i_RSAPublicKey(&rp, &bufp, strlen(keystr)/2);
	prv = NULL;
	pub = keystr;
    }
    
    delete keybuf;
}


WvRSAKey::WvRSAKey(int bits)
{
    size_t size;
    unsigned char *keybuf, *iend;
    
    rsa = RSA_generate_key(bits, 3, NULL, NULL);
    
    size = i2d_RSAPrivateKey(rsa, NULL);
    iend = keybuf = new unsigned char[size];
    i2d_RSAPrivateKey(rsa, &iend);
    
    prv = (char *)malloc(size * 2 + 1);
    hexify(prv, keybuf, size);
    
    iend = keybuf;
    size = i2d_RSAPublicKey(rsa, &iend);
    
    pub = (char *)malloc(size * 2 + 1);
    hexify(pub, keybuf, size);
    
    delete keybuf;
}


WvRSAKey::~WvRSAKey()
{
    if (prv)
	free(prv);
    if (pub)
	free(pub);
    if (rsa)
	RSA_free(rsa);
}



/////////////////////////// WvRSAStream



WvRSAStream::WvRSAStream(WvStream *_slave,
			 WvRSAKey &_my_key, WvRSAKey &_their_key)
		: WvCryptoStream(_slave),
		  my_key(_my_key.private_str(), true),
		  their_key(_their_key.public_str(), false)
{
    // we always want to read encrypted data in multiples of RSA_size.
    slave->queuemin(RSA_size(my_key.rsa));
}


WvRSAStream::~WvRSAStream()
{
    // remove our strange queuing requirements
    slave->queuemin(0);
}


size_t WvRSAStream::uread(void *buf, size_t size)
{
    if (!my_key.rsa)
	return 0;
    
    size_t len, rs = RSA_size(my_key.rsa);
    unsigned char *tbuf = new unsigned char[rs], *tb2 = new unsigned char[rs];
    
    if (size > rs)
	size = rs;
    
    len = WvStreamClone::uread(tbuf, rs);
    if (len < rs)
	return 0;
    
    // FIXME - if too much data arrived in one packet, we discard it!
    decrypt(tbuf, tb2, rs);
    if (decrypt_silly > size)
	decrypt_silly = size;
    memcpy(buf, tb2, decrypt_silly);

    delete tbuf;
    delete tb2;
	   
    return decrypt_silly;
}


size_t WvRSAStream::uwrite(const void *buf, size_t size)
{
    if (!their_key.rsa)
	return 0;
    
    size_t off, rs = RSA_size(my_key.rsa);
    unsigned char *tbuf = new unsigned char[rs];
    
    for (off = 0; off < size; off += rs/2)
    {
	encrypt((char *)buf + off, tbuf,
		size - off < rs/2 ? size - off : rs/2);
	WvStreamClone::uwrite(tbuf, rs);
    }
    
    delete tbuf;
    return size;
}


void WvRSAStream::encrypt(const void *inbuf, void *outbuf, size_t size)
{
    size_t outsz;
    outsz = RSA_public_encrypt(size, (unsigned char *)inbuf, 
			       (unsigned char *)outbuf, their_key.rsa,
			       RSA_PKCS1_PADDING);
}


void WvRSAStream::decrypt(const void *inbuf, void *outbuf, size_t size)
{
    size_t outsz;
    outsz = RSA_private_decrypt(size, (unsigned char *)inbuf,
				(unsigned char *)outbuf, my_key.rsa,
				RSA_PKCS1_PADDING);
    decrypt_silly = outsz;
}
