// Copyright (c) 1996, 1997, 1998, 1999, 2000, 2001, 2002, 2003
//      The Regents of the University of California.  All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that: (1) source code distributions
// retain the above copyright notice and this paragraph in its entirety, (2)
// distributions including binary code include the above copyright notice and
// this paragraph in its entirety in the documentation or other materials
// provided with the distribution, and (3) all advertising materials mentioning
// features or use of this software display the following acknowledgement:
// ``This product includes software developed by the University of California,
// Lawrence Berkeley Laboratory and its contributors.'' Neither the name of
// the University nor the names of its contributors may be used to endorse
// or promote products derived from this software without specific prior
// written permission.
// THIS SOFTWARE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR IMPLIED
// WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.



// secondary proof of concept shell for bro communications
// project begin 6/22/05


	
#include <string>
#include <vector>
#include <iostream>
#include <iomanip>
#include <errno.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <fcntl.h>

#include <sys/time.h>
#include <unistd.h>

#include <broccoli.h>
//#include "broccoli.h"

using std::string;
using std::vector;
using std::cout;
using std::cin;
using std::cerr;

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

// prototypes
static void bro_pong_record(BroConn *conn, BroRecord *rec);
static void return_line(BroConn *conn, BroRecord *rec);
static void return_set(BroConn *conn, BroRecord *rec);
static void prompt();
static void tokenize(const string& str, vector<string>& tokens);
static void print_showhelp();

int main(int argc, char **argv)
	{
	int opt, port, rc;
	BroConn *bc;
	extern char *optarg;
	extern int optind;
	char hostname[512];
	char buf[1024];
	int n;
	string command;

	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);
		}

	// define and 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, "return_line", (BroEventFunc) return_line);
	bro_event_registry_add(bc, "return_set", (BroEventFunc) return_set);
	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-->\n");

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

	while(1)		// main loop
		{
		vector<string> tokens;			// vector to hold parsed command line
		string cmd;				// first argument in command
		BroEvent *ev;				// new event
	  	BroRecord *rec = bro_record_new();	// record (optional) for some events

		bro_conn_process_input(bc);		// this should get procesed on every loop
		int process = 0;			// default unknown command
		prompt();

		if ( fgets(buf, sizeof(buf), stdin) == NULL )
			{
			putchar('\n');
			exit(0);
			}
		command = buf;	// turn buf into c++ string

		tokenize(command, tokens);		// tokenize the raw command
		cmd = tokens[0];			// initial 'key' token

		if ( tokens[0] == "show" )
			{
			if ( tokens[1] == "cpu" )
				{ // general cpu information
				ev = bro_event_new("show_cpu");
				process = 1;
				}
			else if ( tokens[1] == "memory" )
				{ // general memory information
				ev = bro_event_new("show_memory");
				process = 1;
				}
			else if ( tokens[1] == "conn" )
				{ // general connection information
				ev = bro_event_new("show_conn");
				process = 1;
				}
			else if ( tokens[1] == "dropped" )
				{ // has this IP been dropped?
				ev = bro_event_new("host_drop_check");
				uint32 a = (uint32)inet_addr(tokens[2].c_str());
				bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);
				process = 1;
				}
			else if ( tokens[1] == "scan" )
				{ // scan data about the IP
				uint32 a = (uint32)inet_addr(tokens[2].c_str());
				ev = bro_event_new("show_scan");
				bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);
				process = 1;
				}
			else if ( tokens[1] == "lookup" )
				{ // this returns data on an IP/pair if the policy has
				  // been activated.  all testing is done on the policy side.
				if ( tokens[2] == "pair" )
					{
					ev = bro_event_new("lookup_pair");
					uint32 a = (uint32)inet_addr(tokens[3].c_str());
					uint32 a2 = (uint32)inet_addr(tokens[4].c_str());
				  	bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);
				  	bro_event_add_val(ev, BRO_TYPE_IPADDR, &a2);
					}
				else
					{
					ev = bro_event_new("lookup");
					uint32 a = (uint32)inet_addr(tokens[2].c_str());
				  	bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);
					}
				process = 1;
				}
			else if ( tokens[1] == "pcap" )
				{ // print current pcap filter
				ev = bro_event_new("show_pcap_filter");
				process = 1;
				}
			else if ( tokens[1] == "fabric" )
				{ // print current fabric structure
				ev = bro_event_new("show_fabric");
				process = 1;
				}
			else if ( tokens[1] == "notice" )
                                { // print given notice data
                                ev = bro_event_new("show_notice");

                                BroString arg;
                                bro_string_init(&arg);
                                bro_string_set(&arg,tokens[2].c_str());
                                bro_event_add_val(ev, BRO_TYPE_STRING, &arg);
                                bro_string_cleanup(&arg);

                                process = 1;
                                }
			else if ( tokens[1] == "help" || "?" )
				print_showhelp();
			}

		else if ( tokens[0] == "set" )
			{ // this is the scan related functionality
			if ( tokens[1] == "scan" )
				{
				if ( tokens[2] == "on" )
					{
					ev = bro_event_new("scan_on");
					process = 1;
					}
				else if ( tokens[2] == "off" )
					{
					ev = bro_event_new("scan_off");
					process = 1;
					}
				else if ( tokens[2] == "clear" )
					{
					ev = bro_event_new("reset_scan");
					uint32 a = (uint32)inet_addr(tokens[2].c_str());
				  	bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);
					process = 1;
					}
				else
					{
					printf("unknown set scan option\n");
					printf("\tuse: scan {on|off|clear <IP>}\n");
					}
				}
			else if ( tokens[1] == "lookup" )
				{
				if ( tokens[2] == "on" )
					{
					ev = bro_event_new("lookup_on");
					process = 1;
					}
				else if ( tokens[2] == "off" )
					{
					ev = bro_event_new("lookup_off");
					process = 1;
					}
				}
			else if ( tokens[1] == "drop" )
				{
				if ( tokens[2] == "on" )
					{
					ev = bro_event_new("drop_on");
					process = 1;
					}
				else if ( tokens[2] == "off" )
					{
					ev = bro_event_new("drop_off");
					process = 1;
					}
				}
			else if ( tokens[1] == "pcap" )
				{
				// this may need a little work...
				ev = bro_event_new("set_pcap");

				BroString arg;
				bro_string_init(&arg);
				bro_string_set(&arg,tokens[2].c_str());
				bro_event_add_val(ev, BRO_TYPE_STRING, &arg);
				bro_string_cleanup(&arg);

				process = 1;
				}
			else if ( tokens[1] == "help" || "?" )
				//print_sethelp();
				printf("SETHELP\n");
			}

		else if ( tokens[0] == "ping" )
			{
			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 ( tokens[0] == "drop" )
			{
			ev = bro_event_new("host_drop");
			uint32 a = (uint32)inet_addr(tokens[1].c_str());
		  	bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);

			process = 1;
			}
		else if ( tokens[0] == "undrop" )
			{
			ev = bro_event_new("host_undrop");
			uint32 a = (uint32)inet_addr(tokens[1].c_str());
		  	bro_event_add_val(ev, BRO_TYPE_IPADDR, &a);

			process = 1;
			}

		if ( process == 1 )
			{
			bro_event_send(bc, ev);
			bro_event_free(ev);
			bro_conn_process_input(bc);
			process = 0;
			}

		sleep(1);
		}

} // end main

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_line(BroConn *conn, BroRecord *rec)
	{
	BroString *string;
	bro_record_get_nth_val(rec, 0, BRO_TYPE_STRING, &string);

	printf("%s\n", string->str_val);

	conn = NULL;
	}

static void return_set(BroConn *conn, BroRecord *rec)
	{
	BroString *string;
	int iter,i;
printf("in return set\n");
	bro_record_get_nth_val(rec, 0, BRO_TYPE_COUNT, &iter);
printf("in return set iter=%d\n", iter);

	for ( i = 0; i<iter; i++)
		{
		bro_record_get_nth_val(rec, (i+1), BRO_TYPE_STRING, &string);
		printf("%s\n", string->str_val);
		}

	}

void prompt()
	{
	printf("--> ");
	}

void tokenize(const string& str, vector<string>& tokens)
	{
	int num_tokens = 0;
	char delim = '\0';
 
	for ( unsigned int i = 0; i < str.length(); ++i )
		{
		while ( isspace(str[i]) )
		++i;
 
		string next_arg;
 
		if (str[i] == '"' || str[i] == '\'')
			{
			delim = str[i];
			++i;
			}
		else
		delim = '\0';
 
		for ( ; str[i]; ++i )
			{
			if ( delim && str[i] == '\\' &&
				 i < str.length() && str[i+1] == delim )
				{
				++i;
				next_arg.push_back(str[i]);
				}
	 
			else if ( delim && str[i] == delim )
				{
				++i;
				break;
				}
 
			else if ( ! delim && isspace(str[i]) )
				break;
			else
				next_arg.push_back(str[i]);
			}
 
		tokens.push_back(next_arg);
		}
	}

void print_showhelp()
	{
	printf("show\n"); 
	printf("\tcpu\n");
	printf("\tmemory|mem\n");
	printf("\tconn\n");
	printf("\tdroped [ip]\n");
	printf("\tscan\n");
	printf("\tlookup [ip]\n");
	printf("\tlookup [ip] [ip]\n");
	printf("\tpcap\n");
	printf("\n");
	}