// $Id: main.cc,v 1.80 2003/09/01 23:42:29 vern Exp $
//
// Copyright (c) 1995, 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.

#include "config.h"

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>

#include "input.h"
#include "Active.h"
#include "DNS_Mgr.h"
#include "Frame.h"
#include "Scope.h"
#include "Event.h"
#include "File.h"
#include "Logger.h"
#include "Net.h"
#include "NetVar.h"
#include "Var.h"
#include "Timer.h"
#include "Stmt.h"
#include "Debug.h"
#include "DFA.h"
#include "RuleMatcher.h"
#include "Anon.h"
#include "Serializer.h"
#include "EventRegistry.h"

#ifndef HAVE_STRSEP
extern "C" {
char *strsep(char **, const char *);
};
#endif

extern "C" {
#include "setsignal.h"
};

#ifdef USE_MPATROL
#include <mpatrol.h>
#include <mpatrol/heapdiff.h>
#endif

const char* prog;
name_list prefixes;
DNS_Mgr* dns_mgr;
TimerMgr* timer_mgr;
Logger* bro_logger;
Func* log_hook = 0;
Stmt* stmts;
Func* bro_signal = 0;
RuleMatcher* rule_matcher = 0;
PersistenceSerializer* persistence_serializer = 0;
FileSerializer* event_serializer = 0;
EventPlayer* event_player = 0;
EventRegistry* event_registry = 0;
event_source_list external_event_sources;
int signal_val = 0;
int optimize = 0;
int rule_bench = 0;
extern char version[];

void usage()
	{
#ifndef USE_MPATROL
	fprintf(stderr, "bro version %s\n", version);
#else
	fprintf(stderr, "bro/mpatrol version %s\n", version);
#endif

	fprintf(stderr, "usage: %s [options] [file ...]\n", prog);
	fprintf(stderr, "    <file>             | policy file, or read stdin\n");
#ifdef ACTIVE_MAPPING
	fprintf(stderr, "    -a <mapfile>       | use active mapping results\n");
#endif
	fprintf(stderr, "    -d                 | activate policy file debugging\n");
	fprintf(stderr, "    -f <filter>        | tcpdump filter\n");
	fprintf(stderr, "    -h                 | command line help\n");
	fprintf(stderr, "    -i <interface>     | read from given interface\n");
	fprintf(stderr, "    -p <prefix>        | add given prefix to policy file resolution\n");
	fprintf(stderr, "    -r <readfile>      | read from given tcpdump file\n");
	fprintf(stderr, "    -s <rulefile>      | read rules from given file\n");
	fprintf(stderr, "    -t <tracefile>     | activate execution tracing\n");
	fprintf(stderr, "    -w <writefile>     | write to given tcpdump file\n");
	fprintf(stderr, "    -v                 | print version and exit\n");
	fprintf(stderr, "    -A <writefile>     | write transformed trace to given tcpdump file\n");
	fprintf(stderr, "    -C                 | ignore checksums\n");
	fprintf(stderr, "    -D <size>          | DFA state cache size\n");
	fprintf(stderr, "    -F                 | force DNS\n");
	fprintf(stderr, "    -L                 | benchmark for rules\n");
	fprintf(stderr, "    -O                 | optimize policy script\n");
	fprintf(stderr, "    -P                 | prime DNS\n");
	fprintf(stderr, "    -R <events.bst>    | replay events\n");
	fprintf(stderr, "    -S                 | enable rule debugging\n");
	fprintf(stderr, "    -T <level>         | set 'RE_level' for rules\n");
#if 0
	fprintf(stderr, "    -V                 | log also to screen\n");
#endif
	fprintf(stderr, "    -W                 | activate watchdog timer\n");

#ifdef USE_MPATROL
	fprintf(stderr, "    -m                 | show leak table [mpatrol]\n");
	fprintf(stderr, "    -M                 | show unfreed heap diffs [mpatrol]\n");
#endif
	fprintf(stderr, "    -x <file.bst>      | print contents of state file\n");
	fprintf(stderr, "    -X <file.bst>      | print contents of state file as XML\n");

	fprintf(stderr, "    $BROPATH           | file search path (%s)\n", bro_path());
	fprintf(stderr, "    $BRO_PREFIXES      | prefix list (%s)\n", bro_prefixes());

	exit(1);
	}

void done_with_network()
	{
	// Cancel any pending alarms (watchdog, in particular).
	(void) alarm(0);

	EventHandlerPtr net_done = internal_handler("net_done");
	if ( net_done )
		{
		val_list* args = new val_list;
		args->append(new Val(timer_mgr->Time(), TYPE_TIME));
		mgr.QueueEvent(net_done, args);
		mgr.Drain();
		}

	timer_mgr->Expire();
	mgr.Drain();

	net_finish(1);
	}

void terminate_bro()
	{
	EventHandlerPtr bro_done = internal_handler("bro_done");
	if ( bro_done )
		mgr.QueueEvent(bro_done, new val_list);

	// Flush state once before draining in case something goes wrong ...
	const char* f = (const char*) state_file->AsString()->CheckString();

	persistence_serializer->WriteAll(f);

	timer_mgr->Expire();
	mgr.Drain();

	// Now the final flush.
	persistence_serializer->WriteAll(f);

	delete timer_mgr;
	delete bro_logger;
	delete dns_mgr;
	delete persistence_serializer;
	delete event_player;
	delete event_serializer;
	delete event_registry;
	}

void termination_signal()
	{
	Val sval(signal_val, TYPE_COUNT);
	message("received termination signal");
	net_get_final_stats();
	done_with_network();
	terminate_bro();
	net_delete();
	delete rule_matcher;
	exit(0);
	}

RETSIGTYPE sig_handler(int signo)
	{
	if ( signal_val )
		internal_error("double signal");

	if ( (signo == SIGTERM || signo == SIGINT) && processing_start_time == 0.0 )
		// We received the termination signal and we're right
		// now waiting for more packets - we're not in the middle
		// of processing a packet or events, so there are no
		// race conditions to worry about with handling the signal
		// right now (rather than waiting perhaps a long time for
		// pcap to give us the next buffer of packets.
		termination_signal();

	signal_val = signo;
	return RETSIGVAL;
	}

#ifdef USE_MPATROL

RETSIGTYPE sig_mpatrol(int signo)
	{
	static heapdiff hd;
	static bool hd_init = false;
	static bool leak_init = false;

	__mp_summary();

	switch ( signo )
		{
		case SIGUSR1:
			// Show leaks since last SIGUSR1.
			if ( leak_init )
				{
				__mp_stopleaktable();
				__mp_leaktable(50, MP_LT_UNFREED, MP_LT_COUNTS);
				__mp_clearleaktable();
				}

			__mp_startleaktable();

			leak_init = true;
			break;

		case SIGUSR2:
			// Show heap differences since last SIGUSR2.
			if ( hd_init )
				heapdiffend(hd);
			heapdiffstart(hd, HD_UNFREED | HD_VIEW | HD_FULL);
			hd_init = true;
			break;

		case SIGCONT:
			// Stop mpatrol logging.
			
			if ( leak_init )
				{
				__mp_stopleaktable();
				__mp_leaktable(50, MP_LT_UNFREED, MP_LT_COUNTS);
				__mp_clearleaktable();
				}

			if ( hd_init )
				heapdiffend(hd);
			break;

		default:
			internal_error("Unexpected signal in sig_mpatrol()");
		}

	return RETSIGVAL;
	}

#endif // USE_MPATROL


static void bro_new_handler()
	{
	out_of_memory("new");
	}

int main(int argc, char** argv)
	{
	name_list interfaces;
	name_list read_files;
	name_list rule_files;
	char* writefile = 0;
	char* transformed_writefile = 0;
	char* bst_file = 0;
	char* events_file = 0;
	int to_xml = 0;
	int do_watchdog = 0;
	int override_ignore_checksums = 0;
	int rule_debug = 0;
	int RE_level = 4;

#ifdef USE_MPATROL
	int leaktable = 0;
	int heapdiff = 0;
#endif

	enum DNS_MgrMode dns_type =
		getenv("BRO_DNS_FAKE") ? DNS_FAKE : DNS_DEFAULT;
	RETSIGTYPE (*oldhandler)(int);

	prog = argv[0];

	prefixes.append("");	// "" = "no prefix"

	char* p = getenv("BRO_PREFIXES");
	if ( p )
		add_to_name_list(p, ':', prefixes);

	string active_file;
	extern char* optarg;
	extern int optind, opterr;
	opterr = 0;

	int op;

#ifndef USE_MPATROL
	while ( (op = getopt(argc, argv, "A:a:D:f:i:p:r:w:b:l:R:s:T:t:x:X:CFLOPSWdhv")) != EOF )
#else
	while ( (op = getopt(argc, argv, "A:a:D:f:i:p:r:w:b:l:R:s:T:t:x:X:CFLOPSWdhvmM")) != EOF )
#endif
		switch ( op ) {
		case 'a':
#ifdef ACTIVE_MAPPING
			fprintf(stderr, "Using active mapping file %s.\n", optarg);
			active_file = optarg;
#else
			fprintf(stderr, "Bro not compiled for active mapping.\n");
			exit(1);
#endif
			break;

		case 'd':
			fprintf(stderr, "Policy file debugging ON.\n");
			g_policy_debug = true;
			break;

		case 'f':
			user_pcap_filter = optarg;
			break;

		case 'i':
			interfaces.append(optarg);
			break;

		case 'K':
#ifdef DEBUG_BRO
			hash_md5(strlen(optarg), (const u_char*) optarg,
				 shared_hmac_md5_key);
			hmac_key_set = 1;
#else
			fprintf(stderr, "-K is valid only when bro is compiled with DEBUG_BRO\n");
#endif
			break;

		case 'p':
			prefixes.append(optarg);
			break;

		case 'r':
			read_files.append(optarg);
			break;

		case 's':
			rule_files.append(optarg);
			break;

		case 't':
			g_trace_state.SetTraceFile(optarg);
			g_trace_state.TraceOn();
			break;

		case 'w':
			writefile = optarg;
			break;

		case 'A':
			transformed_writefile = optarg;
			break;

		case 'C':
			override_ignore_checksums = 1;
			break;

		case 'D':
			dfa_state_cache_size = atoi(optarg);
			break;

		case 'F':
			if ( dns_type != DNS_DEFAULT )
				usage();
			dns_type = DNS_FORCE;
			break;

		case 'L':
			++rule_bench;
			break;

		case 'O':
			optimize = 1;
			break;

		case 'P':
			if ( dns_type != DNS_DEFAULT )
				usage();
			dns_type = DNS_PRIME;
			break;

		case 'R':
			events_file = optarg;
			break;

		case 'S':
			rule_debug = 1;
			break;

		case 'T':
			RE_level = atoi(optarg);
			break;

		case 'W':
			do_watchdog = 1;
			break;

		case 'h':
			usage();
			break;

		case 'v':
			fprintf(stderr, "%s version %s\n", prog, version);
			exit(0);
			break;

#ifdef USE_MPATROL
		case 'm':
			leaktable = 1;
			break;

		case 'M':
			heapdiff = 1;
			break;
#endif

		case 'x':
			bst_file = optarg;
			break;

		case 'X':
			bst_file = optarg;
			to_xml = 1;
			break;


		default:
			usage();
			break;
		}

	init_random_seed();
	DEBUG_MSG("HMAC key: %s\n", md5_digest_print(shared_hmac_md5_key));

	if ( interfaces.length() > 0 && read_files.length() > 0 )
		usage();

	timer_mgr = new CQ_TimerMgr();

	add_input_file("bro.init");

	if ( optind == argc && read_files.length() == 0 && ! bst_file )
		add_input_file("-");

	while ( optind < argc )
		add_input_file(argv[optind++]);

	if ( ! load_mapping_table(active_file.c_str()) )
		{
		fprintf(stderr, "Could not load active mapping file %s\n",
			active_file.c_str());
		exit(1);
		}

	dns_mgr = new DNS_Mgr(dns_type);

	persistence_serializer = new PersistenceSerializer;
	event_registry = new EventRegistry;

	if ( events_file )
		{
		event_player = new EventPlayer();
		if ( event_player->Replay(events_file) )
			external_event_sources.append(event_player);
		}

	push_scope(0);
	yyparse();

	if ( nerr > 0 )
		{
		delete dns_mgr;
		exit(1);
		}

	// Parse rule files defined on the script level.
	char* script_rule_files =
		copy_string(internal_val("signature_files")->AsString()->CheckString());

	char* tmp = script_rule_files;
	char* s;
	while ( (s = strsep(&tmp, " \t")) )
		if ( *s )
			rule_files.append(s);

	if ( rule_files.length() > 0 )
		{
		rule_matcher = new RuleMatcher(RE_level);
		if ( ! rule_matcher->ReadFiles(rule_files) )
			{
			delete dns_mgr;
			exit(1);
			}

		if ( rule_debug )
			rule_matcher->PrintDebug();
		}

	delete [] script_rule_files;

	if ( g_policy_debug )
		// ### Add support for debug command file.
		dbg_init_debugger(0);

	dns_mgr->Init();

	Val* bro_log_file = internal_val("bro_log_file");
	bro_logger = new Logger("bro",
		bro_log_file ? bro_log_file->AsFile() : new BroFile(stderr));

	if ( read_files.length() == 0 && interfaces.length() == 0 )
		{
		Val* interfaces_val = internal_val("interfaces");
		if ( interfaces_val )
			{
			char* interfaces_str =
				interfaces_val->AsString()->ExpandedString();

			if ( interfaces_str[0] != '\0' )
				add_to_name_list(interfaces_str, ' ', interfaces);

			delete [] interfaces_str;
			}
		}

	if ( writefile && interfaces.length() == 0 && read_files.length() == 0 )
		usage();

	// Only enable actual syslog'ing for live network.
	bro_logger->SetEnabled((interfaces.length() > 0));

	if ( dns_type != DNS_PRIME )
		net_init(interfaces, read_files, writefile, transformed_writefile,
				 user_pcap_filter ? user_pcap_filter : "tcp or udp", do_watchdog);

	log_hook = internal_func("log_hook");
	bro_signal = internal_func("bro_signal");

	if ( ! g_policy_debug )
		{
		(void) setsignal(SIGTERM, sig_handler);
		(void) setsignal(SIGINT, sig_handler);
#ifdef USE_MPATROL
		(void) setsignal(SIGUSR1, sig_mpatrol);
		(void) setsignal(SIGUSR2, sig_mpatrol);
		(void) setsignal(SIGCONT, sig_mpatrol);
#endif
		}

	// Cooperate with nohup(1).
	if ( (oldhandler = setsignal(SIGHUP, sig_handler)) != SIG_DFL )
		(void) setsignal(SIGHUP, oldhandler);

#ifndef USE_MPATROL
	set_new_handler(bro_new_handler);
#endif

	no_handler = internal_func("no_handler");
	if ( ! no_handler )
		no_handler = internal_func("no_handler_default");

	if ( dns_type == DNS_PRIME )
		{
		dns_mgr->Verify();
		dns_mgr->Resolve();

		if ( ! dns_mgr->Save() )
			bro_logger->Log("**Can't update DNS cache");

		mgr.Drain();
		delete dns_mgr;
		exit(0);
		}

	// Just read state file from disk.
	if ( bst_file )
		{
		if ( to_xml )
			{
			BinarySerializationFormat* b =
				new BinarySerializationFormat();
			XMLSerializationFormat* x = new XMLSerializationFormat();
			ConversionSerializer s(b, x);
			s.Convert(bst_file, "/dev/stdout");
			}
		else
			{
			FileSerializer s;
			s.Read(bst_file);
			}

		exit(0);
		}

	const char* f = (const char*) state_file->AsString()->Bytes();
	persistence_serializer->ReadAll(f);

	EventHandlerPtr bro_init = internal_handler("bro_init");
	if ( bro_init )	//### this should be a function
		mgr.QueueEvent(bro_init, new val_list);

	if ( stmts )
		{
		stmt_flow_type flow;
		Frame f(current_scope()->Length(), 0, 0);
		g_frame_stack.push_back(&f);
		stmts->Exec(&f, flow);
		g_frame_stack.pop_back();
		}

	if ( override_ignore_checksums )
		ignore_checksums = 1;

	mgr.Drain();

	if ( interfaces.length() > 0 || read_files.length() > 0 || event_player )
		{
#ifdef USE_MPATROL
		if ( leaktable )
			{
			fputs( "Starting leak table...\n", stderr );
			sig_mpatrol(SIGUSR1);
			}

		if ( heapdiff )
			{
			fputs( "Starting heap diffs...\n", stderr );
			sig_mpatrol(SIGUSR2);
			}

#endif

		net_run();
		done_with_network();
		terminate_bro();
		net_delete();

#ifdef USE_MPATROL
		fputs( "Stopping mpatrol logging...", stderr );
		sig_mpatrol(SIGCONT);
#endif

		}
	else
		terminate_bro();

	delete rule_matcher;

	return 0;
	}
