#include <sys/types.h>
#include <sys/time.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>


#include <broccoli.h>

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#define NARGS 64

char *host_default = "127.0.0.1";
char *port_default = "47758";
char *host_str;
char *port_str;

// functions and notes
// utility functions
char *PING =	"ping";	 	// ping bro instance connected to
char *STATS =	"show-cpu";	// cpu data 
char *MEM =	"show-memory";	// memory data
char *CONN =	"show-conn";	// conection data
char *HELP =	"help";

// scan detect related functions
char *ISHDROP =	"is-dropped";	// see if provided IP has been dropped
char *SCAN =	"show-scan";	// show scan related numbers assosciated with an IP
char *DROP =	"drop";		// drop an IP
char *UNDROP =	"undrop";	// take an IP off the dropped host list, see below or this is pointless
char *SCANRESET = "reset-scan";	// reset a given IP address' scan related values
char *ALLOWDROP = "allow-drop";	// turn drop functionality on (default bahavior = F)
char *SCANPAIR = "scan-pair";	// 

// sample policy modifications
// note that *what* can be changes is more or less limited to tables and sets given that most of
// what bro considers as dynamic can not be modified after runtime start.
char *ONSCAN = "onscan";	// turn scan detection on
char *OFFSCAN = "offscan";	// turn scan detection off

char *TEST ="test";
char *AC ="ac";
char *NU ="null";

// prototypes
size_t bufsplit(char *buf, size_t n, char **a);
static void bro_pong_record(BroConn *conn, BroRecord *rec);
static void addr_check_record(BroConn *conn, BroRecord *rec);
static void return_cpu(BroConn *conn, double *interval);
static void return_memory(BroConn *conn, BroRecord *rec);
static void return_host_drop_check(BroConn *conn, BroRecord *rec);
static void return_scan(BroConn *conn, BroRecord *rec);
static void return_conn(BroConn *conn, BroRecord *rec);
static void return_host_drop(BroConn *conn, BroRecord *rec);

// ------------------------------------------------- //
int
main(int argc, char **argv)
	{
	int opt, port;
	BroConn *bc;
	extern char *optarg;
	extern int optind;
	char hostname[512];

	host_str = host_default;
	port_str = port_default;

	bro_debug_calltrace = 1;
	bro_debug_messages  = 1;

	while ( (opt = getopt(argc, argv, "p:h?")) != -1)
		{
		switch (opt)
			{
			case 'p':
				port_str = optarg;
				break;

			//default:
				//usage();
			}
		}

	argc -= optind;
	argv += optind;

	if (argc > 0)
		host_str = argv[0];

	port = strtol(port_str, NULL, 0);

	if (errno == ERANGE)
		{
		printf("Please provide a port number with -p.\n");
		exit(-1);
		}

	snprintf(hostname, 512, "%s:%s", host_str, port_str);

	/* Connect to Bro */
	if (! (bc = bro_connect_str(hostname, BRO_CFLAG_RECONNECT | BRO_CFLAG_ALWAYS_QUEUE)))
		{
		printf("Could not connect to Bro at %s:%s.\n", host_str, port_str);
		exit(-1);
		}

	// register the response events with the host we have connected to
	bro_event_registry_add(bc, "pong", (BroEventFunc) bro_pong_record);
	bro_event_registry_add(bc, "addr_checkR", (BroEventFunc) addr_check_record);
	bro_event_registry_add(bc, "return_cpu", (BroEventFunc) return_cpu);
	bro_event_registry_add(bc, "return_memory", (BroEventFunc) return_memory);
	bro_event_registry_add(bc, "return_host_drop_check", (BroEventFunc) return_host_drop_check);
	bro_event_registry_add(bc, "return_scan", (BroEventFunc) return_scan);
	bro_event_registry_add(bc, "return_conn", (BroEventFunc) return_conn);
	bro_event_registry_add(bc, "return_host_drop", (BroEventFunc) return_host_drop);

	bro_event_registry_request(bc);

	if (! bro_conn_await_handshake(bc, 10))
		{
		printf("Could not complete handshake successfully.\n");
		exit(-1);
		}
	printf("Handshake complete, starting shell...\n");
	printf(" deliniate arguments w/ tabs\n");

	// we are now connected to the bro instance, begin the command line fun

	//char** cp;
	int n, process;
	char *args[NARGS];
	char command[BUFSIZ];

	for (;;)
		{
		bro_conn_process_input(bc);
		printf("--> ");

		if ( fgets(command, sizeof(command), stdin) == NULL )
			{
			putchar('\n');
			exit(0);
			}

		// split up the command into words
		n = bufsplit(command, NARGS, args);

		// ignore blank lines
		if ( **args == '\0' )
			continue;
		//int t;
		//for ( t=0;t<=n;++t)
		//	printf("ARG(%d): |%s|\n", t,args[t]);

		// for now we assme that the command form is 
		// command <optional argument>
		// not all that flexable... 
		BroEvent *ev;
	  	BroRecord *rec = bro_record_new();
		bro_conn_process_input(bc);

		process = 0;

		if ( strcmp(args[0],PING) == 0 )
			{
			ev = bro_event_new("ping");
			int seq = 1;
	  		double timestamp = bro_util_current_time();

	  		bro_record_add_val(rec, "seq", BRO_TYPE_COUNT, &seq);
	  		bro_record_add_val(rec, "src_time", BRO_TYPE_TIME, &timestamp);
	  		bro_event_add_val(ev, BRO_TYPE_RECORD, rec);
	  		bro_record_free(rec);
			process=1;
			} 
		else if ( strcmp(args[0],AC) == 0 )
			{
			uint32 a = (uint32)inet_addr(args[1]);
			ev = bro_event_new("addr_check");

		  	bro_record_add_val(rec, "checkA", BRO_TYPE_IPADDR, &a);
		  	bro_event_add_val(ev, BRO_TYPE_RECORD, rec);
	  		bro_record_free(rec);
			process=1;
			}
		else if ( strcmp(args[0], STATS) == 0 )
			{
			ev = bro_event_new("show_cpu");
			process=1;
			}
		else if ( strcmp(args[0], MEM) == 0 )
			{
			ev = bro_event_new("show_memory");
			process=1;
			}
		else if ( strcmp(args[0], NU) == 0 )
			{
			ev = bro_event_new("null_event");
			process=1;
			}
		else if ( strcmp(args[0], ISHDROP) == 0 )
			{
			uint32 a = (uint32)inet_addr(args[1]);
			ev = bro_event_new("host_drop_check");

		  	bro_record_add_val(rec, "checkA", BRO_TYPE_IPADDR, &a);
		  	bro_event_add_val(ev, BRO_TYPE_RECORD, rec);
	  		bro_record_free(rec);
			process=1;
			}
		else if ( strcmp(args[0], DROP) == 0 )
			{
			uint32 a = (uint32)inet_addr(args[1]);
			ev = bro_event_new("host_drop");

		  	bro_record_add_val(rec, "checkA", BRO_TYPE_IPADDR, &a);
		  	bro_event_add_val(ev, BRO_TYPE_RECORD, rec);
	  		bro_record_free(rec);
			process=1;
			}
		else if ( strcmp(args[0], UNDROP) == 0 )
			{
			uint32 a = (uint32)inet_addr(args[1]);
			ev = bro_event_new("host_undrop");

		  	bro_record_add_val(rec, "checkA", BRO_TYPE_IPADDR, &a);
		  	bro_event_add_val(ev, BRO_TYPE_RECORD, rec);
	  		bro_record_free(rec);
			process=1;
			}

		else if ( strcmp(args[0], SCAN) == 0 )
			{
			// this looks at scan data assosciated with a 
			// single IP address
			uint32 a = (uint32)inet_addr(args[1]);
			ev = bro_event_new("show_scan");

		  	bro_record_add_val(rec, "checkA", BRO_TYPE_IPADDR, &a);
		  	bro_event_add_val(ev, BRO_TYPE_RECORD, rec);
	  		bro_record_free(rec);
			process=1;
			}
		else if ( strcmp(args[0], CONN) == 0 )
			{
			ev = bro_event_new("show_conn");
			process=1;
			}
		else if ( strcmp(args[0], ALLOWDROP) == 0 )
			{
			ev = bro_event_new("allow_drop");
			process=1;
			}


		if ( process != 0 )
			{
			bro_event_send(bc, ev);
			bro_event_free(ev);

			bro_conn_process_input(bc);
			}
		else
			{
			printf("bad command\n");
			}

		// this is in place to make sure that we wait around for
		// the external bro to (optionaly) return something
		sleep(1);
		
		} // end of shell loop
	bro_disconnect(bc);
  
	} // end of main



size_t bufsplit(char *buf, size_t n, char **a)
	{
	int i; 
	int nsplit;
	static char *splitch = "\t\n";

	if ( buf != NULL && n == 0 )
		{ 
		splitch = buf;
		return(1);
		}

	nsplit = 0;
	while ( nsplit < n )
		{
		a[nsplit++] = buf;

		if ((buf = strpbrk(buf, splitch)) == NULL)
			break;

		*(buf++) = '\0';

		if (*buf == '\0' )
			break;
		}

	buf = strrchr(a[nsplit-1], '\0');

	for (i=nsplit; i<n; i++)
		a[i] = buf;

	return(nsplit);
	}


static void
addr_check_record(BroConn *conn, BroRecord *rec)
{
 uint32 *ret;
      
  if (! bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &ret))
    {
      printf("Error getting return value\n");
      return;
    }
      
	if ( *ret == 1 )
		printf("address seen\n");
	else
		printf("address not seen\n");

 conn = NULL;
}

static void
bro_pong_record(BroConn *conn, BroRecord *rec)
{
  double now = bro_util_current_time();
  double *src_time, *dst_time;
  uint32 *seq;

  if (! bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &seq))
    {
      printf("Error getting sequence count from event.\n");
      return;
    }

  if (! bro_record_get_nth_val(rec, 1, BRO_TYPE_TIME, &src_time))
    {
      printf("Error getting src time from event.\n");
      return;
    }
  
  if (! bro_record_get_nth_val(rec, 2, BRO_TYPE_TIME, &dst_time))
    {
      printf("Error getting dst time from event.\n");
      return;
    }
  
  printf("pong event from %s: seq=%i, time=%f/%f s\n",
	 host_str, *seq, *dst_time - *src_time,
	 now - *src_time);
  
  conn = NULL;
}

static void return_cpu(BroConn *conn, double *interval)
	{
	// 
	printf("total CPU time/total clock time: %f\n", *interval);

	conn = NULL;
	}

static void return_memory(BroConn *conn, BroRecord *rec)
	{
	uint32 *mem, *minor_faults, *major_faults, *num_swap;
	uint32 *blocking_input, *blocking_output, *num_context; 

	bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &mem);
	bro_record_get_nth_val(rec, 1, BRO_TYPE_COUNT, &minor_faults);
	bro_record_get_nth_val(rec, 2, BRO_TYPE_COUNT, &major_faults);
	bro_record_get_nth_val(rec, 3, BRO_TYPE_COUNT, &num_swap);
	bro_record_get_nth_val(rec, 4, BRO_TYPE_COUNT, &blocking_input);
	bro_record_get_nth_val(rec, 5, BRO_TYPE_COUNT, &blocking_output);
	bro_record_get_nth_val(rec, 6, BRO_TYPE_COUNT, &num_context);

	printf("memory data: \n");
	printf("\ttotal mem (k):   %d\n", *mem);
	printf("\tminor faults:    %d\n", *minor_faults);
	printf("\tmajor faults:    %d\n", *major_faults);
	printf("\tnum swap:        %d\n", *num_swap);
	printf("\tblocking input:  %d\n", *blocking_input);
	printf("\tblocking output: %d\n", *blocking_output);
	printf("\tnum context:     %d\n", *num_context);

	conn = NULL;
	}

static void
return_host_drop_check(BroConn *conn, BroRecord *rec)
{
 uint32 *ret;
      
  if (! bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &ret))
    {
      printf("Error getting return value\n");
      return;
    }
      
	if ( *ret == 1 )
		printf("Address Dropped\n");
	else
		printf("Address not dropped\n");

 conn = NULL;
}

static void return_scan(BroConn *conn, BroRecord *rec)
	{
	uint32 *backscatter_peers, *distinct_peers, *poss_scan_source;
	//uint32 *num_distinct_ports;

	bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &backscatter_peers);
	bro_record_get_nth_val(rec, 1, BRO_TYPE_COUNT, &distinct_peers);
	bro_record_get_nth_val(rec, 2, BRO_TYPE_COUNT, &poss_scan_source);
	//bro_record_get_nth_val(rec, 3, BRO_TYPE_COUNT, &num_distinct_ports);

	printf("scan data:\n");
	printf("\thost count: %d\n", *distinct_peers);
	printf("\tbackscatter count: %d\n", *backscatter_peers);
	printf("\tpossible scan source: %d\n", *poss_scan_source);

	conn = NULL;
	}

static void return_conn(BroConn *conn, BroRecord *rec)
	{
	uint32 *num_TCP_conns, *num_UDP_conns, *num_ICMP_conns;
	uint32 *num_fragments, *max_TCP_conns, *max_UDP_conns;
	uint32 *max_ICMP_conns, *max_fragments;

	bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &num_TCP_conns);
	bro_record_get_nth_val(rec, 1, BRO_TYPE_COUNT, &num_UDP_conns);
	bro_record_get_nth_val(rec, 2, BRO_TYPE_COUNT, &num_ICMP_conns);
	bro_record_get_nth_val(rec, 3, BRO_TYPE_COUNT, &num_fragments);
	bro_record_get_nth_val(rec, 4, BRO_TYPE_COUNT, &max_TCP_conns);
	bro_record_get_nth_val(rec, 5, BRO_TYPE_COUNT, &max_UDP_conns);
	bro_record_get_nth_val(rec, 6, BRO_TYPE_COUNT, &max_ICMP_conns);
	bro_record_get_nth_val(rec, 7, BRO_TYPE_COUNT, &max_fragments);

	printf("connection data:\n");
	printf("\t      tcp (current/max): %d/%d \n", *num_TCP_conns, *max_TCP_conns);
	printf("\t      udp (current/max): %d/%d \n", *num_UDP_conns, *max_UDP_conns);
	printf("\t     icmp (current/max): %d/%d \n", *num_ICMP_conns, *max_ICMP_conns);
	printf("\tfragments (current/max): %d/%d \n", *num_fragments, *max_fragments);

	conn = NULL;
	}

static void return_host_drop(BroConn *conn, BroRecord *rec)
	{
	uint32 *ret;
      
	if (! bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &ret))
		{
		printf("Error getting return value\n");
		return;
		}

	switch (*ret)
		{
		case 0:
			printf("policy forbids host dropping\n");
			break;

		case 1:
			printf("host in no drop address space\n");
			break;

		case 2:
			printf("drop script run on host\n");
			break;

		case 3:
			printf("host has been already dropped\n");
			break;

		default:
			break;

		}

	conn = NULL;

	}