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

#include "webfs.h"
#include "json.h"
#include "acme.h"
#include "misc.h"
#include "rsa.h"
#include "factotum.h"

Directory*
parse_directory(char *dir)
{
	Webfs *w;
	Directory *d;

	w = webfs(dir);
	if(w == nil)
		return nil;

	char *bod = webfsget(w);
	webfsfree(w);
	if(bod == nil)
		return nil;

	JSON *j = jsonparse(bod);
	free(bod);
	if(j == nil)
		return nil;

	d = mallocz(sizeof(Directory), 1);
	if(d == nil){
		jsonfree(j);
		return nil;
	}

	d->directory = estrdup(dir);
	d->newreg = jgetstr(j, "new-reg");
	d->newauthz = jgetstr(j, "new-authz");
	d->newcert = jgetstr(j, "new-cert");
	d->revoke = jgetstr(j, "revoke-cert");

	if(d->newreg == nil ||
		d->newauthz == nil ||
		d->newcert == nil ||
		d->revoke == nil)
			goto Error;

	jsonfree(j);

	return d;

Error:
	jsonfree(j);
	free(d->newreg);
	free(d->newauthz);
	free(d->newcert);
	free(d->revoke);
	return nil;
}

JSON*
registration_json(char *agree)
{
	JSON *reg;

	reg = jalloc(JSONObject);
	jappend(reg, estrdup("resource"), jstrdup("new-reg"));
	jappend(reg, estrdup("agreement"), jstrdup(agree));

	return reg;
}

char*
get_nonce(char *directory)
{
	char *key = "Replay-Nonce";
	char *val;
	Webfs *w;

	w = webfs(directory);
	if(w == nil)
		return nil;

	free(webfsget(w));
	val = webfsgethdr(w, key);
	if(val == nil){
		webfsfree(w);
		return nil;
	}

	return val;
}

char*
signed_request(RSApub *pub, char *nonce, char *payload)
{
	char *pa64, *e, *m, *h, *pr, *pr64, *u, *s64, *data;
	uchar *s, digest[SHA2_256dlen];
	int slen;

	print("payload\n%s\n", payload);
	pa64 = base64url((uchar*)payload, strlen(payload));

	e = mpint2b64u(pub->ek);
	m = mpint2b64u(pub->n);
	h = json_fmt_header_rsa(e, m);

	print("header\n%s\n", h);

	pr = json_fmt_protected_rsa(e, m, nonce);

	print("protected\n%s\n", pr);

	pr64 = base64url((uchar*)pr, strlen(pr));

	u = smprint("%s.%s", pr64, pa64);
	if(u == nil)
		sysfatal("smprint: %r");

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

	s64 = data = nil;

	s = factotum_sign(pub, digest, SHA2_256dlen, &slen);
	if(s == nil)
		goto Error;
	s64 = base64url(s, slen);

	data = json_fmt_signed(h, pr64, pa64, s64);

Error:
	free(s64);
	free(s);
	free(u);
	free(pr64);
	free(pr);
	free(h);
	free(m);
	free(e);
	free(pa64);

	return data;
}

static char*
simplereq(Directory *d, RSApub *key, char *url, char *payload)
{
	char *nonce = get_nonce(d->directory);
	if(nonce == nil){
		werrstr("unable to get nonce: %r");
		return nil;
	}

	char *body = signed_request(key, nonce, payload);
	free(nonce);

	Webfs *w = webfs(url);
	if(w == nil){
		free(body);
		return nil;
	}

	char *resp = webfspost(w, body);
	free(body);
	webfsfree(w);

	return resp;
}

// do a registration. returns registration json object on success.
char*
donewreg(Directory *d, RSApub *key, char *agreement)
{
	JSON *j = registration_json(agreement);
	if(j == nil)
		return nil;

	char *ag = smprint("%J", j);
	jsonfree(j);
	if(ag == nil)
		return nil;

	char *resp = simplereq(d, key, d->newreg, ag);
	free(ag);
	return resp;
}

char*
donewcert(Directory *d, RSApub *key, uchar *cert, int certlen)
{
	char *cert64, *payload, *nonce, *body, *resp;
	Webfs *w;

	cert64 = base64url(cert, certlen);
	payload = json_fmt_newcert(cert64);
	free(cert64);

	nonce = get_nonce(d->directory);
	if(nonce == nil){
		werrstr("unable to get nonce: %r");
		return nil;
	}

	body = signed_request(key, nonce, payload);

	free(nonce);
	free(payload);

	w = webfs(d->newcert);
	if(w == nil){
		free(body);
		return nil;
	}

	resp = webfspost(w, body);
	free(body);

	if(resp == nil){
		webfsfree(w);
		return nil;
	}

	print("cert resp: %s\n", resp);

	free(resp);

	char *loc = webfsgethdr(w, "Location");
	webfsfree(w);
	if(loc == nil)
		return nil;

	return loc;
}

char*
donewauthz(Directory *d, RSApub *key, char *value)
{
	char *payload = json_fmt_newauthz("dns", value);
	char *resp = simplereq(d, key, d->newauthz, payload);
	free(payload);
	return resp;
}

int
dochallenge(Directory *d, RSApub *key, ACMEChallenge *c, char *www)
{
	int n;
	char path[512], content[512], *thumb, *payload;

	path[0] = 0;
	content[0] = 0;

	thumb = jthumbprint(key);

	snprint(path, sizeof(path), "%s/%s", www, c->token);
	snprint(content, sizeof(content), "%s.%s\n", c->token, thumb);

	print("creating %s...\n", path);

	int fd = create(path, OWRITE|ORCLOSE, 0755);
	if(fd < 0){
		free(thumb);
		return 0;
	}

	n = write(fd, content, strlen(content));
	if(n < 0)
		goto Error;

	payload = json_fmt_challenge(c->token, thumb);

	print("sending challenge request...\n");

	char *resp = simplereq(d, key, c->uri, payload);
	free(payload);
	if(resp == nil)
		goto Error;

	free(resp);

	print("waiting for challenge to succeed...\n");

	// now we naively sit here and wait for it to work..
	while(c->status != 1 && c->retries++ < 6){
		print("%lud... ", c->retries);
		Webfs *w = webfs(c->uri);
		resp = webfsget(w);
		if(resp != nil){
			//print("%s:\n%s\n", c->uri, resp);
			JSON *j = jsonparse(resp);
			if(j != nil){
				char *st = jgetstr(j, "status");
				if(st != nil){
					if(strcmp(st, "valid") == 0){
						c->status = 1;
					}
				}
				free(st);
			}
			jsonfree(j);
		}
		free(resp);
		webfsfree(w);

		sleep(c->status == 0 ? 10 : 0);
	}

	print("ok\n");

	return 1;
Error:
	free(thumb);
	close(fd);
	return 0;
}

int
rsa_verify_pub(RSApub *pub)
{
	int rv;

	if(pub->n == nil){
		werrstr("missing public modulus");
		return 0;
	}

	if(mpcmp(pub->ek, mptwo) == -1){
		werrstr("public exponent too small");
		return 0;
	}

	mpint *mpmax = itomp(1<<31-1, nil);
	assert(mpmax != nil);

	rv = mpcmp(pub->ek, mpmax);
	mpfree(mpmax);
	if(rv == 1){
		werrstr("public exponent too large");
		return 0;
	}

	return 1;
}

int
rsa_verify_key(RSApriv *p)
{
	mpint *modulus = nil;
	mpint *congruence = nil, *de = nil, *pminus1 = nil;

	if(!rsa_verify_pub(&p->pub)){
		return 0;
	}

	// Check that Πprimes == n.
	if(mpcmp(p->p, mpone) == -1){
		werrstr("p < 1");
		goto fail;
	}

	if(mpcmp(p->q, mpone) == -1){
		werrstr("q < 1");
		goto fail;
	}

	modulus = mpcopy(mpone);
	assert(modulus != nil);
	mpmul(p->p, modulus, modulus);
	mpmul(p->q, modulus, modulus);

	if(mpcmp(modulus, p->pub.n) != 0){
		werrstr("p*q != n");
		goto fail;
	}

	// Check that de ≡ 1 mod p-1, for each prime.
	// This implies that e is coprime to each p-1 as e has a multiplicative
	// inverse. Therefore e is coprime to lcm(p-1,q-1,r-1,...) =
	// exponent(ℤ/nℤ). It also implies that a^de ≡ a mod p as a^(p-1) ≡ 1
	// mod p. Thus a^de ≡ a mod n for all a coprime to n, as required.
	congruence = mpnew(0);
	de = mpcopy(p->pub.ek);
	mpmul(de, p->dk, de);
	pminus1 = mpnew(0);

	mpsub(p->p, mpone, pminus1);
	mpmod(de, pminus1, congruence);
	if(mpcmp(congruence, mpone) != 0){
		werrstr("p mod check failed");
		goto fail;
	}

	mpsub(p->q, mpone, pminus1);
	mpmod(de, pminus1, congruence);
	if(mpcmp(congruence, mpone) != 0){
		werrstr("q mod check failed");
		goto fail;
	}

	return 1;

fail:
	mpfree(modulus);
	mpfree(congruence);
	mpfree(de);
	mpfree(pminus1);
	return 0;
}