#include <u.h>
#include <libc.h>
#include <bio.h>
#include <mp.h>
#include <libsec.h>

char *openssh_magic = "openssh-key-v1";

#define GET4(p)	(u32int)(p)[3] | (u32int)(p)[2]<<8 | (u32int)(p)[1]<<16 | (u32int)(p)[0]<<24

int
vunpack(uchar *p, int n, char *fmt, va_list a)
{
	uchar *p0 = p, *e = p+n;
	u32int u;
	mpint *m;
	void *s;

	for(;;){
		switch(*fmt++){
		case '\0':
			return p - p0;
		case '_':
			if(++p > e) goto err;
			break;
		case '.':
			*va_arg(a, void**) = p;
			break;
		case 'b':
			if(p >= e) goto err;
			*va_arg(a, int*) = *p++;
			break;
		case 'm':
			if(p+4 > e) goto err;
			u = GET4(p), p += 4;
			if(u > e-p) goto err;
			m = va_arg(a, mpint*);
			betomp(p, u, m), p += u;
			break;
		case 's':
			if(p+4 > e) goto err;
			u = GET4(p), p += 4;
			if(u > e-p) goto err;
			*va_arg(a, void**) = p;
			*va_arg(a, int*) = u;
			p += u;
			break;
		case '[':
			s = va_arg(a, void*);
			u = va_arg(a, int);
			if(u > e-p) goto err;
			memmove(s, p, u);
			p += u;
			break;
		case 'u':
			if(p+4 > e) goto err;
			u = GET4(p);
			*va_arg(a, int*) = u;
			p += 4;
			break;
		}
	}
err:
	return -1;
}

int
unpack(uchar *p, int n, char *fmt, ...)
{
	va_list a;
	va_start(a, fmt);
	n = vunpack(p, n, fmt, a);
	va_end(a);
	return n;
}

static void
rsadump(uchar *buf, int len)
{
	int rv;
	mpint *n, *e, *d, *p, *q, *iqmp;
	char *comment;
	int commentlen;
	RSApriv *key;

	n = mpnew(0);
	e = mpnew(0);
	d = mpnew(0);
	p = mpnew(0);
	q = mpnew(0);
	iqmp = mpnew(0);


	rv = unpack(buf, len, "mmmmmms",
		n, e, d, iqmp, p, q, &comment, &commentlen);

	if(rv < 0)
		sysfatal("malformed private key");

	fmtinstall('B', mpfmt);
	quotefmtinstall();

	key = rsafill(n, e, d, p, q);
	if(key == nil)
		sysfatal("rsafill: %r");

	Biobuf bio;
	Binit(&bio, 1, OWRITE);

	if(Bprint(&bio, "key proto=rsa service=ssh role=client size=%d ek=%B !dk=%B n=%B !p=%B !q=%B !kp=%B !kq=%B !c2=%B comment=%.*q\n",
		mpsignif(key->pub.n), key->pub.ek,
		key->dk, key->pub.n, key->p, key->q,
		key->kp, key->kq, key->c2, commentlen, comment) < 0 || Bterm(&bio) < 0)
		sysfatal("write: %r");
}

void
usage(void)
{
	sysfatal("usage");
}

void
main(int argc, char *argv[])
{;
	char *file, *buf;
	uchar *bin, *p;
	int fd, len, rv;
	long n, tot;

	ARGBEGIN{
	}ARGEND

	fmtinstall('H', encodefmt);

	if(argc > 1)
		usage();

	if(argc == 1)
		file = argv[0];
	else
		file = "#d/0";

	if((fd = open(file, OREAD)) < 0)
		sysfatal("open %s: %r", file);
	buf = nil;
	tot = 0;
	for(;;){
		buf = realloc(buf, tot+8192);
		if(buf == nil)
			sysfatal("realloc: %r");
		if((n = read(fd, buf+tot, 8192)) < 0)
			sysfatal("read: %r");
		if(n == 0)
			break;
		tot += n;
	}
	buf[tot] = 0;

	bin = decodePEM(buf, "OPENSSH PRIVATE KEY", &len, nil);
	if(bin == nil)
		sysfatal("bad key: %r");

	p = bin;

	n = strlen(openssh_magic);
	if(strncmp((char*)p, openssh_magic, strlen(openssh_magic)) != 0)
		sysfatal("bad magic");

	p += n+1;
	len -= n+1;

	char *cipher, *kdf, *kdfopts, *pub, *priv;
	int cipherlen, kdflen, kdfoptslen, numkeys, publen, privlen;

	rv = unpack(p, len, "sssuss", &cipher, &cipherlen,
		&kdf, &kdflen, &kdfopts, &kdfoptslen, &numkeys,
		&pub, &publen, &priv, &privlen);

	if(rv < 0)
		sysfatal("bad key");

	if(strncmp(cipher, "none", 4) != 0)
		sysfatal("key is encrypted");

	if(numkeys != 1)
		sysfatal("multikey not supported");

	char *type, *privdata;
	int check1, check2, typelen;

	rv = unpack((uchar*)priv, privlen, "uus.", &check1, &check2, &type, &typelen, &privdata);
	if(rv < 0)
		sysfatal("malformed private key");

	if(check1 != check2)
		sysfatal("bad magic");

	if(strncmp(type, "ssh-rsa", 7) != 0)
		sysfatal("not rsa");

	n = privlen - rv;

	rsadump((uchar*)privdata, n);

	exits(nil);
}