#include <u.h>
#include <libc.h>
#include <ctype.h>
#include <fcall.h>
#include <thread.h>
#include <9p.h>
#include <ip.h>
#include "ber.h"
#include "snmp.h"


int alarmflag;
int readflag;

char *pduerror[] = {
	"noError",
	"tooBig",
	"noSuchName",
	"badValue",
	"readOnly",
	"genErr",
	"noAccess",
	"wrongType",
	"wrongLength",
	"wrongEncoding",
	"wrongValue",
	"noCreation",
	"inconsistenValue",
	"resourceUnavailable",
	"commitFailed",
	"undoFailed",
	"authorizationError",
	"notWritable",
	"inconsistentName"
};
 


/* debug or error */
int
error(int level, char *format, ... )
{
	va_list args;
	char lasterror[ERRMAX];

	va_start(args, format);
	vsnprint(lasterror,ERRMAX, format, args);
	va_end(args);

	werrstr(lasterror);
	switch(level) {
		case 1:
			fprint(2,"Debug:%r\n");
			return 0;
			break;
		case 2:
			fprint(2,"Error: %r\n");
			abort();
			break;
	}

	return 0;
}


/* creates a new snmp request id */
/* based on rsc code */
int
mkreqid(void) 
{
	return (nsec() & 0x7FFE) + 1;
}

void
freesession(Session *s)
{
	free(s->host);
	free(s->port);
	free(s->rocomm);
	free(s->rwcomm);
	free(s);
}

void
freevarbind(Tuple *t, int len)
{
	int i;

	if ( t == nil || len < 1 )
		return;

	for(i=0;i<len;i++) {
		free(t[i].oid);
		free(t[i].value);
	}

	free(t);
}
		
void
freepacket(Packet *p)
{
	if ( p == nil )
		return;
	freevarbind(p->pdu.varbind,p->pdu.nbind);

	if ( p->comm != nil )
		free(p->comm);

	free(p);
}

int
mkwalk(Packet * p, Packet *o)
{

	int i;
	
	p->ver = o->ver;
	p->comm = o->comm;
	p->pdu.type = GetNextRequest;
	p->pdu.reqid = o->pdu.reqid;
	p->pdu.errstat = 0;
	p->pdu.erridx = 0;
	p->pdu.nbind= o->pdu.nbind;

	freevarbind(p->pdu.varbind,p->pdu.nbind);

	p->pdu.varbind = (Tuple*)emalloc(sizeof(Tuple)*p->pdu.nbind);

	for(i=0;i<p->pdu.nbind;i++){
		p->pdu.varbind[i].oid = smprint("%s",o->pdu.varbind[i].oid);
		p->pdu.varbind[i].value = nil;
	}

	return 1;
}



/* The End-Of-MIB message is not always supported */
/* so we need some heuristics to do correct walks */
/* All oids of a packet must complain to be able to walk */
int
canwalk(Packet *orig, Packet *p1, Packet *p2)
{

	int origlen, oid1len;
	int i;
	
	if ( p1 == nil || p2 == nil ) {
		error(DEBUG,"canwalk(): if no packet we can walk to get one");
		return 1;
	}

	for (i=0;i<p1->pdu.nbind;i++) {

		oid1len = strlen(p1->pdu.varbind[i].oid);
		origlen = strlen(orig->pdu.varbind[i].oid);

		if ( strncmp(p1->pdu.varbind[i].oid, p2->pdu.varbind[i].oid, oid1len) == 0 ) {
			error(DEBUG,"canwalk(): error: p1 and p2 are the same");
			return 0;
		}

		if ( strncmp(orig->pdu.varbind[i].oid,p2->pdu.varbind[i].oid,origlen) != 0 ) {
			error(DEBUG,"canwalk(): error: p2 is out of branch: %s -- %s",orig->pdu.varbind[i].oid,p2->pdu.varbind[i].oid);
			return 0;
		}
	}

	/* if i am here, i can walk other packet */
	return 1;

}

/* PDU is an IMPLICIT SEQUENCE
  * and that means the SEQUENCE id
  * is changed using one of those
  * specified and them should be encoded as they
  * should appear in the packet, ie, no more
  * transformations are done */

Elem
mkpdu(int type, Elist *el)
{
	Elem e;

	e.tag.class = type;
	e.tag.num = SEQUENCE;
	e.val.tag = VSeq;
	e.val.u.seqval = el;

	return e;
}

/* this is used to compose set-request message */
/* only int and octectstring */
Elem
char2asn(char *data)
{
	/* for normal requests */
	if ( data == nil ) {
		error(DEBUG,"char2asn(): data = nil");
		return Null();
	}

	/* for set-request only octect-strings and integers */
	if ( data[0] < '0' || data[0] > '9' ) 
		return mkoctet((uchar*)data,strlen(data));
	else 
		return mkint(atoi(data));
	

}

Elem
oid2asn(char *s)
{
	Elem e;
	Ints *o;
	char *f[MAX_ELEM];
	int i, n;
	char *aux;

	aux = strdup(s);
	n = getfields(aux, f, MAX_ELEM, 0, ".");

	if ( n < 0 || atoi(f[0]) < 1 ) {
		free(aux);
		return e;
	}

	o = (Ints*) emalloc9p(sizeof(Ints)*n);
	o->len = n;
	for(i=0;i<n;i++)
		o->data[i] = atoi(f[i]);
	
	e = mkoid(o);

	error(DEBUG,"oid2asn(): oid = %s o->len = %d",s,o->len);
	free(aux);
	freeints(o);
	return e;

}

char*
asn2char(Elem e)
{
	Value v;
	int i;
	Fmt fmt;
	char *aux;
	mpint *m;

	fmtstrinit(&fmt);

	v = e.val;

	/* chek if i have a valid type before trying to decode it */
	if ( e.tag.class > 30 ) {
		switch(e.tag.class) {
		case 	IpAddress: fmtprint(&fmt,"%V",v.u.octetsval->data); break;
		case Counter:
		case Gauge:  
		case TimeTicks:
		case Opaque:
		case NsapAddress:
		case Counter64:
		case UIInteger32:
			m = betomp(v.u.octetsval->data, v.u.octetsval->len, nil);
			aux = mptoa(m,10,nil,0);
			fmtprint(&fmt,"%s",aux);
			mpfree(m);
			free(aux);
			break;
		case EndOfMIB:
			fmtprint(&fmt,"End of MIB");
			break;
		default:
			fmtprint(&fmt,"unknown class type %x and val tag %x",e.tag.class,e.val.tag);
			break;
		}

		return fmtstrflush(&fmt); 
	}

	/* TODO: complete all known types */
	if ( e.tag.class == Universal )
		switch(v.tag){
		case VBool: fmtprint(&fmt,"BOOLEAN");break;
		case VInt: fmtprint(&fmt,"%d",v.u.intval); break;
		case VOctets: /* heuristic is need to know if it is a printable string or not */	
			if ( isalnum(v.u.octetsval->data[0]) ) {
				aux = (char*)emalloc(v.u.octetsval->len+1);
				memcpy(aux,v.u.octetsval->data, v.u.octetsval->len);
				aux[v.u.octetsval->len] = '\0';
				fmtprint(&fmt,"%s",aux);
				free(aux);
			} else 
				for(i=0;i<v.u.octetsval->len;i++)
					fmtprint(&fmt,"%.2x",v.u.octetsval->data[i]);
			break;
		case VBigInt: fmtprint(&fmt,"BINGINT");break;
		case VReal: fmtprint(&fmt,"REAL");break;	
		case VOther: fmtprint(&fmt,"OTHER"); break;
		case VBitString: fmtprint(&fmt,"BITSTRING");break;
		case VNull: fmtprint(&fmt,"NULL");break;
		case VEOC: fmtprint(&fmt,"EOC");break;
		case VObjId:
			for(i = 0; i<v.u.objidval->len; i++)
				fmtprint(&fmt,"%d%s", v.u.objidval->data[i],(i != (v.u.objidval->len-1)) ? ".":"");	
			break; 
		case VString: fmtprint(&fmt,"%s",v.u.stringval); break;
		case VSeq: fmtprint(&fmt,"SEQUENCE");break;
		case VSet: fmtprint(&fmt,"SET");break;
		default:
			fmtprint(&fmt,"unknown class type %x and val tag %x",e.tag.class,e.val.tag);
			break;
		}

	return fmtstrflush(&fmt);
}

/* ----
Var-Bind-List =SEQUENCE 		
	Var-Bind = SEQUENCE		
		oid	OBJECT IDENTIFIER	
		value	NULL
		oid	OBJECT IDENTIFIER
		value NULL
---*/
Elem
mkvarbind(Packet *p)
{
	Elist *z=nil;
	int i;
	Elem oid;
	Elem value;
	
	for(i=0; i < p->pdu.nbind; i++) {	
		value = char2asn(p->pdu.varbind[i].value);
		oid = oid2asn(p->pdu.varbind[i].oid);
		z = mkel(mkseq(mkel(oid, mkel(value, nil))), z);
	}

	return mkseq(z);
}


/* ----
getBulkRequest PDU ASN1
	Request-ID INTEGER
	Non-Repeaters INTEGER
	Max-Repetition INTEGER
	Var-Bind-List =SEQUENCE 		
---- */

Elist *
mkgetbulkreq(Packet *p)
{
	Elist *e;

	e = mkel(mkpdu((int)p->pdu.type,
			mkel(mkint(p->pdu.reqid),
			mkel(mkint(p->pdu.nrep),
			mkel(mkint(p->pdu.maxrep),
			mkel(mkvarbind(p),nil))))),nil);

	return e;

}

/* -----
getNextRequest PDU ASN1				
		Request-ID INTEGER			
		Error-Status INTEGER			
		Error-Index INTEGER			
		Var-Bind-List =SEQUENCE 		
-----*/	
			
Elist *
mkgetnreq(Packet *p)
{
	Elist *e;
	e =  mkel(mkpdu((int)p->pdu.type,
			mkel(mkint(p->pdu.reqid),
			mkel(mkint(p->pdu.errstat),
			mkel(mkint(p->pdu.erridx),
			mkel(mkvarbind(p),nil))))),nil);
	return e;
}


/* ASN.1 snmp packet encoding 	*/
/* message = SEQUECE	*/
/*		version INTEGER	*/	
/*		community OCTECT_STRING	*/
/*		PDU = IMPLICIT SEQUENCE	*/

int
encpkt(Packet *p, Bytes **pbytes)
{
	int err;
	Bytes *bsnmp;
	Elem snmp;
	Elem ver;
	Elem comm;
	Elist * pdu;
	
	switch(p->pdu.type){
	case GetRequest: 
	case SetRequest:
	case GetNextRequest: pdu = mkgetnreq(p); break;
	case GetBulkRequest: pdu = mkgetbulkreq(p); break;
	default: error(DEBUG,"encpkt(): unknown op\n");return 0; break;
	}
	
	if ( p->comm == nil ) {
		error(DEBUG,"encpkt(): community not found");
		return 0;
	}
	ver = mkint(p->ver);
	comm = mkoctet((uchar*)p->comm,strlen(p->comm));
	snmp = mkseq(mkel(ver, mkel(comm, pdu)));

	err = encode(snmp, &bsnmp);
	if ( err != ASN_OK ) {
		error(DEBUG,"encpkt(): cannot encode the pkt: %d %r\n",err);
		freevalfields(&snmp.val);
		return 0;
	}
	
	*pbytes = bsnmp;
	freevalfields(&snmp.val);
	return 1;
}

/* translates an asn1 stream to the pdu struct */
int
dec_pdu(Elist *el, Pdu * p)
{
	Elem *e, *last;
	Elist *pel;
	Elist *elvar = nil;
	Elist *elbind = nil;
	int tabsize, i;

	e = &el->hd;

	if( is_int(e,&p->reqid) ) {
		pel = el->tl; 
		e = &pel->hd;
		if ( is_int(e,&p->errstat) ) {
			pel = pel->tl; 
			e= &pel->hd;
			if ( is_int(e,&p->erridx) ) {
				pel = pel->tl;
				e = &pel->hd;
			} else {
				error(DEBUG,"dec_pdu(): reqid not found");
				return 0;
			}
		} else {
			error(DEBUG,"dec_pdu(): errstat not found");
			return 0;
		}
	}else {
		error(DEBUG,"dec_pdu(): erridx not found");
		return 0;
	}
	p->nbind = 0;

	if ( is_seq(e,&elvar) ) {
		e = &elvar->hd;
		tabsize = elistlen(elvar);
		if ( tabsize <= 0 ) {
			error(DEBUG,"dec_pdu(): varbind length %d",tabsize);
			return 0;
		}
		p->varbind = (Tuple*)emalloc(sizeof(Tuple)*tabsize);
		for(i=0;i<tabsize;i++) {
			if ( is_seq(e,&elbind) ) {
				pel=elbind;
				last = &pel->hd;
				pel = pel->tl;
				if ( last->val.tag == VObjId ) {
					/* store them in reverse order to correct what the ber code did */
					p->varbind[tabsize-1-i].oid = asn2char(*last);
					p->varbind[tabsize-1-i].value = asn2char(pel->hd);
					p->nbind++;
				} else {
					free(p->varbind);
					error(DEBUG,"dec_pdu(): Error decoding var list");
					return 0;
					}
			} else {
				free(p->varbind);
				error(DEBUG,"dec_pdu(): Error decoding bind list");
				return 0;
				}
			elvar = elvar->tl;
			e = &elvar->hd;
		}
	} else {
		error(DEBUG,"dec_pdu(): Error decoding var list");
		return 0;
		}

	return 1;
}

/* decode a pdu with a quick hack: */
/* the asn1 code treats the unknown objects as octetstring; */
/* so we can encode our pdu with the octet-string tag */
/* and change that tag with the sequence-type that is the correct */
/* ( pdu definition says it is IMPLICIT SEQUENCE and that is a */
/* sequence with the correspondant pdu id 0xa1..0xa8*/

int
pdu2asn(Elem *epdu, Elem *e)
{
	int err;
	Bytes *ocpdu;
 
	if ( encode(*epdu,&ocpdu) == ASN_OK)
		*ocpdu->data = 0x30; /* pdu is a sequence, so we make it that */
	else {
		error(DEBUG,"pdu2asn(): error encoding octet-pdu: ");
		freebytes(ocpdu);
		return 0;
	}
	
	 
	/* decode what we encoded, to obtain the pdu */
	err = decode(ocpdu->data,ocpdu->len,e);
	freebytes(ocpdu);
	
	if ( err != ASN_OK ) {
		error(DEBUG,"pdu2asn(): error decoding pdu: ");
		return 0;
	}

	return 1;
}

int
decpkt(uchar *buff, int n, Packet *p) 
{
	Bytes *bbuff, *bcomm;
	Elem emess, *ever, *ecomm, pdu, *epdu;
	Elist *elmess, *el, *elpdu;
	int pver;

	bbuff = makebytes(buff, n);

	if (decode(bbuff->data,bbuff->len,&emess) != ASN_OK) {
		error(DEBUG,"decpkt: error in ASN.1 decode");
		freebytes(bbuff);
		return 0;
	}

	freebytes(bbuff);

	if( !is_seq(&emess,&elmess)  ) {
		error(DEBUG,"decpkt: cannot recognize message");
		freevalfields(&emess.val);
		return 0;
	}
	if ( elistlen(elmess) != 3 ) {
		error(DEBUG,"decpkt: cannot recognize message: n = %d",elistlen(elmess));
		freevalfields(&emess.val);
		return 0;
	}


	ever = &elmess->hd; el = elmess->tl;
	ecomm = &el->hd; el = el->tl;
	epdu = &el->hd;

	if ( !is_int(ever,&pver) ) {
		error(DEBUG,"decpkt: error decoding version field");
		freevalfields(&emess.val);
		return 0;
	}

	p->ver = pver;
	

	if ( !is_octetstring(ecomm,&bcomm) ) {
		error(DEBUG,"decpkt: error decoding community field");
		freevalfields(&emess.val);
		return 0;
	}
	

	p->comm =  (char*)emalloc9p(bcomm->len);
	memcpy(p->comm,(char*)bcomm->data,bcomm->len);
	p->comm[bcomm->len-1]='\0';

	if ( ! pdu2asn(epdu,&pdu) ) {
		error(DEBUG,"decpkt: error in pdu2asn");
		freevalfields(&pdu.val);
		freevalfields(&emess.val);
		return 0;
	}

	if ( !is_seq(&pdu,&elpdu) ){
		error(DEBUG,"decpkt: error decoding pdu");
		freevalfields(&pdu.val);
		freevalfields(&emess.val);
		return 0;
	}

	if ( ! dec_pdu(elpdu,&p->pdu) ) {
		error(DEBUG,"decpkt(): error decoding pdu field");
		freevalfields(&emess.val);
		return 0;
	}

	freevalfields(&emess.val);
	freevalfields(&pdu.val);
	return 1;
}

/* timeout
  * Code from rsc */
int
alarmtr(void*,  char *why)
{

	if(!readflag)
		return 0;

	if(strcmp(why, "alarm") == 0) {
		alarmflag++;
		return 1;
	}
	return 0;
}



int 
talk(Session *s, Packet *req, Packet *recv)
{
	int n, reqid,fd;
	uchar recvbuf[DMAX_PKT];
	Bytes * snmp = nil;
	
	if ( req->pdu.type == SetRequest )
		req->comm = strdup(s->rwcomm);
	else
		req->comm = strdup(s->rocomm);

	if (! encpkt(req, &snmp)) {
		error(DEBUG,"talk(): Error encoding Packet to ASN1");
		freebytes(snmp);
		return 0;
	}
	reqid = req->pdu.reqid;


	if ((fd = dial(netmkaddr(s->host, "udp", s->port), 0, 0, 0)) < 0) {
		error(DEBUG,"talk(): error dialing %s",s->host);
		freebytes(snmp);
		return 0;
	}

	
	do {	
		n = write(fd, snmp->data, snmp->len);
		
		if (n != snmp->len) {
			error(DEBUG,"talk(): Error writing to %d",fd);
			freebytes(snmp);
			return 0;
		}

		freebytes(snmp);

		/* set alarm to exit the while when timeout reached */
		alarmflag = 0;
		readflag = 1;
		threadnotify(alarmtr, 1);
		alarm(s->timeout*1000);

		do {
			n = read(fd, recvbuf, DMAX_PKT);

			if (n > 0) {
				if ( ! decpkt(recvbuf, n, recv) )
					continue;
				else if ( reqid == recv->pdu.reqid ) {
					close(fd);
					alarm(0);
					return 1;
				}
			}
		} while(!alarmflag);

		if (alarmflag) {
				s->retries--;
				error(DEBUG,"talk(): Timeout %d s. Wil retry %d times",s->timeout,s->retries);
		}
		

	} while (s->retries > 0);
	
	close(fd);
	
	alarm(0);

	return 0;
}


/* from rsc */
char * 
dumppkt(Packet *p) 
{

	int i,e;
	Fmt fmt;

	fmtstrinit(&fmt);
	
	for(i=0;i<p->pdu.nbind;i++) {
		e = p->pdu.errstat;
		if ( p->pdu.erridx != i+1 )
			e=noError;
		fmtprint(&fmt,"((%s) (%s) (%s))\n", p->pdu.varbind[i].oid, p->pdu.varbind[i].value,pduerror[e]);
	}

	return fmtstrflush(&fmt);
}


char *
doget(Session *s, Packet * req)
{
	char *buff;
	Packet *recv;
	recv = (Packet*)emalloc(sizeof(Packet));

	if ( talk(s, req, recv) ) {
		buff = dumppkt(recv);
		freepacket(recv);
		return buff;
		
	} else {
		error(DEBUG,"doget(): error talking snmp: %r");
		freepacket(recv);
		return nil;
	}
	
}


char *
dowalk(Session *s, Packet *orig)
{
	Packet *recv;
	Packet *req;
	Packet *last;
	char *aux;
	Fmt fmt;

	fmtstrinit(&fmt);

	recv= (Packet*)emalloc9p(sizeof(Packet));

	if ( ! talk(s, orig, recv) ) {
			error(DEBUG,"dowalk(): error talking snmp");
			freepacket(recv);
			return nil;
	}
	aux = dumppkt(recv);
	fmtprint(&fmt,"%s", aux);
	free(aux);
	last = nil;

	while ( canwalk(orig,last, recv) ) {
		
		req = (Packet*)emalloc9p(sizeof(Packet));

		mkwalk(req,recv);

		freepacket(last);
	
		last = recv;
		recv = (Packet*)emalloc9p(sizeof(Packet));
		

		if ( ! talk(s, req, recv) ) {
			error(DEBUG,"dowalk(): error talking snmp");
			freepacket(recv);
			return nil;
		}
		aux = dumppkt(recv);
		fmtprint(&fmt,"%s", aux);
		free(aux);
		freepacket(req);
	
	}
	freepacket(last);
	freepacket(recv);
	return fmtstrflush(&fmt);
}

char *
dobulk(Session *s, Packet * req)
{
	return dowalk(s, req);
	
}

char *
doset(Session *s, Packet * req)
{
	return doget(s, req);

}

char*
dogetn(Session *s, Packet *req)
{
	return doget(s,req);
}

char *
dosnmp(Session *s, Packet *req)
{
	char *buff;

	/* GetNextRequest is used to walk the snmp tree, */
	/* while WalkRequest is used to identify the operation */
	switch(req->pdu.type) {
		case GetRequest:
			buff = doget(s, req);
			break;
		case GetNextRequest:
			buff = dogetn(s,req);
			break;
		case WalkRequest:
			req->pdu.type = GetNextRequest;
			buff = dowalk(s,req);
			break;
		case GetBulkRequest:
			buff = dobulk(s, req);
			break;
		case SetRequest:
			buff = doset(s, req);
			break;
		default:
			buff = smprint("(dosnmp(): Unknown PDU type %X)\n",req->pdu.type);
			break;
	}

	if ( buff == nil )
		buff =smprint("(%r)\n");

	return buff;
}