#include <u.h>
#include <libc.h>
#include <draw.h>
#include <mp.h>
#include <libsec.h>
#include "dat.h"
#include "fns.h"
#define DBG if(0)
//#undef DBG
//#define DBG

enum
{
	ExponentSize		= 4,
	PaddingSize		= 8,
	MaxKeySize		= 16,
	RandomSize		= 32,
	ModulusSize		= 64,
	MaxModulusSize	= 256,
	MacSize 			= 8,
};

static
struct SecContext
{
	uchar crandom[RandomSize];	/* client random */
	uchar srandom[RandomSize];	/* server random */
	RSApub *spub;				/* server public key */
	int klen;					/* symmetric key length ≤ MaxKeySize */
	uchar ekey[MaxKeySize];		/* encryption key */
	uchar dkey[MaxKeySize];		/* decryption key */
	uchar skey[MaxKeySize];		/* signing key */
	uchar iekey[MaxKeySize]; 	/* initial encryption key */
	uchar idkey[MaxKeySize]; 	/* initial decryption key */
	int ecount;				/* encrypted packet counter */
	int dcount;				/* decrypted packet counter */
	RC4state e;
	RC4state d;
} sc;

static uchar pad1[40] = {
	54, 54, 54, 54, 54, 54, 54, 54,
	54, 54, 54, 54, 54, 54, 54, 54,
	54, 54, 54, 54, 54, 54, 54, 54,
	54, 54, 54, 54, 54, 54, 54, 54,
	54, 54, 54, 54, 54, 54, 54, 54
};
static uchar pad2[48] = {
	92, 92, 92, 92, 92, 92, 92, 92,
	92, 92, 92, 92, 92, 92, 92, 92,
	92, 92, 92, 92, 92, 92, 92, 92,
	92, 92, 92, 92, 92, 92, 92, 92,
	92, 92, 92, 92, 92, 92, 92, 92,
	92, 92, 92, 92, 92, 92, 92, 92
};

enum
{
	ENCRYPTION_METHOD_NONE	= 0x00,
	ENCRYPTION_METHOD_40BIT	= 0x01,
	ENCRYPTION_METHOD_128BIT	= 0x02,
	ENCRYPTION_METHOD_56BIT	= 0x08,
	ENCRYPTION_METHOD_FIPS	= 0x10,

	ENCRYPTION_LEVEL_NONE		= 0,
	ENCRYPTION_LEVEL_LOW		= 1,
	ENCRYPTION_LEVEL_CLIENT_COMPATIBLE	= 2,
	ENCRYPTION_LEVEL_HIGH		= 3,
	ENCRYPTION_LEVEL_FIPS		= 4,
};

int
sizesechdr(int secflags)
{
	if(secflags&SEC_ENCRYPT)
		return  12;	// flags[4] mac[8]
	else if(secflags)
		return 4;		// flags[4]
	return 0;
}

uchar*
prebuf(uchar* buf, int nb, int ndata, int chanid, int secflags)
{
	int n, len, shdsize;
	uchar *p, *ep;

	if(chanid==0)
		chanid = GLOBALCHAN;

	shdsize = sizesechdr(secflags);
	len = TPDATAFIXLEN+8+shdsize+ndata;
	if(len>nb){
		werrstr("%s: provided %d, need %d, data %d", Esmall, nb, len, ndata);
		return nil;
	}
	ep = buf+len;

	ndata = len-TPDATAFIXLEN;
	n = mktpdat(buf, nb, ndata);
	if(n < 0)
		sysfatal("mktpdat: %r");
	p = buf+TPDATAFIXLEN;

	ndata -= 8;
	if(mkmcssdr(p, ep-p, ndata, chanid) < 0)
		sysfatal("mkmcssdr: %r");
	p += 8;

	if(shdsize > 0)
		PLONG(p, secflags);
	return p + shdsize;
}

/* 48-byte hash used to generate master secret (6.1) and key material (6.2.2). */
void
saltedhash(uchar* out, uchar* in, uchar* crandom, uchar* srandom, uchar salt)
{
	DigestState *ds;
	uchar ihash[SHA1dlen], pad[4];
	int i;

	for(i = 0; i < 3; i++){
		memset(pad, salt+i, i+1);
		ds = sha1(pad, i+1, nil, nil);
		sha1(in, 48, nil, ds);
		sha1(crandom, 32, nil, ds);
		sha1(srandom, 32, ihash, ds);
		ds = md5(in, 48, nil, nil);
		md5(ihash, SHA1dlen, out+i*16, ds);
	}
}

/* 16-byte transformation used to generate export keys (6.2.2). */
void
finalhash(uchar* out, uchar* in, uchar* crandom, uchar* srandom)
{
	DigestState *ds;

	ds = md5(in, 16, nil, nil);
	md5(crandom, 32, nil, ds);
	md5(srandom, 32, out, ds);
}

static void
reduceto40(uchar* key)
{
	key[0] = 0xd1;
	key[1] = 0x26;
	key[2] = 0x9e;
}

static void
setsecrets(uchar* crandom, uchar* srandom, ulong emethod)
{
	uchar premaster[48], master[48], keyblock[48];

	memcpy(premaster, crandom, 24);
	memcpy(premaster+24, srandom, 24);
	saltedhash(master, premaster, crandom, srandom, 'A');
	saltedhash(keyblock, master, crandom, srandom, 'X');

	memcpy(sc.skey, keyblock, 16);
	finalhash(sc.dkey, keyblock+16, crandom, srandom);
	finalhash(sc.ekey, keyblock+32, crandom, srandom);

	switch(emethod){
	case ENCRYPTION_METHOD_40BIT:
		sc.klen = 8;
		reduceto40(sc.skey);
		reduceto40(sc.dkey);
		reduceto40(sc.ekey);
		break;
	case ENCRYPTION_METHOD_128BIT:
		sc.klen = 16;
		break;
	default:
		sysfatal("unsupported encryption method %lud", emethod);
	}

	memcpy(sc.iekey, sc.ekey, sc.klen);
	memcpy(sc.idkey, sc.dkey, sc.klen);

	setupRC4state(&sc.e, sc.ekey, sc.klen);
	setupRC4state(&sc.d, sc.dkey, sc.klen);
}

/* Generate MAC (5.2.3.1) */
void
mac(uchar* amac, int nmac, uchar* data, int n, uchar* key, int nkey)
{
	uchar ihash[SHA1dlen], ohash[MD5dlen], buf[4];
	DigestState *ds;

	ds = sha1(key, nkey, nil, nil);
	sha1(pad1, sizeof pad1, nil, ds);
	PLONG(buf, n);
	sha1(buf, 4, nil, ds);
	sha1(data, n, ihash, ds);
	ds = md5(key, nkey, nil, nil);
	md5(pad2, sizeof pad2, nil, ds);
	md5(ihash, SHA1dlen, ohash, ds);
	memcpy(amac, ohash, nmac);
}

void
rdpmac(uchar* amac, int nmac, uchar* data, int n)
{
	mac(amac, nmac, data, n, sc.skey, sc.klen);
}

static void
updatekey(uchar* key, uchar* ukey, int n)
{
	uchar shasig[SHA1dlen], dmd5[MD5dlen];
	DigestState *ds;
	RC4state c;

	ds = sha1(ukey, n, nil, nil);
	sha1(pad1, sizeof pad1, nil, ds);
	sha1(key, n, shasig, ds);

	ds = md5(ukey, n, nil, nil);
	md5(pad2, sizeof pad2, nil, ds);
	md5(shasig, 20, dmd5, ds);
	memcpy(key, dmd5, sc.klen);

	setupRC4state(&c, key, n);
	rc4(&c, key, sc.klen);
	if (sc.klen == 8)
		reduceto40(key);
}

void
rdpencrypt(uchar* a, int n)
{
	rc4(&sc.e, a, n);
	if(++sc.ecount == 4096){
		updatekey(sc.ekey, sc.iekey, sc.klen);
		setupRC4state(&sc.e, sc.ekey, sc.klen);
		sc.ecount = 0;
DBG fprint(2, "ekey updated\n");
	}
}

void
rdpdecrypt(uchar* a, int n)
{
	rc4(&sc.d, a, n);
	if(++sc.dcount == 4096){
		updatekey(sc.dkey, sc.idkey, sc.klen);
		setupRC4state(&sc.d, sc.dkey, sc.klen);
		sc.dcount = 0;
DBG fprint(2, "dkey updated\n");
	}
}

uchar*
rdpdecryptv(uchar* a, int n)
{
	uchar *seal, mac[MacSize];

	seal = a;
	a += MacSize;
	n -= MacSize;

	rdpdecrypt(a, n);
	rdpmac(mac, MacSize, a, n);
	if(memcmp(mac, seal, MacSize) != 0){
		werrstr("message authentication failed");
		return nil;
	}

	return a;
}

int
getecrandom(uchar* buf, int nb)
{
	int n;
	mpint *c;

	c = letomp(sc.crandom, sizeof(sc.crandom), nil);
	rsaencrypt(sc.spub, c, c);
	n = mptole(c, buf, nb, nil);
	mpfree(c);
	return n;
};

/* 2.2.1.10 Client Security Exchange PDU */
void
sendsecxchg(void)
{
	uchar ecr[MaxModulusSize], buf[MaxModulusSize+128], *p;
	ulong n, nb, ndata, len;

	nb = sizeof(buf);
	n = getecrandom(ecr, sizeof(ecr));
	ndata = n+4+PaddingSize;

	p = prebuf(buf, nb, ndata, 0, SEC_EXCHANGE_PKT);
	if(p == nil)
		sysfatal("bufsec: %r");
	
	PLONG(p+0, ndata-4);
	memcpy(p+4, ecr, n);
	memset(p+4+n, 0, PaddingSize);

	len = p-buf+ndata;
	writen(rd.fd, buf, len);
}



enum {
	RSAkeyTag	= 6,
	RSAsealTag	= 8,
	CERT_CHAIN_VERSION_1 = 1,
	CERT_CHAIN_VERSION_2 = 2,
	SIGNATURE_ALG_RSA = 1,
	KEY_EXCHANGE_ALG_RSA = 1,
};

/* [MS-RDPBCGR] 2.2.1.4.3.1.1 Server Proprietary Certificate */
static RSApub*
scanpcert(uchar* p, uchar* ep)
{
	uchar *q;
	ushort tag, len;
	ulong version, sigalg, keyalg, magic, keylen;
	RSApub *rsa = nil;

	version = GLONG(p);
	sigalg = GLONG(p+4);
	keyalg = GLONG(p+8);
	p += 12;

	assert(version == CERT_CHAIN_VERSION_1);
	assert(sigalg == SIGNATURE_ALG_RSA);
	assert(keyalg == KEY_EXCHANGE_ALG_RSA);

	while(p < ep){
		tag = GSHORT(p);
		len = GSHORT(p+2);
		p += 4;
		q = p+len;

		switch (tag){
		case RSAkeyTag:
			magic = GLONG(p);
			keylen = GLONG(p+4);
			p += 8+PaddingSize;
			keylen -= PaddingSize;
			if(magic != 0x31415352){
				werrstr("bad RSA magic: %#lx", magic);
				return nil;
			}
			if(keylen<ModulusSize || keylen>MaxModulusSize){
				werrstr("bad server public key size: %lud bits", keylen*8);
				return nil;
			}
			rsa = rsapuballoc();
			rsa->ek = letomp(p, ExponentSize, nil);
			rsa->n = letomp(p+ExponentSize, keylen, nil);
			break;
		case RSAsealTag:
			/* fixme: verify key signature */
			break;
		}
		p = q;
	}
	return rsa;
}

/* [MS-RDPELE] 2.2.1.4.2 X.509 Certificate Chain (X509 _CERTIFICATE_CHAIN) */
static RSApub*
scanX509chain(uchar* p, uchar* ep)
{
	ulong ncert, len;
	RSApub* rsa;

	ncert = GLONG(p);
	p += 4;
	if(ncert < 2){
		werrstr("X.509 certificate chain too short");
		return nil;
	}
	for(; ncert > 1; ncert--){
		len = GLONG(p);
		p += 4+len;
		if(p >= ep){
			werrstr("short data");
			return nil;
		}
	}

	len = GLONG(p);
	if(p+4+len >= ep){
		werrstr(Eshort);
		return nil;
	}
	rsa = X509toRSApub(p+4, len, nil, 0);
	if(rsa == nil){
		asn1dump(p+4, len);
		werrstr("could not parse X509: %r");
		return nil;
	} else {
//DBG fmtinstall('B', mpfmt);
//DBG X509dump(p+4, len);
	}
	/* fixme: validate server certificate */
	return rsa;
}

/* [MS-RDPBCGR] 2.2.1.4.3.1 Server Certificate */
static RSApub*
servercert(uchar* p, ulong n)
{
	ulong certver;
	certver = GLONG(p);
	if(certver == CERT_CHAIN_VERSION_1)
		return scanpcert(p, p+n);
	else
		return scanX509chain(p+4, p+n);
}

/*
 * 2.2.1.4.3 Server Security Data (TS_UD_SC_SEC1)
 * http://msdn.microsoft.com/en-us/library/cc240518.aspx
 */
void
scansrvsec(uchar* buf, int nbytes)
{
	ulong emethod, elevel, srandomlen, certlen;
	uchar *p, *ep;

	p = buf+4;		/* skip the user data header */
	ep = buf+nbytes;
	if(p+8 > ep)
		sysfatal(Eshort);

	emethod = GLONG(p);
	elevel = GLONG(p+4);
	if(elevel == ENCRYPTION_METHOD_NONE)
		return;
	srandomlen = GLONG(p+8);
	certlen = GLONG(p+12);
	p += 16;

	if(srandomlen != RandomSize)
		sysfatal("random len %lud, expected %d\n", srandomlen, RandomSize);
	if(p+srandomlen > ep)
		sysfatal(Eshort);
	memcpy(sc.srandom, p, srandomlen);
	p += srandomlen;

	if(p+certlen > ep)
		sysfatal(Eshort);
	sc.spub = servercert(p, certlen);
	if(sc.spub == nil)
		sysfatal("error extracting server public key: %r");

	genrandom(sc.crandom, sizeof sc.crandom);
	setsecrets(sc.crandom, sc.srandom, emethod);
}