/*
 * [MS-RDPELE] 2.2.2 Licensing PDU (TS_LICENSING_PDU)
 * http://msdn.microsoft.com/en-us/library/cc241913.aspx
 *
 * Most licensing crypto code below is broken - just good enough to
 * get through the basic protocol.
 */
#include <u.h>
#include <libc.h>
#include <mp.h>
#include <libsec.h>
#include <draw.h>
#include "dat.h"
#include "fns.h"
#define DBG if(0)

static void	scanlrq(uchar*,uchar*);
static void scanlchal(uchar*,uchar*);
static void	scanlnew(uchar*,uchar*);

enum
{
	RandomSize		= 32,
	MaxKeySize		= 16,
	PaddingSize		= 8,
	ModulusSize		= 64,

	LicenseTokenSize	= 10,
	LicenseHwidSize	= 20,
	LicenseMacSize	= 16,

	LICENSE_REQUEST=				0x01,
	PLATFORM_CHALLENGE=			0x02,
	NEW_LICENSE=				0x03,
	UPGRADE_LICENSE=				0x04,
	LICENSE_INFO=				0x12,
 	NEW_LICENSE_REQUEST=		0x13,
	PLATFORM_CHALLENGE_RESPONSE=	0x15,
	ERROR_ALERT=				0xFF,
};

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 */
	RC4state e;
	RC4state d;
} lc;

void
scanlicensepdu(uchar* p, uchar* ep)
{
	uchar tag;

	tag = p[0];

	switch (tag){
	case LICENSE_REQUEST:
		scanlrq(p+4, ep);
		break;
	case PLATFORM_CHALLENGE:
		scanlchal(p+4, ep);
		break;
	case NEW_LICENSE:
		scanlnew(p+4, ep);
		break;
	case UPGRADE_LICENSE:
		break;
	case ERROR_ALERT:
		break;
	}
}


/* 2.2.2.3 Client License Information (CLIENT_LICENSE_INFO) */
void
sendlicense(uchar * client_random, uchar * rsa_data,
		uchar * licence_data, int licence_size, uchar * hwid, uchar * signature)
{
	uchar buf[4096], *p;
	int len, nb, ndata;

	nb = sizeof(buf);
	ndata = 24+licence_size+RandomSize+ModulusSize
		+PaddingSize+LicenseHwidSize+LicenseMacSize;
	p = prebuf(buf, nb, ndata, 0, SEC_LICENSE_PKT);
	if(p == nil)
		sysfatal("bufsec: %r");
	len = p-buf+ndata;

	// 2.2.1.12.1.1 Licensing Preamble (LICENSE_PREAMBLE)
	p[0] = LICENSE_INFO;
	p[1] = 2;			// version: 2=RDP4, 3=RDP5+
	PSHORT(p+2, ndata);

	PLONG(p+4, 1);		// PreferredKeyExchangeAlg: 1=RSA
	PLONG(p+8, 0x02010000);	// PlatformId: vendor=MS, clientos=NT4
	memcpy(p+12, client_random, RandomSize);	// ClientRandom[32]
	p += 12+RandomSize;

	// EncryptedPreMasterSecret[*]
	PSHORT(p+0, 0);
	PSHORT(p+2, (ModulusSize+PaddingSize));
	memcpy(p+4, rsa_data, ModulusSize);
	p += 4+ModulusSize;
	memset(p, PaddingSize, 0);
	p += PaddingSize;

	// LicenseInfo[*]
	PSHORT(p+0, 1);
	PSHORT(p+2, licence_size);
	memcpy(p+4, licence_data, licence_size);
	p += 4+licence_size;

	// EncryptedHWID[*]
	PSHORT(p+0, 1);
	PSHORT(p+2, LicenseHwidSize);
	memcpy(p+4, hwid, LicenseHwidSize);
	p += 4+LicenseHwidSize;
	
	memcpy(p, signature, LicenseMacSize);	// MACData[16]
	p += LicenseMacSize;

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

/* 2.2.2.2 Client New License Request (CLIENT_NEW_LICENSE_REQUEST) */
enum {
	BB_CLIENT_USER_NAME_BLOB=		0x000F,
	BB_CLIENT_MACHINE_NAME_BLOB=	0x0010,
};
static void
reqlicense(uchar* crandom, uchar* rsa_data, char* user, char *host)
{
	uchar buf[180], *p;
	int len, nb, ndata, usersize, hostsize;

	usersize = strlen(user)+1;
	hostsize = strlen(host)+1;
	ndata = 24+usersize+hostsize+RandomSize+ModulusSize+PaddingSize;
	nb = sizeof(buf);

	p = prebuf(buf, nb, ndata, 0, SEC_LICENSE_PKT);
	if(p == nil)
		sysfatal("reqlicense: %r");
	len = p-buf+ndata;
	// 2.2.1.12.1.1 Licensing Preamble (LICENSE_PREAMBLE)
	p[0] = NEW_LICENSE_REQUEST;
	p[1] = 2;			// version: 2=RDP4, 3=RDP5+
	PSHORT(p+2, ndata);

	PLONG(p+4, 1);		// PreferredKeyExchangeAlg: 1=RSA
	PLONG(p+8, 0xff010000);	// PlatformId
	memcpy(p+12, crandom, RandomSize);	// ClientRandom[32]
	p += 12+RandomSize;

	// EncryptedPreMasterSecret[*]
	PSHORT(p+0, 0);
	PSHORT(p+2, (ModulusSize+PaddingSize));
	memcpy(p+4, rsa_data, ModulusSize);
	p += 4+ModulusSize;
	memset(p, PaddingSize, 0);
	p += PaddingSize;

	// ClientUserName[*]
	PSHORT(p+0, BB_CLIENT_USER_NAME_BLOB);
	PSHORT(p+2, usersize);
	memcpy(p+4, user, usersize);
	p+= 4+usersize;

	// ClientMachineName[*]
	PSHORT(p+0, BB_CLIENT_MACHINE_NAME_BLOB);
	PSHORT(p+2, hostsize);
	memcpy(p+4, host, hostsize);
	p+= 4+hostsize;

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

static int
loadlicense(unsigned char **data)
{
	char *home, *path;
	int fd, n;
	uchar buf[8192], *p, *ep;

	home = getenv("home");
	if (home == nil)
		return -1;
	path = smprint("/sys/log/rdplic.d/%s", rd.local);
	if(path == nil)
		return -1;
	fd = open(path, OREAD);
	free(path);
	if(fd < 0)
		return -1;
	p = buf;
	ep = buf + sizeof(buf);
	while((n = read(fd, p, ep-p)) > 0)
		p += n;
	close(fd);
	n = p-buf;
	*data = malloc(n);
	if(*data == nil)
		sysfatal("readlicense: strdup: %r");
	memcpy(*data, buf, n);
	return n;
}

static void
savelicense(uchar *data, int len)
{
	char *home, *path;
	int fd;

	home = getenv("home");
	if (home == nil)
		return;
	path = smprint("/sys/log/rdplic.d/%s", rd.local);
	if(path == nil)
		sysfatal("writelicense: smprint: %r");
	fd = create(path, OWRITE | OTRUNC, 0600);
	if(fd < 0){
		fprint(2, "writelicense: create: %s: %r\n", path);
		free(path);
		return;
	}
	if(write(fd, data, len) != len){
		fprint(2, "writelicense: write: %s: %r\n", path);
		remove(path);
	}
	close(fd);
	free(path);
}

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

	/*
	 * [MS-RPDELE] 5.1.3 Generating the Licensing Encryption and MAC Salt Keys
	 */
	saltedhash(master, premaster, crandom, srandom, 'A');
	saltedhash(keyblock, master, srandom, crandom, 'A');
	memcpy(lc.skey, keyblock, 16);
	finalhash(lc.ekey, keyblock+16, crandom, srandom);
}

static void
gethwid(uchar * hwid)
{
	PLONG(hwid, 2);
	strncpy((char *) (hwid + 4), rd.local, LicenseHwidSize - 4);
}

static void
scanlrq(uchar *p, uchar *ep)
{
	uchar null_data[ModulusSize];
	uchar *srandom;
	uchar amac[LicenseMacSize];
	uchar hwid[LicenseHwidSize];
	uchar *license;
	int nlicense;

	if(p+RandomSize > ep)
		sysfatal("scanlrq: %s", Eshort);

	srandom = p;
	p += RandomSize;
	USED(p);

	genrandom(lc.crandom, sizeof lc.crandom);
	setlicensesecrets(lc.crandom, srandom, null_data);

	nlicense = loadlicense(&license);
	if(nlicense <= 0){
DBG fprint(2, "requesting a license\n");
		reqlicense(lc.crandom, null_data, rd.user, rd.local);
		return;
	}

DBG fprint(2, "presenting a license\n");
	gethwid(hwid);
	mac(amac, 16, hwid, sizeof(hwid), lc.skey, 16);

	setupRC4state(&lc.e, lc.skey, 16);
	rc4(&lc.e, hwid, sizeof(hwid));

	sendlicense(null_data, null_data, license, nlicense, hwid, amac);
	free(license);
	return;

}

/* 2.2.2.5 Client Platform Challenge Response (CLIENT_PLATFORM_CHALLENGE_RESPONSE) */
static void
respondlchal(uchar * token, uchar * crypt_hwid, uchar * signature)
{
	uchar buf[128], *p;
	int len, nb, ndata;

	ndata = 8+LicenseTokenSize+4+LicenseHwidSize+LicenseMacSize;
	nb = sizeof(buf);

	p = prebuf(buf, nb, ndata, 0, SEC_LICENSE_PKT);
	if(p == nil)
		sysfatal("prebuf: %r");
	// 2.2.1.12.1.1 Licensing Preamble (LICENSE_PREAMBLE)
	p[0] = PLATFORM_CHALLENGE_RESPONSE;
	p[1] = 2;			// version: 2=RDP4, 3=RDP5+
	PSHORT(p+2, ndata);
	PSHORT(p+4, 1);
	PSHORT(p+6, LicenseTokenSize);
	memcpy(p+8, token, LicenseTokenSize);
	p += 8+LicenseTokenSize;
	PSHORT(p+0, 1);
	PSHORT(p+2, LicenseHwidSize);
	memcpy(p+4, crypt_hwid, LicenseHwidSize);
	p += 4+LicenseHwidSize;
	memcpy(p, signature, LicenseMacSize);
	p += LicenseMacSize;

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

/* Process an authentication request packet */
static void
scanlchal(uchar *p, uchar *ep)
{
	uchar *in_token;
	uchar out_token[LicenseTokenSize], decrypt_token[LicenseTokenSize];
	uchar hwid[LicenseHwidSize];
	uchar sealed_buffer[LicenseTokenSize + LicenseHwidSize];
	uchar out_sig[LicenseMacSize];
	uint tokenlen;

	if(p+8 > ep)
		sysfatal("scanlchal: %s", Eshort);

	tokenlen = GSHORT(p+6);
	if(tokenlen != LicenseTokenSize){
		fprint(2, "scanlchal: bad token len %d\n", tokenlen);
		return;
	}

	in_token = p+8;

	memcpy(out_token, in_token, LicenseTokenSize);

	setupRC4state(&lc.d, lc.ekey, 16);
	memcpy(decrypt_token, in_token, LicenseTokenSize);
	rc4(&lc.d, decrypt_token, LicenseTokenSize);

	gethwid(hwid);
	memcpy(sealed_buffer, decrypt_token, LicenseTokenSize);
	memcpy(sealed_buffer + LicenseTokenSize, hwid, LicenseHwidSize);

	mac(sealed_buffer, sizeof(sealed_buffer), lc.skey, 16, out_sig, 16);

	setupRC4state(&lc.e, lc.ekey, 16);
	rc4(&lc.e, hwid, LicenseHwidSize);

DBG fprint(2, "responding to the licensing authentication challenge\n");
	respondlchal(out_token, hwid, out_sig);
}

/* Process an licence issue packet */
static void
scanlnew(uchar *p, uchar *ep)
{
	uint length;
	uint check;
	int i;
DBG fprint(2, "got a license\n");

	if(p+2*1+2 > ep){
		werrstr(Eshort);
		return;
	}

	p += 2*1;
	length = GSHORT(p);
	p += 2;

	if(p+length > ep){
		werrstr(Eshort);
		return;
	}

	setupRC4state(&lc.d, lc.ekey, 16);
	rc4(&lc.d, p, length);

	check = GSHORT(p);
	p += 2;
	if (check != 0)
		return;

	p += 2*1;

	/* advance to 4th string */
	length = 0;
	for(i = 0; i < 4; i++){
		p += length;
		length = GLONG(p);
		p += 4;
		if(p+length > ep){
			werrstr(Eshort);
			return;
		}
	}

	savelicense(p, length);
	rd.licensed = 1;
}