#include <u.h>
#include <libc.h>
#include <bio.h>
#include <ctype.h>
#include <auth.h>
#include <fcall.h>
#include <mp.h>
#include <libsec.h>
#include <thread.h>
#include <9p.h>
#include "mysql.h"

#define PATH(type, n)		((type)|((n)<<8))
#define TYPE(path)		((int)(path) & 0xFF)
#define NUM(path)		((uint)(path)>>8)

enum
{
	Qroot,
	Qdatabases,
	Qprocesses,
	Qclone,
	Qn,
		Qctl,
		Qquery,
		Qdata,
		Qformat,
};

typedef struct Client Client;
struct Client {
	int num;	/* slot number */
	int ref;	/* reference count */

	Results *res;	/* results from last query */
	int idx;	/* data line number to read next */
	Layout layout;	/* layout of tables, nil for default */
	char *db;	/* the database we want to use */
};

typedef struct Info Info;
struct Info {
	Results *res;
	int idx;
	Layout layout;
};

static Client **client;
static int nclient;
static Sess *S;

int Verbose;
int Debug;

/* used in mysql V3.23-4.0 though later databases may still have old passwords */
static void
hash(ulong res[2], char *pass)
{
	char *p;
	ulong nr, nr2, tmp, add;

	add = 7;
	nr = 1345345333;
	nr2 = 0x12345671;
	for(p = pass; *p; p++){
		if(*p == '\t' || *p == ' ')
			continue;

		tmp = *p;
		nr ^= (((nr & 63) +add) *tmp) + (nr << 8);
		nr2 += (nr2 << 8) ^ nr;
		add += tmp;
	}
	res[0] = nr & 0x7fffffff; 
	res[1] = nr2 & 0x7fffffff;
}

static uchar
rnd(ulong state[2])
{
	state[0] = (state[0] * 3 + state[1]) % 0x3FFFFFFFL;
	state[1] = (state[0] + state[1] + 33) % 0x3FFFFFFFL;
	return floor((state[0] / (double)0x3FFFFFFFL) *31.0);

}

static void
scramble3x(char *resp, char *salt, char *pass)
{
	int ns;
	char extra, *r;
	ulong hp[2], hs[2], state[2];

	hash(hp, pass);
	hash(hs, salt);
	state[0] = (hp[0] ^ hs[0]) & 0x3FFFFFFFL;
	state[1] = (hp[1] ^ hs[1]) & 0x3FFFFFFFL;

	r = resp;
	ns = strlen(salt);
	while(r < resp+ns)
		*r++ = rnd(state) +64;

	r = resp;
	extra = rnd(state);
	while(r < resp+ns)
		*r++ ^= extra;
}

/* used in mysql V4.1 and later */
static void
scramble4x(uchar *resp, char *salt1, char *salt2, char *pass)
{
	uchar *r, *s, *t;
	DigestState *ds;
	uchar tmp1[SHA1dlen], tmp2[SHA1dlen], tmp3[SHA1dlen];
	
	sha1((uchar *)pass, strlen(pass), tmp1, nil);
	sha1(tmp1, sizeof(tmp1), tmp2, nil);
	ds = sha1((uchar *)salt1, strlen(salt1), nil, nil);
	ds = sha1((uchar *)salt2, strlen(salt2), nil, ds);
	sha1(tmp2, sizeof(tmp2), tmp3, ds);

	r = resp;
	s = tmp3;
	t = tmp1;
	while(r < resp+Nauth)
		*r++ = *s++ ^ *t++;
}


int
newclient(void)
{
	int i;
	Client *c;

	for(i=0; i<nclient; i++)
		if(client[i]->ref==0)
			return i;

	if(nclient%16 == 0)
		client = erealloc9p(client, (nclient+16)*sizeof(client[0]));

	c = emalloc9p(sizeof(Client));
	memset(c, 0, sizeof(*c));
	c->num = nclient;
	client[nclient++] = c;
	return c->num;
}

typedef struct Tab Tab;
struct Tab {
	char *name;
	ulong mode;
};

static Tab tab[] = {
	"/",		DMDIR|0555,
	"databases",	0444,
	"processes",	0444,
	"clone",	0666,
	nil,		DMDIR|0555,
	"ctl",		0666,
	"query",	0666,
	"data",		0666,
	"format",	0444,
};

static void
responderrstr(Req *r)
{
	char e[ERRMAX];
	*e = 0;
	rerrstr(e, sizeof e);
	respond(r, e);
}

static void
fillstat(Dir *d, uvlong path)
{
	Tab *t;

	memset(d, 0, sizeof(*d));
	d->uid = estrdup9p("mysql");
	d->gid = estrdup9p("mysql");
	d->qid.path = path;
	d->atime = d->mtime = time(0);
	t = &tab[TYPE(path)];
	if(t->name)
		d->name = estrdup9p(t->name);
	else{
		d->name = smprint("%ud", NUM(path));
		if(d->name == nil)
			sysfatal("out of memory");
	}
	d->qid.type = t->mode>>24;
	d->mode = t->mode;
}

static int
rootgen(int i, Dir *d, void*)
{
	int j;

	j = i+Qroot+1;
	if(j <= Qclone){
		fillstat(d, j);
		return 0;
	}
	j = i-Qn+1;
	if(j < nclient){
		fillstat(d, PATH(Qn, j));
		return 0;
	}
	return -1;
}

static int
clientgen(int i, Dir *d, void *aux)
{
	Client *c;

	c = aux;
	i += Qn+1;
	if(i <= Qformat){
		fillstat(d, PATH(i, c->num));
		return 0;
	}
	return -1;
}

static void
inforead(Req *r)
{
	char *buf;
	Info *i = r->fid->aux;

	if(r->ifcall.offset == 0)
		i->idx = 0;

	buf = fmtdata(i->res, &i->layout, i->idx);
	if(buf == nil){
 		r->ofcall.count = 0;
		respond(r, nil);
		return;
	}

	i->idx++;
	r->ifcall.offset = 0;
	readstr(r, buf);
	free(buf);
	respond(r, nil);
	return;
}

static void
ctlread(Req *r, Client *c)
{
	char buf[32];

	snprint(buf, sizeof(buf), "%d", c->num);
	readstr(r, buf);
	respond(r, nil);
	return;
}

static void
ctlwrite(Req *r, Client *c)
{
	char *f[3], *s;
	int nf;

	s = emalloc9p(r->ifcall.count+1);
	memmove(s, r->ifcall.data, r->ifcall.count);
	s[r->ifcall.count] = '\0';

	nf = tokenize(s, f, 3);
	if(nf == 0){
		free(s);
		respond(r, nil);
		return;
	}

	if(strcmp(f[0], "use") == 0){
		if(nf != 2)
			goto Badarg;
		else
		if(mysql_use(S, f[1]) == -1){
			free(s);
			responderrstr(r);
			return;
		}
		free(c->db);
		free(S->db);
		c->db = estrdup9p(f[1]);
		S->db = estrdup9p(f[1]);
	}
	else
	if(strcmp(f[0], "headings") == 0){
		if(nf != 1)
			goto Badarg;
		c->layout.headings = 1;
	}
	else
	if(strcmp(f[0], "noheadings") == 0){
		if(nf != 1)
			goto Badarg;
		c->layout.headings = 0;
	}
	else
	if(strcmp(f[0], "padded") == 0){
		if(nf != 1)
			goto Badarg;
		c->layout.delimited = 0;
	}
	else
	if(strcmp(f[0], "delimited") == 0){
		if(nf != 1)
			goto Badarg;
		c->layout.delimited = 1;
	}
	else
	if(strcmp(f[0], "rowsep") == 0){
		if(nf != 2)
			goto Badarg;
		c->layout.rowsep = f[1][0];
	}
	else
	if(strcmp(f[0], "colsep") == 0){
		if(nf != 1)
			goto Badarg;
		c->layout.colsep = f[1][0];
	}

	free(s);
	respond(r, nil);
	return;
Badarg:
	free(s);
	respond(r, "bad ctl arguments");
}

static void
querywrite(Req *r, Client *c)
{
	char *s;

	s = emalloc9p(r->ifcall.count+1);
	memmove(s, r->ifcall.data, r->ifcall.count);
	s[r->ifcall.count] = 0;

	if(c->res){
		freeres(c->res);
		c->res = nil;
	}

	if(c->db == nil){
		respond(r, "no database selected");
		free(s);
		return;
	}
	
	if(strcmp(S->db, c->db) != 0){
		if(mysql_use(S, c->db) == -1){
			responderrstr(r);
			free(s);
			return;
		}
		free(S->db);
		S->db = estrdup9p(c->db);
	}

	if(mysql_query(S, s, &c->res) == -1){
		responderrstr(r);
		free(s);
		return;
	}
	widths(c->res);
	r->ofcall.count	= r->ifcall.count;
	respond(r, nil);	
	free(s);	
}

static void
dataread(Req *r, Client *c)
{
	char *buf;

	if(r->ifcall.offset == 0)
		c->idx = 0;

	buf = fmtdata(c->res, &c->layout, c->idx);
	if(buf == nil){
		r->ofcall.count = 0;
		respond(r, nil);
		return;
	}
	c->idx++;
	r->ifcall.offset = 0;
	readstr(r, buf);
	free(buf);
	respond(r, nil);
	return;
}

static void
formatread(Req *r, Client *c)
{
	char *buf;

	if(r->ifcall.offset == 0)
		c->idx = 0;

	buf = fmtfields(c->res, c->idx);
	if(buf == nil){
		r->ofcall.count = 0;
		respond(r, nil);
		return;
	}
	c->idx++;
	r->ifcall.offset = 0;
	readstr(r, buf);
	free(buf);
	respond(r, nil);
	return;
}

/***************************************************/


static void
fsattach(Req *r)
{
	if(r->ifcall.aname && r->ifcall.aname[0]){
		respond(r, "invalid attach specifier");
		return;
	}
	r->fid->qid.path = PATH(Qroot, 0);
	r->fid->qid.type = QTDIR;
	r->fid->qid.vers = 0;
	r->ofcall.qid = r->fid->qid;
	respond(r, nil);
}

static void
fsstat(Req *r)
{
	fillstat(&r->d, r->fid->qid.path);
	respond(r, nil);
}

static char*
fswalk1(Fid *fid, char *name, Qid *qid)
{
	int i, n;
	char buf[32];
	ulong path;

	path = fid->qid.path;
	if(!(fid->qid.type&QTDIR))
		return "walk in non-directory";

	if(strcmp(name, "..") == 0){
		switch(TYPE(path)){
		case Qn:
			qid->path = PATH(Qroot, 0);
			qid->type = tab[Qroot].mode>>24;
			return nil;
		case Qroot:
			return nil;
		default:
			return "bug in fswalk1";
		}
	}

	i = TYPE(path)+1;
	for(; i<nelem(tab); i++){
		if(i==Qn){
			n = atoi(name);
			snprint(buf, sizeof buf, "%d", n);
			if(n < nclient && strcmp(buf, name) == 0){
				qid->path = PATH(i, n);
				qid->type = tab[i].mode>>24;
				return nil;
			}
			break;
		}
		if(strcmp(name, tab[i].name) == 0){
			qid->path = PATH(i, NUM(path));
			qid->type = tab[i].mode>>24;
			return nil;
		}
		if(tab[i].mode&DMDIR)
			break;
	}
	return "directory entry not found";
}

static void
fsopen(Req *r)
{
	static int need[4] = { 4, 2, 6, 1 };
	ulong path;
	int n;
	Tab *t;
	Info *i;
	Results *res;

	/*
	 * lib9p already handles the blatantly obvious.
	 * we just have to enforce the permissions we have set.
	 */
	path = r->fid->qid.path;
	t = &tab[TYPE(path)];
	n = need[r->ifcall.mode&3];
	if((n&t->mode) != n){
		respond(r, "permission denied");
		return;
	}

	switch(TYPE(path)){
	case Qdatabases:
		if(mysql_query(S, "show databases", &res) == -1){
			responderrstr(r);
			break;
		}
		i = emalloc9p(sizeof(Info));
		widths(res);
		i->res = res;
		r->fid->aux = i;
		respond(r, nil);
		break;
	case Qprocesses:
		if(mysql_ps(S, &res) == -1){
			responderrstr(r);
			break;
		}
		i = emalloc9p(sizeof(Info));
		widths(res);
		i->res = res;
		i->layout.headings = 1;
		r->fid->aux = i;
		respond(r, nil);
		break;
	case Qclone:
		n = newclient();
		path = PATH(Qctl, n);
		r->fid->qid.path = path;
		r->ofcall.qid.path = path;
		if(chatty9p)
			fprint(2, "open clone => path=%lux\n", path);
		t = &tab[Qctl];
		/* fall through */
	default:
		if(t-tab >= Qn)
			client[NUM(path)]->ref++;
		respond(r, nil);
		break;
	}
}


static void
fsread(Req *r)
{
	char e[ERRMAX];
	ulong path;

	path = r->fid->qid.path;
	switch(TYPE(path)){
	default:
		snprint(e, sizeof(e), "bug in fsread type=%d", TYPE(path));
		respond(r, e);
		return;
	case Qroot:
		dirread9p(r, rootgen, nil);
		respond(r, nil);
		break;
	case Qdatabases:
		inforead(r);
		break;
	case Qprocesses:
		inforead(r);
		break;
	case Qn:
		dirread9p(r, clientgen, client[NUM(path)]);
		respond(r, nil);
		break;
	case Qctl:
		ctlread(r, client[NUM(path)]);
		break;
	case Qdata:
		dataread(r, client[NUM(path)]);
		break;
	case Qformat:
		formatread(r, client[NUM(path)]);
		break;
	}
}

static void
fswrite(Req *r)
{
	ulong path;
	char e[ERRMAX];

	path = r->fid->qid.path;
	switch(TYPE(path)){
	default:
		snprint(e, sizeof e, "bug in fswrite type=%d", TYPE(path));
		respond(r, e);
		break;
	case Qquery:
		querywrite(r, client[NUM(path)]);
		break;
	case Qctl:
		ctlwrite(r, client[NUM(path)]);
		break;
//	case Qdata:
//		datawrite(r, client[NUM(path)]);
//		break;
	}
}


static void
fsdestroyfid(Fid *fid)
{
	Info *srcs;

	if(fid->qid.type == Qdatabases || fid->qid.type == Qprocesses){
		srcs = fid->aux;
		freeres(srcs->res);
		free(srcs);
	}
}

Srv fs = 
{
.attach=	fsattach,
.walk1=		fswalk1,
.open=		fsopen,
.read=		fsread,
.write=		fswrite,
.stat=		fsstat,
.destroyfid=	fsdestroyfid,
};


static void
ding(void *u, char *msg)
{
	USED(u);

	if(strstr(msg, "alarm"))
		noted(NCONT);
	noted(NDFLT);
}

void
usage(void)
{
	fprint(2, "usage: %s [-v] [-k name=value...] [-m mntpt] [-s srvname] [-a] [-b] host\n", argv0);
	exits("usage");
}


void
main(int argc, char *argv[])
{
	int flag;
	UserPasswd *up;
	char resp3x[Nauth];
	uchar resp4x[Nauth];
	char *host, *keyp, *mtpt, *svs;

	flag = 0;
	svs = nil;
	keyp = "";
	mtpt = "/mnt/db";
	ARGBEGIN{
	case 'D':
		chatty9p++;
		break;
	case 'v':
		Verbose++;
		break;
	case 'd':
		Debug++;
		break;
	case 'k':
		keyp = EARGF(usage());
		break;
	case 'm':
		mtpt = EARGF(usage());
		break;
	case 's':
		svs = EARGF(usage());
		break;
	case 'a':
		flag |= MAFTER;
		break;
	case 'b':
		flag |= MBEFORE;
		break;
	case '?':
		usage();
		break;
	default:
		fprint(2, "unrecognized option\n");
		usage();
	}ARGEND

	if (argc < 1)
		usage();
	host = argv[0];

	/*
	 * We must get our credentials before opening the session to the database
	 * as the server will timeout after a seccond or two without authentication.
	 * This is not a problem if your key is already in factotum but causes the
	 * session to fail if it is not.
	 */
	if ((up = auth_getuserpasswd(auth_getkey,
	   "server=%s proto=pass service=mysql %s", host, keyp)) == nil)
		sysfatal("cannot get credentials - %r\n");

	notify(ding);
	if((S = mysql_open(host)) == nil)
		sysfatal("%s: session: %r\n", host);
	
	scramble3x(resp3x, S->salt1, up->passwd);
	scramble4x(resp4x, S->salt1, S->salt2, up->passwd);
	if(*up->passwd){
		if(mysql_auth(S, up->user, resp3x, resp4x) != 0)
			sysfatal("%s: auth: %r\n", host);
	}
	else{
		if(mysql_auth(S, up->user, "", nil) != 0)
			sysfatal("%s: auth: %r\n", host);
	}
	memset(up->passwd, 0xff, strlen(up->passwd));

	if(Verbose)
		print("connected host=%s user=%s protocol=v%d server=v%s\n",
			host, up->user, S->proto, S->server);

 	postmountsrv(&fs, svs, mtpt, flag);
	exits(0);
}