/*
 *  nfreplay :  Reads netflow data from files saved by nfcapd
 *  			and sends them to another host.
 *
 *  Copyright (c) 2004, SWITCH - Teleinformatikdienste fuer Lehre und Forschung
 *  All rights reserved.
 *  
 *  Redistribution and use in source and binary forms, with or without 
 *  modification, are permitted provided that the following conditions are met:
 *  
 *   * Redistributions of source code must retain the above copyright notice, 
 *	 this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright notice, 
 *	 this list of conditions and the following disclaimer in the documentation 
 *	 and/or other materials provided with the distribution.
 *   * Neither the name of SWITCH 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 BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
 *  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 
 *  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
 *  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
 *  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
 *  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
 *  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
 *  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
 *  POSSIBILITY OF SUCH DAMAGE.
 *  
 *  $Author: peter $
 *
 *  $Id: nftrack.c 67 2007-03-06 10:16:19Z peter $
 *
 *  $LastChangedRevision: 67 $
 *	
 */

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

#include "config.h"

#ifdef HAVE_STDINT_H
#include <stdint.h>
#endif


#include "version.h"
#include "nf_common.h"
#include "nffile.h"
#include "rbtree.h"
#include "nftree.h"
#include "nfdump.h"
#include "util.h"
#include "grammar.h"

#include "nftrack_stat.h"
#include "nftrack_rrd.h"

// We have 288 slot ( 1 day ) for stat record
#define NUM_AV_SLOTS 288

/* Externals */
extern int yydebug;

/* Global Variables */
FilterEngine_data_t	*Engine;
int 		byte_mode, packet_mode;
uint32_t	byte_limit, packet_limit;	// needed for linking purpose only

/* Local Variables */
static char const *rcsid 		  = "$Id: nftrack.c 67 2007-03-06 10:16:19Z peter $";

/* Function Prototypes */
static void usage(char *name);

static data_row *process(char *filter);

/* Functions */
static void usage(char *name) {
		printf("usage %s [options] [\"filter\"]\n"
					"-h\t\tthis text you see right here\n"
					"-l\t\tLast update of Ports DB\n"
					"-V\t\tPrint version and exit.\n"
					"-I\t\tInitialize Ports DB files.\n"
					"-d <db_dir>\tPorts DB directory.\n"
					"-r <input>\tread from file. default: stdin\n"
					"-p\t\tOnline output mode.\n"
					"-s\t\tCreate port statistics for timeslot -t\n"
					"-t <time>\tTimeslot for statistics\n"
					"-S\t\tCreate port statistics for last day\n"
					"-w <file>\twrite output to file\n"
					"-f <filter>\tfilter syntaxfile\n"
					, name);
} /* usage */

/* 
 * Some C code is needed for daemon code as well as normal stdio code 
 * therefore a generic LogError is defined, which maps in this case
 * to syslog
 */
void LogError(char *format, ...) {
va_list var_args;
char string[512];

    va_start(var_args, format);
    vsnprintf(string, 511, format, var_args);
    va_end(var_args);
    syslog(LOG_ERR, "%s", string);

} // End of LogError

static data_row *process(char *filter) {
data_block_header_t flow_header;					
master_record_t		master_record;
common_record_t		*flow_record, *in_buff;
int i, rfd, done, ret;
uint32_t	NumRecords, buffer_size;
data_row * 	port_table;

	rfd = GetNextFile(0, 0, 0, NULL);
	if ( rfd < 0 ) {
		if ( errno ) 
			perror("Can't open file for reading");
		return NULL;
	}

	// prepare read and send buffer
	buffer_size = BUFFSIZE;
	in_buff = (common_record_t *) malloc(buffer_size);
	if ( !in_buff ) {
		perror("Memory allocation error");
		close(rfd);
		return;
	}

	port_table    = (data_row *)calloc(65536, sizeof(data_row));
    if ( !port_table) {
        perror("Memory allocation error");
        close(rfd);
        return;
    }

    memset((void *)port_table, 0, 65536 * sizeof(data_row));

	// setup Filter Engine to point to master_record, as any record read from file
	// is expanded into this record
	Engine->nfrecord = (uint64_t *)&master_record;

	done	 	= 0;
	while ( !done ) {
		ret = read(rfd, &flow_header, sizeof(data_block_header_t));
		if ( ret == 0 ) {
            // EOF of rfd
            rfd = GetNextFile(rfd, 0, 0, NULL);
            if ( rfd < 0 ) {
                if ( rfd == FILE_ERROR )
                    fprintf(stderr, "Can't read from file '%s': %s\n",GetCurrentFilename(), strerror(errno) );
                done = 1;
            } 
			continue;
		} else if ( ret == -1 ) {
			perror("Error reading data");
			close(rfd);
			return;
		}

		if ( flow_header.id != DATA_BLOCK_TYPE_1 ) {
			fprintf(stderr, "Can't process block type %u\n", flow_header.id);
			continue;
		}

		NumRecords = flow_header.NumBlocks;
		if ( flow_header.size > buffer_size ) {
			void *tmp;
			// Actually, this should never happen, but catch it anyway
			if ( flow_header.size > MAX_BUFFER_SIZE ) {
				// this is most likely corrupt
				fprintf(stderr, "Corrupt data file: Requested buffer size %u exceeds max. buffer size.\n", flow_header.size);
				break;
			}
			// make it at least the requested size
			buffer_size = flow_header.size;
			tmp = realloc((void *)in_buff, buffer_size);
			if ( !tmp ) {
				fprintf(stderr, "Can't reallocate buffer to %u bytes: %s\n", buffer_size, strerror(errno));
				break;
			}
			in_buff = (common_record_t *)tmp;
		}
		ret = read(rfd, in_buff, flow_header.size);
		if ( ret == 0 ) {
			done = 1;
			break;
		} else if ( ret == -1 ) {
			perror("Error reading data");
			close(rfd);
			return;
		}

		if ( flow_header.size != ret ) {
			// Ups - this was a short read - most likely reading from the stdin pipe
			// loop until we have requested size
			size_t	request_size, total_size;
			void *read_ptr;

			total_size 	 = ret;
			request_size = flow_header.size - total_size;
			read_ptr 	 = (void *)((pointer_addr_t)in_buff + total_size);
			do {
				ret = read(rfd, read_ptr, request_size);
				if ( ret == 0 ) {
					break;
				} else if ( ret < 0 ) {
					perror("Error reading data");
					break;
				}
				total_size += ret;
				if ( total_size < flow_header.size ) {
					request_size = flow_header.size - ret;
					read_ptr 	 = (void *)((pointer_addr_t)in_buff + total_size);
					request_size = flow_header.size - total_size;
				}
			} while ( ret > 0 && ( total_size < flow_header.size ));
			
			if ( total_size != flow_header.size ) {
				// still unsuccessful
				fprintf(stderr, "Short read for netflow records: Expected %i, got %lu bytes!\n",flow_header.size, total_size );
				break;
			} else {
				// continue
				ret = flow_header.size;
			}
		}

		flow_record = in_buff;

		for ( i=0; i < NumRecords; i++ ) {
			// if no filter is given, the result is always true
			ExpandRecord( flow_record, &master_record);

			// filter netflow record with user supplied filter
			flow_record->mark = (*Engine->FilterEngine)(Engine);

			if ( flow_record->mark == 0 ) { // record failed to pass the filter
				// increment pointer by number of bytes for netflow record
				flow_record = (common_record_t *)((pointer_addr_t)flow_record + flow_record->size);	
				// go to next record
				continue;
			}

			// Add to stat record
			if ( master_record.prot == 6 ) {
				port_table[master_record.dstport].proto[tcp].type[flows]++;
				port_table[master_record.dstport].proto[tcp].type[packets]	+= master_record.dPkts;
				port_table[master_record.dstport].proto[tcp].type[bytes]	+= master_record.dOctets;
			} else if ( master_record.prot == 17 ) {
				port_table[master_record.dstport].proto[udp].type[flows]++;
				port_table[master_record.dstport].proto[udp].type[packets]	+= master_record.dPkts;
				port_table[master_record.dstport].proto[udp].type[bytes]	+= master_record.dOctets;
			}
			/* else
				other protocols are not handled
			 */

			// Advance pointer by number of bytes for netflow record
			flow_record = (common_record_t *)((pointer_addr_t)flow_record + flow_record->size);	
		}
	} // while

	return port_table;

} // End of process


int main( int argc, char **argv ) {
struct stat stat_buff;
char c, *wfile, *rfile, *Rfile, *Mdirs, *ffile, *filter, *timeslot, *DBdir;
char datestr[64];
int ffd, ret, DBinit, AddDB, GenStat, AvStat, output_mode;
unsigned int lastupdate, topN;
data_row *port_table;
time_t	when;
struct tm * t1;

	wfile = rfile = Rfile = Mdirs = ffile = filter = DBdir = timeslot = NULL;
	DBinit = AddDB = GenStat = AvStat = 0;
	lastupdate = output_mode = 0;
	topN = 10;
	while ((c = getopt(argc, argv, "d:hln:pr:st:w:AIM:R:SV")) != EOF) {
		switch (c) {
			case 'h':
				usage(argv[0]);
				exit(0);
				break;
			case 'I':
				DBinit = 1;
				break;
			case 'M':
				Mdirs = strdup(optarg);
				break;
			case 'R':
				Rfile = strdup(optarg);
				break;
			case 'd':
				DBdir = strdup(optarg);
				ret  = stat(DBdir, &stat_buff);
				if ( !(stat_buff.st_mode & S_IFDIR) ) {
					fprintf(stderr, "No such directory: %s\n", DBdir);
					exit(255);
				}
				break;
			case 'l':
				lastupdate = 1;
				break;
			case 'n':
				topN = atoi(optarg);
				if ( topN < 0 ) {
					fprintf(stderr, "TopnN number %i out of range\n", topN);
					exit(255);
				}
				break;
			case 'p':
				output_mode = 1;
				break;
			case 'r':
				rfile = strdup(optarg);
				break;
			case 'w':
				wfile = strdup(optarg);
				break;
			case 's':
				GenStat = 1;
				break;
			case 't':
				timeslot = optarg;
				if ( !ISO2UNIX(timeslot) ) {
					exit(255);
				}
				break;
			case 'A':
				AddDB = 1;
				break;
			case 'S':
				AvStat = 1;
				break;
			default:
				usage(argv[0]);
				exit(0);
		}
	}

	if (argc - optind > 1) {
		usage(argv[0]);
		exit(255);
	} else {
		/* user specified a pcap filter */
		filter = argv[optind];
	}

	openlog(argv[0] , LOG_CONS|LOG_PID, LOG_DAEMON);

	if ( !filter && ffile ) {
		if ( stat(ffile, &stat_buff) ) {
			perror("Can't stat file");
			exit(255);
		}
		filter = (char *)malloc(stat_buff.st_size);
		if ( !filter ) {
			perror("Memory error");
			exit(255);
		}
		ffd = open(ffile, O_RDONLY);
		if ( ffd < 0 ) {
			perror("Can't open file");
			exit(255);
		}
		ret = read(ffd, (void *)filter, stat_buff.st_size);
		if ( ret < 0   ) {
			perror("Error reading file");
			close(ffd);
			exit(255);
		}
		close(ffd);
	}

	if ( !DBdir ) {
		fprintf(stderr, "DB directory required\n");
		exit(255);
	}

	InitStat(DBdir);

	if ( !filter )
		filter = "any";

	Engine = CompileFilter(filter);
	if ( !Engine ) 
		exit(254);
	
	if ( DBinit ) {
		when = time(NULL);
		when -= ((when % 300) + 300);
		InitStatFile(when, NUM_AV_SLOTS);
		if ( !CreateRRDBs(DBdir, when) ) {
			fprintf(stderr, "Init DBs failed\n");
			exit(255);
		}
		fprintf(stderr, "Port DBs initialized.\n");
		exit(0);
	}

	if ( lastupdate ) {
		when = RRD_LastUpdate(DBdir);
		if ( !when )
			exit(255);
		t1 = localtime(&when);
		strftime(datestr, 63, "%b %d %Y %T", t1);
		printf("Last Update: %i, %s\n", when, datestr);
		exit(0);
	}

	port_table = NULL;
	if ( Mdirs || Rfile || rfile ) {
		SetupInputFileSequence(Mdirs, rfile, Rfile);
		port_table = process(filter);
//		Lister(port_table);
		if ( !port_table ) {
			exit(255);
		}
		if ( AddDB ) {
			if ( !timeslot ) {
				fprintf(stderr, "Timeslot required!\n");
				exit(255);
			}
			UpdateStat(port_table, ISO2UNIX(timeslot));
			RRD_StoreDataRow(DBdir, timeslot, port_table);
		}
	}

	if ( AvStat ) {
		port_table = GetStat();
		if ( !port_table ) {
			fprintf(stderr, "Unable to get port table!\n");
			exit(255);
		}
		// DoStat
		Generate_TopN(port_table, topN, NUM_AV_SLOTS, 0, output_mode, wfile);

	} 


	if ( GenStat ) {
		when = ISO2UNIX(timeslot);
		if ( !port_table ) {
			if ( !timeslot ) {
				fprintf(stderr, "Timeslot required!\n");
				exit(255);
			}
			port_table = RRD_GetDataRow(DBdir, when);
		}
		if ( !port_table ) {
			fprintf(stderr, "Unable to get port table!\n");
			exit(255);
		}
		// DoStat
		Generate_TopN(port_table, topN, 1, when, output_mode, wfile);
	}

	CloseStat();

	return 0;
}
