/**
 * This program retrieves all of the I/O that a given pid 
 * (and its children) wrote to stdout, stderr and stdin.
 * 
 * Usage: get_io pid rollcount
 * 
 * Currently, this should handle everything *except* for a few cases:
 * 1. Right now, the close_on_exec flag for fd's is checked when FORK
 *    occurs, rather than when execve occurs. This must be changed, 
 *    because fd's do not actually get closed until execve occurs.
 *
 * 2. Clone system calls that result in interleaved openings and closings
 *    of file descriptors can result in imperfect replay of a session.
 *    Fixing this will require an architectural change in the program.
 * 
 */
#include <unistd.h>
#include <sys/types.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mysql/mysql.h>
#include <fcntl.h>
#include "DList.h"

// Database-connection strings.
#define MYSQL_HOST   "localhost"     // database server name
#define MYSQL_DB     "log_db"        // database name
#define MYSQL_USER   "audit_user"    // database user id

#define MAX_QUERY_LEN 60000

// Gets the list of events that effect i/o beteen the time of the fork that
// starts a process and the time of the fork that starts its child.
// Format returned:
//
// [date]--[oldfd]--[newfd]--[rc]--[syscall]
#define GET_DUPS_AND_WRITES "SELECT event.date, oldfd, newfd, 0, syscall, dup.id FROM dup, event WHERE \
 dup.parent = event.id AND \
 event.pid = %d AND \
 event.roll_count = %d AND \
 ( dup.cmd = 0 OR dup.cmd = 2 OR dup.cmd IS NULL OR event.syscall = 6) AND \
 event.rc >= 0 \
 UNION \
 SELECT event.date, fd, 0, 0, syscall , io.id FROM io, event WHERE \
 io.parent = event.id AND \
 event.pid = %d AND \
 event.roll_count = %d AND \
 event.syscall = 4 AND \
 event.rc >= 0 \
 UNION \
 SELECT event.date, 0, 0, rc, syscall, event.id FROM event WHERE \
 event.pid = %d AND \
 event.roll_count = %d AND \
 (syscall = 120 OR syscall = 2 OR syscall = 190 OR syscall = 11) \
 ORDER BY event.date"							 

// Get the i/o for a given event.
#define GET_IO "SELECT date, count, data FROM event, io WHERE io.id = %llu AND event.id = io.parent"

// Get the process (pid, rollcount) generated by the given fork.
#define GET_FORK "SELECT rc, child_roll_count FROM event WHERE event.id = %llu"

// Get the fcntl information for a given event.
#define GET_FCNTL "SELECT oldfd, newfd, cmd, arg FROM dup WHERE id = %llu"

// This string is used for debugging purposes.
char tab_string[100] = "";

// This structure holds fd->inherit tuples.
typedef struct {
	int fd;
	int inherited;
} fd_inherit;

// This structure holds io events.
typedef struct {
	unsigned long long timestamp;
	unsigned long long count;
	char * data;
} io_event;

/**
 * Forward-declarations.
 */
int get_io (MYSQL * db, int pid, int rollcount, DList * cur_fds, DList * io_list);


/**
 * Print usage and exit
 */
void usage(char * progname) {
	fprintf (stderr, "Usage: %s pid rollcount\n", progname);
	exit (0);
}

/**
 * Function to compare two io events. (works by timestamp)
 */
int ioComparator (void * io1, void * io2) {
	unsigned long long ts1, ts2;
	ts1 = ((io_event *)(io1))->timestamp;
	ts2 = ((io_event *)(io2))->timestamp;
	
	if (ts1 > ts2) return 1;
	if (ts1 < ts2) return -1;
	return 0;
}


/**
 * Function to compare two integers.
 */
int fdComparator (void * fd1, void * fd2) {
	int i1 = ((fd_inherit *)(fd1))->fd;
	int i2 = ((fd_inherit *)(fd2))->fd;
	return i1 - i2;
}

/**
 * Makes an io event out of a row.
 */ 
io_event * make_io_event (MYSQL_ROW row) {
	io_event * rc = (io_event *)malloc (sizeof(io_event));
	unsigned long long count;
	unsigned long long timestamp;
	char * data;
	
	timestamp = strtoull (row[0], NULL, 10);
	count     = strtoull (row[1], NULL, 10);
	data      = (char *)malloc (count);
	memcpy (data, row[2], count);
	
	rc->data = data;
	rc->timestamp = timestamp;
	rc->count = count;
	
	return rc;
}


/**
 * Deletes an io event.
 */
void delete_io_event (io_event * e) {
	free (e->data);
	free (e);
}


/**
 * Handles write events.
 * Returns 0 on success.
 */
int handle_write (MYSQL * db, DList * fd_list, int fd, unsigned long long id, DList * io_list) {
	int rc;
	int loopcounter = 0;

	// Debugging
	fprintf (stderr, "%swrite (%d, %s)\n", tab_string, fd, fd_list->contains((void *)&fd) ? "Y" : "N");
	
	// If the file-descriptor is in the list of current fd's,
	// then we need to output the data.
	if (fd_list->contains((void *)(&fd))) {
		MYSQL_RES * io_res;
		MYSQL_ROW   io_row;
		char query_string[MAX_QUERY_LEN];
		
		snprintf (query_string, MAX_QUERY_LEN, GET_IO, id);
		rc = mysql_query(db, query_string);
		if (rc) {
			fprintf (stderr, "Unable to execute write query: %s\n", mysql_error(db));
			return 1;
		}
		
		io_res = mysql_store_result(db);
		if (!io_res) {
			fprintf (stderr, "Unable to allocate write result set: %s\n", mysql_error(db));
			return 1;
		}
		
		// Insert the result into our list of io events.
		while ( (io_row = mysql_fetch_row(io_res)) ) {
			
			unsigned long long count = strtoull (io_row[1], NULL, 10);
			
			if (loopcounter) {
				fprintf (stderr, "Warning: Too many events returned by write query. Bailing.\n");
				exit (1);
			}
			
			io_list->insertNode(new DNode ((void *)make_io_event (io_row)));
			
			fprintf (stderr, "%sio-->[%d, %llu]\n", tab_string, fd, count);
						
			loopcounter++;
		}
		
		mysql_free_result (io_res);
	}
	
	return 0;
}


/**
 * Gets a pointer to the fd_inherit data for a 
 * given fd.
 */
fd_inherit * get_inherit (DList * list, int fd) {
	fd_inherit * cur;
	DNode * cur_node;
	list->startIteration (DL_ITERATE_FORWARD);

	while ( (cur_node = ((DNode *)list->getNext())) ) {
		cur = (fd_inherit *)cur_node->getData();
		if (cur->fd == fd) {
			return cur;
			break;
		}
	}
	return NULL;
}

/**
 * Sets the inherit permission for a given fd.
 */ 
void set_inherit (DList * list, int fd, int permission) {
	fd_inherit * cur = get_inherit (list, fd);
	if (cur) {
		cur->inherited = permission;
	}
}


/**
 * Debugging method: Prints out the list of fds
 */
void print_list (DList * list) {
	list->startIteration (DL_ITERATE_FORWARD);
	DNode * cur;
	fprintf (stderr, "%s\t\t\t\t(", tab_string);
	while ( (cur = list->getNext())) {
		fprintf(stderr, "%d, ", *(int *)(cur->getData()));
	}
	fprintf (stderr, ")\n");
	
}


/**
 * Adds a file-descriptor to a list.
 */
void add_fd (DList * list, int fd) {
	
	fd_inherit   * newint  = (fd_inherit *)malloc (sizeof(fd_inherit));
	DNode * newNode = new DNode (newint);
	newint->fd = fd;
	newint->inherited = 1;
	list->pushFront (newNode);
	
	fprintf(stderr, "%s[adding   %d]\n", tab_string, fd);
	print_list(list);
}


/**
 * Deletes a file descriptor from the list.
 */ 
void delete_fd (DList * list, int fd) {
	DNode * cur;
	
	list->startIteration (DL_ITERATE_FORWARD);
	while ( (cur = list->getNext()) ) {
		if ( !fdComparator( cur->getData(), (void *)&fd) ) {
			list->removeNode(cur);
			break;
		}
	}

	fprintf(stderr, "%s[deleting %d]\n", tab_string, fd);
	print_list(list);
}


/**
 * Create a list of only fd's that are to be inherited.
 */
DList * get_inherited_fds (DList * in) {
	DNode *      cur;
	fd_inherit * cur_data;
	DList *      r = new DList();
	
	in->startIteration(DL_ITERATE_FORWARD);
	while ( (cur = in->getNext()) ) {
		cur_data = (fd_inherit *)cur->getData();
		if (cur_data->inherited) {
			fd_inherit * fd = (fd_inherit *)malloc (sizeof(fd_inherit));
			fd->fd = cur_data->fd;
			fd->inherited = 1;
			r->pushBack(new DNode( (void *)fd));
		}
	}
	r->setComparator (fdComparator);
	
	return r;
}


/**
 * Handles dup events.
 */
int handle_dup (MYSQL * db, DList * fd_list,
				int oldfd, int newfd, unsigned long long id) {
	
	int is_open  = (oldfd == -1);
    int is_close = (newfd == -1);
	int is_dup   = (oldfd != -1 && newfd != -1);
	
	// Debugging
	
	if (is_dup) {
		fprintf (stderr, "%sdup (%d, %d)\n", tab_string, oldfd, newfd);

		// If the old fd is in the list and the newfd is not, add the newfd, 
		// and set its inherit permissions to be the same as the oldfd.
		if ( fd_list->contains( (void *)&oldfd) && !fd_list->contains ( (void *)&newfd) ) {
			fd_inherit * old_inherit = get_inherit(fd_list, oldfd);
			add_fd (fd_list, newfd);
			set_inherit (fd_list, newfd, old_inherit->inherited);
		}
		// If the new fd is in the list and the old fd is not, delete the new fd.
		else if ( fd_list->contains( (void *)&newfd) && !fd_list->contains ( (void *)&oldfd) ) {
			delete_fd (fd_list, newfd);
		}
	} else if (is_open) {
		//add_fd (fd_list, newfd);
	} else if (is_close) {
		fprintf (stderr, "%sclose (%d)\n", tab_string, oldfd);
		delete_fd (fd_list, oldfd);
	} else {
		fprintf (stderr, "%sdup: unknown type.\n", tab_string);
	}
	
	return 0;
}

/**
 * Handles fork events.
 * Returns 0 on success.
 */
int handle_fork (MYSQL * db, DList * fd_list, unsigned long long id, DList * io_list, int syscall) {
	int rc;
	int loopcounter = 0;
	MYSQL_RES * fork_res;
	MYSQL_ROW   fork_row;
	DList *     inherit_list = NULL;
	char query_string[MAX_QUERY_LEN];

	// Execute the query to retrieve the fork event.
	snprintf (query_string, MAX_QUERY_LEN, GET_FORK, id);
	rc = mysql_query (db, query_string);
	if (rc) {
		fprintf(stderr, "Unable to execute fork query: %s\n", mysql_error(db));
		return 1;
	}
	
	fork_res = mysql_store_result(db);
	if (!fork_res) {
		fprintf (stderr, "Unable to allocate fork result: %s\n", mysql_error(db));
		return 1;
	}
	
	// This loop should only get executed once, as we expect 
	// only ONE (and EXACTLY one) event to match the id.
	while ( (fork_row = mysql_fetch_row(fork_res)) ) {
		unsigned long long rollcount;
		unsigned long long pid;
		
		if (loopcounter) {
			fprintf (stderr, "Too many rows retrieved for fork event. Aborting.\n");
			exit (1);
		}
		
		if (!fork_row[0] || !fork_row[1]) {
			fprintf (stderr, "Unable to get (rollcount, pid) for child at event %llu.\nPerhaps the post-processing database is uninitialized?\n", id);
			exit (1);
		}
		
		pid = strtoull (fork_row[0], NULL, 10);
		rollcount = strtoull (fork_row[1], NULL, 10);
		
		// Debugging info.
		strcat(tab_string, " ");
		switch (syscall) {
		case 2:
			fprintf (stderr, "%sfork (%llu, %llu)\n", tab_string, pid, rollcount);
			break;
		case 190:
			fprintf (stderr, "%svfork (%llu, %llu)\n", tab_string, pid, rollcount);
			break;
		case 120:
			fprintf (stderr, "%svfork (%llu, %llu)\n", tab_string, pid, rollcount);
			break;
		}
		
		inherit_list = get_inherited_fds (fd_list);

		// Call get_io on this process. Hurray for recursion!
		get_io(db, pid, rollcount, inherit_list, io_list);
		
		delete inherit_list;
		
		// Debugging info.
		tab_string[strlen(tab_string)-1] = 0;
	}
	
	mysql_free_result(fork_res);
	
	return 0;
}

/**
 * Handles fcntl events.
 */
int handle_fcntl (MYSQL * db, DList * fd_list, unsigned long long id) {
	int rc;
	int loopcounter = 0;
	MYSQL_RES * fcntl_res;
	MYSQL_ROW   fcntl_row;
	char query_string [MAX_QUERY_LEN];	
	
	snprintf (query_string, MAX_QUERY_LEN, GET_FCNTL, id);
	rc = mysql_query (db, query_string);
	if (rc) {
		fprintf (stderr, "Unable to execute fcntl query: %s\n", mysql_error(db));
		return 1;
	}
	
	fcntl_res = mysql_store_result(db);
	if (!fcntl_res) {
		fprintf (stderr, "Unable to allocate fcntl result: %s\n", mysql_error(db));
		return 2;
	}
	
	while ( (fcntl_row = mysql_fetch_row (fcntl_res)) ) {
		int oldfd, newfd, cmd, arg;
		oldfd = fcntl_row[0] ? strtol (fcntl_row[0], NULL, 10) : -1;
		newfd = fcntl_row[1] ? strtol (fcntl_row[1], NULL, 10) : -1;
		cmd   = fcntl_row[2] ? strtol (fcntl_row[2], NULL, 10) : -1;
		arg   = fcntl_row[3] ? strtol (fcntl_row[3], NULL, 10) : -1;

		if (loopcounter != 0) {
			fprintf (stderr, "Warning: fcntl query returned more than one result. Bailing.\n");
		}
		
		if (!fcntl_row[2]) {
			fprintf(stderr, "Invalid data returned by fcntl call. Bailing.\n");
			exit (1);
		}
		
		// If it's a dup, call handle_dup to deal with it.
		// From "man 2 fcntl": "The close_on_exec flag of the copy is off, meaning
		// that it will be closed on exec."
		// ... So, we set "inherited" to false, to emmulate system behaviour.
		if (cmd == F_DUPFD) {
			fprintf (stderr, "%sfcntl: calling dup\n", tab_string);
			handle_dup (db, fd_list, oldfd, newfd, id);
			set_inherit (fd_list, newfd, 0);
		}
		
		// If it's a setfd, modify the inherit_list appropriately.
		if (cmd == F_SETFD) {
			// If they're setting the close-on-exec flag to true,
			// add it to the inherit list (if it's not already in).
			// Otherwise, remove it from the list.
			if ( (arg & FD_CLOEXEC) ) {
				fprintf (stderr, "%sfcntl: set_inherit (%d, %d)\n", tab_string, oldfd, 0);
				set_inherit (fd_list, oldfd, 1);
			} else {
				fprintf (stderr, "%sfcntl: set_inherit (%d, %d)\n", tab_string, oldfd, 1);
				set_inherit (fd_list, oldfd, 0);
			}
		}
		loopcounter += 1;
	}
	
	mysql_free_result(fcntl_res);
	return 0;
}

/**
 * This routine does the work. It outputs all the 
 * io events for the current pid/rollcount.
 * Assumes that the database is already connected.
 * Returns 0 on success
 */
int get_io (MYSQL * db, int pid, int rollcount, DList * cur_fds, DList * io_list) {
	int rc;
	char query_string [MAX_QUERY_LEN];
	MYSQL_RES * events;
	MYSQL_ROW cur_event;
	
	DList * fd_list;        // List of currently open fd's.
	
	// Copy the list of fd's.	
	fd_list      = cur_fds->clone();
	
	print_list(fd_list);

	snprintf (query_string, MAX_QUERY_LEN, GET_DUPS_AND_WRITES,
			  pid, rollcount,
			  pid, rollcount,
			  pid, rollcount);
	
	rc = mysql_query (db, query_string);
	if (rc) {
		fprintf (stderr, "Unable to execute event query: %s\n", mysql_error(db));
		return 1;
	}
	
	events = mysql_store_result(db);
	if (!events) {
		fprintf (stderr, "Unable to allocate event result set: %s\n", mysql_error(db));
		return 2;
	}
	
	// Loop through the events.
	while ( (cur_event = mysql_fetch_row (events) ) ) {
		unsigned long long cur_date;
		unsigned long long id;
		int oldfd;
		int newfd;
		int event_rc;
		int syscall;
		
		cur_date = strtoull (cur_event[0], NULL, 10);
		id       = strtoull (cur_event[5], NULL, 10);
		oldfd    = cur_event[1] ? strtol   (cur_event[1], NULL, 10) : -1;
		newfd    = cur_event[2] ? strtol   (cur_event[2], NULL, 10) : -1;
		event_rc = strtol   (cur_event[3], NULL, 10);
		syscall  = strtol   (cur_event[4], NULL, 10);
		
		// Check the type of syscall to determine the action to take.
		switch (syscall) {
		case 11:
			
			break;
		case 120:
		case 2:
		case 190:
			handle_fork  (db, fd_list, id, io_list, syscall);
			break;
		case 4:
			handle_write (db, fd_list, oldfd, id, io_list);
			break;
		case 41:
		case 63:
		case 6:
			handle_dup   (db, fd_list, oldfd, newfd, id);
			break;
		case 55:
		case 221:
			handle_fcntl (db, fd_list, id); 
			break;
		}		
	}
	
    mysql_free_result (events);
	delete fd_list;
		
	return 0;
}

/**
 * Prints out the i/o data.
 */
void print_io (DList * io_list) {
	DNode * cur;
	io_event * cur_event;
	char temp_str[100];
	int fd = fileno (stdout);
	
	io_list->startIteration(DL_ITERATE_BACKWARD);
	while ( (cur = io_list->getNext()) ) {
		cur_event = (io_event *)(cur->getData());
		snprintf (temp_str, 100, "%llu", cur_event->timestamp);
		write (fd, temp_str, strlen(temp_str));
		write (fd, "\t", 1);
		
		fprintf (stderr, "%s\n", temp_str);
		
		snprintf (temp_str, 100, "%llu", cur_event->count);
		write (fd, temp_str, strlen(temp_str));
		write (fd, "\t", 1);
		write (fd, cur_event->data, cur_event->count);
		
		free (cur_event->data);
		free (cur_event);
	}
	
}


/**
 * Command-line routine.
 * Connects to database, then calls get_io.
 */ 
int main (int argc, char * argv[]) {
	int rc;                                  // Utility variable.
	MYSQL * db = NULL;                       // The database connection.
	MYSQL * temp = NULL;                     // Temp database connection.
	int pid;                                 // Pid +
	int rollcount;                           // Rollcount = unique process identifier.
	DList * init_fds;                        // List of initial open FD's.
	DList * io_list;

	// Array of initial open FD's
	int * init_fd_array[3];
	fd_inherit fds [3] = { {0, 1}, {1,1}, {2, 1} };
	
	if (argc != 3) {
		usage(argv[0]);
	}
	
	// Get the pid and rollcount.
	pid = atoi (argv[1]);
	rollcount = atoi (argv[2]);
	
	// Initialize database stuff.
	db = mysql_init (db);
	if (!db) {
		perror ("Unable to allocate memory.\n");
		return 2;
	}
	
	temp = mysql_real_connect (db, MYSQL_HOST, MYSQL_USER, NULL, "log_db", 0, NULL, 0);
	if (temp == NULL || mysql_errno(db)) {
		fprintf (stderr, "Unable to connect MySQL server: %s\n", mysql_error(db));
		return 3;
	}
	
	rc = mysql_select_db (db, MYSQL_DB);
	if (rc) {
		fprintf (stderr, "Unable to select database: %s\n", mysql_error(db));
		return 4;
	}

	// Set up the list of initial open fds.
	init_fd_array[0] = (int *) malloc (sizeof(fd_inherit));
	init_fd_array[1] = (int *) malloc (sizeof(fd_inherit));
	init_fd_array[2] = (int *) malloc (sizeof(fd_inherit));
	memcpy (init_fd_array[0], &fds[0], sizeof(fd_inherit));
	memcpy (init_fd_array[1], &fds[1], sizeof(fd_inherit));
	memcpy (init_fd_array[2], &fds[2], sizeof(fd_inherit));
	
	init_fds = new DList();
	init_fds->pushFront (new DNode (init_fd_array[0]));
    init_fds->pushFront (new DNode (init_fd_array[1]));
	init_fds->pushFront (new DNode (init_fd_array[2]));
	
	// Set the comparator function on the list.
	init_fds->setComparator (fdComparator);
	
	// Set up the io list.
	io_list = new DList();
	io_list->setComparator (ioComparator);

	// Do the work.
	rc = get_io (db, pid, rollcount, init_fds, io_list);

	print_io (io_list);
	
	delete io_list;
	
	return rc;
}