#include <u.h>
#include <libc.h>
#include <auth.h>
#include <mp.h>
#include <libsec.h>
#include <json.h>

#include "json.h"
#include "misc.h"

RSApub*
parse_rsa(char *keydata)
{
	char *a;
	Attr *att;
	RSApub *pub;

	att = _parseattr(keydata);
	if(att == nil)
		return nil;

	pub = rsapuballoc();

	if((a=_strfindattr(att, "ek"))==nil || (pub->ek=strtomp(a, nil, 16, nil))==nil)
		goto Error;
	if((a=_strfindattr(att, "n"))==nil || (pub->n=strtomp(a, nil, 16, nil))==nil)
		goto Error;

	return pub;

Error:
	rsapubfree(pub);
	werrstr("bad rsa public key");
	return nil;
}

RSApriv*
getrsakey(char *keydata, int needprivate, Attr **pa)
{
	char *p;
	int sz;
	RSApriv *key;
	int regen;
	Attr *a;

	key = mallocz(sizeof(RSApriv), 1);
	if(key == nil)
		return nil;

	if(strncmp(keydata, "key ", 4) != 0){
		werrstr("bad key format");
		return nil;
	}

	regen = 0;
	a = _parseattr(keydata+4);
	if(a == nil){
		werrstr("empty key");
		return nil;
	}
	if((p = _strfindattr(a, "proto")) == nil){
		werrstr("no proto");
		return nil;
	}
	if(strcmp(p, "rsa") != 0){
		werrstr("proto not rsa");
		return nil;
	}
	if((p = _strfindattr(a, "ek")) == nil){
		werrstr("no ek");
		return nil;
	}
	if((key->pub.ek = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		werrstr("bad ek");
		return nil;
	}
	if((p = _strfindattr(a, "n")) == nil){
		werrstr("no n");
		return nil;
	}
	if((key->pub.n = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		werrstr("bad n");
		return nil;
	}
	if((p = _strfindattr(a, "size")) == nil)
		fprint(2, "warning: missing size; will add\n");
	else if((sz = strtol(p, &p, 10)) == 0 || *p != 0)
		fprint(2, "warning: bad size; will correct\n");
	else if(sz != mpsignif(key->pub.n))
		fprint(2, "warning: wrong size (got %d, expected %d); will correct\n",
			sz, mpsignif(key->pub.n));
	if(!needprivate)
		goto call;
	if((p = _strfindattr(a, "!dk")) == nil){
		werrstr("no !dk");
		return nil;
	}
	if((key->dk = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		werrstr("bad !dk");
		return nil;
	}
	if((p = _strfindattr(a, "!p")) == nil){
		werrstr("no !p");
		return nil;
	}
	if((key->p = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		werrstr("bad !p");
		return nil;
	}
	if((p = _strfindattr(a, "!q")) == nil){
		werrstr("no !q");
		return nil;
	}
	if((key->q = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		werrstr("bad !q");
		return nil;
	}
	if((p = _strfindattr(a, "!kp")) == nil){
		fprint(2, "warning: no !kp\n");
		regen = 1;
		goto regen;
	}
	if((key->kp = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		fprint(2, "warning: bad !kp\n");
		regen = 1;	
		goto regen;
	}
	if((p = _strfindattr(a, "!kq")) == nil){
		fprint(2, "warning: no !kq\n");
		regen = 1;	
		goto regen;
	}
	if((key->kq = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		fprint(2, "warning: bad !kq\n");
		regen = 1;	
		goto regen;
	}
	if((p = _strfindattr(a, "!c2")) == nil){
		fprint(2, "warning: no !c2\n");
		regen = 1;	
		goto regen;
	}
	if((key->c2 = strtomp(p, &p, 16, nil)) == nil || *p != 0){
		fprint(2, "warning: bad !c2\n");
		regen = 1;	
		goto regen;
	}
regen:
	if(regen){
		RSApriv *k2;

		k2 = rsafill(key->pub.n, key->pub.ek, key->dk, key->p, key->q);
		if(k2 == nil){
			werrstr("regenerating chinese-remainder parts failed: %r");
			return nil;
		}
		key = k2;
	}
call:
	a = _delattr(a, "ek");
	a = _delattr(a, "n");
	a = _delattr(a, "size");
	a = _delattr(a, "!dk");
	a = _delattr(a, "!p");
	a = _delattr(a, "!q");
	a = _delattr(a, "!c2");
	a = _delattr(a, "!kp");
	a = _delattr(a, "!kq");
	if(pa)
		*pa = a;
	return key;
}

char*
jthumbprint(RSApub *pub)
{
	char *e, *m, *t;
	uchar digest[SHA2_256dlen];

	e = mpint2b64u(pub->ek);
	m = mpint2b64u(pub->n);
	t = json_fmt_thumb_rsa(e, m);

	free(e);
	free(m);

	sha2_256((uchar*)t, strlen(t), digest, nil);

	free(t);

	return base64url(digest, SHA2_256dlen);
}

static uchar sha256prefix[19] = {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20};

// assumes sha256.
uchar*
signpkcs1v15(RSApriv *key, uchar *in, int inlen, int *outlen)
{
	int hashlen, tlen, clen, k, i;
	uchar digest[SHA2_256dlen], *em;
	mpint *message, *c;

	hashlen = SHA2_256dlen;
	tlen = sizeof(sha256prefix) + hashlen;

	k = (mpsignif(key->pub.n)+7)/8;

	if(k < tlen+11){
		werrstr("message too long");
		return nil;
	}

	sha2_256(in, inlen, digest, nil);

	em = mallocz(k, 1);
	if(em == nil)
		sysfatal("mallocz: %r");

	em[1] = 1;

	for(i = 2; i < k-tlen-1; i++){
		em[i] = 0xFF;
	}

	memcpy(em + (k - tlen), sha256prefix, sizeof(sha256prefix));
	memcpy(em + (k - hashlen), digest, hashlen);

	message = betomp(em, k, nil);
	print("digest: %B\n", message);
	c = rsadecrypt(key, message, nil);

	clen = (mpsignif(c)+7)/8;

	memset(em, 0, k - clen);

	mptober(c, em, clen);

	*outlen = k;

	return em;
}