#!/usr/bin/perl -w

###########################################################
# ns-color.pl 
# haldar@isi.edu
# Takes event/color pairs and colors pkts in nam output files
# Is used to color pkts to identify events in nam
############################################################

use warnings;

# some global declarations
local (@events, %colors, %pkt_ids);
local ($INFILE, $OUTFILE);
local (@buffer, $line);
local ($BUFFER_SIZE) = 1000;

# location of fields in nam output files as generated by ns
# change these values accordingly if nam output format is changed
local ($EVENT_FIELD) = 0;        # position event column
local ($EVENT_TYPE_FIELD) = 6;   # event-type column
local ($COLOR_FIELD) = 16;       # position of pkt color
local ($PID_FIELD) = 14;         # position of pkt-id
local ($ESRC_FIELD) = 18;
local ($EDST_FIELD) = 19;
local ($SEQNUM_FIELD) = 20;

%event_offset_table = (
    
    'TIMEOUT'        => '+1',     # timeout is followed by slowstart, same event
    'SLOW_START'     => '+1',
    'FAST_RECOVERY'  => '+1',
    'FAST_RETX'      => '-1',
    'RENO_FAST_RETX' => '-1',
);

# default colors (IN FUTURE) 
#%event_color_table = ();
    
    
sub usage 
{
    print "Usage: ns-color.pl <inputfile> <outputfile> <event-type 1> <color 1> .. <events-type N> <color N>" ;
    exit 0;
}

sub getOptions 
{
    # read command line options and save them into arrays/hashes
    my ($n);
    %colors = () ;
    
    $n = scalar @ARGV;
    #assuming event/color pair
    if ($n <= 2 || ($n > 2 && ($n % 2 != 0)) ) {
	usage();
    }

    # create events and colors arrays/hashes resp.
    for (my $i=2; $i < $n; $i++) {
	push(@events, $ARGV[$i]);
	push (@{$colors{"$ARGV[$i]"}}, $ARGV[$i+=1]) ;
    }
    
    find_event();
}

# find the event and color pkts
sub find_event {
    
    $INFILE = $ARGV[0];
    $OUTFILE = $ARGV[1];
    
    open (INPUT, "<", $INFILE) or die "can't open $INFILE\n";
    open (OUTPUT, ">", $OUTFILE) or die "can't open $OUTFILE\n";
    
    first_pass_for_pkt_id();
    second_pass_for_coloring_pkt();
    
    close(INPUT);
    close(OUTPUT);
}


# The first pass thru the input file for grabbing 
# the pkt-ids related to the specific events
sub first_pass_for_pkt_id {
    %pkt_ids = ();
    my (@field, $et, $etype);
    my ( $pid, $offset, $pos);
    
    # Reading input file
    while ($line = <INPUT>) {
	
        # maintain a buffer of n previous lines
	# skipping lines not starting with +/-/h/r (non-pkt event lines)
	
	@field = split (" ", $line);
	$et = $field[$EVENT_FIELD];
	
	# interested only in packet level nam traces
	if ($et eq 'E' || $et eq '+' || $et eq '-' || $et eq 'h' || $et eq 'r' || $et eq 'd') { 
	    
	    $pos = push_b (@buffer, $line);
	    
	    # if it is a EVENT line
	    if ($field[$EVENT_FIELD] eq 'E') {
		$etype = $field[$EVENT_TYPE_FIELD];
		
		foreach my $key (@events) {
		    
		    # and if an (exact) matching event is found
		    if ($key =~ $etype) {
			
			# First get the events offset value
			$offset = get_offset ($key);

			# Next get pktid using offset from buffered lines
			$pid = pid_from_offset ($offset, $pos);

			#print "pid for ***pkt***: $pid\n";
			store_color_pid ($key, $pid);  			

			# if fast-retx, mark tcp pkts as well
			if ($key =~ /FAST_RETX/) {
			    
			    # mark 3rd tcp pkt after pkt drop
			    # that triggers fast-retx
			    $pid = pid_for_third_TCP ($pos);
			    #print "pid for TCP pkt: $pid\n";
			    store_color_pid ($key, $pid);

			    # also mark the pkt that is fast-retx'ed
			    $pid = pid_from_offset(+1, $pos);
			    #print "pid for FAST-RETX'ED TCP pkt: $pid\n";
			    store_color_pid ($key, $pid);
			    
			} 
			
		    }
		}
	    }
	}
    }
}

# second pass for coloring pkts using pkt-ids
sub second_pass_for_coloring_pkt {
    
    # go back to start of input file
    seek (INPUT, 0, 0);
    while ($line = <INPUT>) {
	my @field = split (" ", $line);
	my $et = $field[$EVENT_FIELD];

	# skip for all lines not starting with d/+/-/h/r
	if ($et eq 'd' || $et eq '+' || $et eq '-' || $et eq 'h' || $et eq 'r') {
	    
	    my $id = $field[$PID_FIELD];
	
	    # for matching pkt-ids
	    foreach my $pid (keys %pkt_ids) {
		if ($pid == $id) {
		    #color the pkt 
		    my @color = @{$pkt_ids{$pid}};
		    $field[$COLOR_FIELD] = $color[0];
		}
	    }
	}
	
	# join back the line
	$line = join (" ", @field);
	
	#print line in output file;
	print OUTPUT $line,"\n";
	
    }
}


# save pkt ids as pid/color pair
sub store_color_pid {
    my $key = shift;
    my $pid = shift;

    # Now need to get color for this event type
    my @color = @{$colors{$key}};
			
    # create a hash of pkt-ids using pid/color pairs 
    push (@{$pkt_ids{"$pid"}}, $color[0]);
}

sub pid_for_third_TCP {
    
    my $pos = shift;
    
    # get seq num/src/dst-id from 3rd ack
    my @list = get_seq_num (-1, $pos);
    
    # switch src and dst for ack as now we look for 
    # TCP pkts
    my $sid = $list[0];
    $list[0] = $list[1];
    $list[1] = $sid;
    
    # get seq num for 3rd tcp pkt after drop
    $list[2] += 4; 

    # get pkt-id for TCP pkt
    my $pid = pid_from_seq_num ( @list );
    if ($pid == -1) {
	print "match for seqnum not found\n";
	exit 1;
    } else {
	return ($pid);
    }
}

# returns sequence no and src/dest pair for pkt with given offset
sub get_seq_num {

    my $offset = shift; # offset wrt current line pos
    my $pos = shift;    # cuurent line position
    my ($seq, $srcid, $dstid, @vlist);

    # if offset is -ve, read from buffer
    if ($offset < 0) {
	my $line = $buffer[$pos + $offset];
	my @fields = split(" ", $line);
	
	$seq = $fields[$SEQNUM_FIELD];

	my @tmp = split ("{", $fields[$ESRC_FIELD]); 
	$srcid = $tmp[1];

	$dstid = $fields[$EDST_FIELD];
	
	@vlist = ($srcid, $dstid, $seq);

	# if offset is +ve, read from input file	
    } else {
	@vlist = seq_from_infile ($offset);
    }
    return @vlist;
}

sub seq_from_infile {
    my $offset = shift;
    my $i = 0;
    my ($et, $line, @fields);
    my ($vlist, $seq, $srcid, $dstid);
    
    while (($i < $offset) && ($line = <INPUT>)) {
	# read line and place in buffer
	@fields = split(" ", $line);
	$et = $fields[$EVENT_FIELD];

	
	# interested only in packet level nam traces
	if ($et eq '+' || $et eq '-' || $et eq 'h' || $et eq 'r' || $et eq 'd') { 
	    push_b (@buffer, $line);
	    $i++;
	}
    }
    
    $seq = $fields[$SEQ_FIELD];
    $srcid = $fields[$SRCID_FIELD];
    $dstid = $fields[$DSTID_FIELD];
    $vlist = ($srcid, $dstid, $seq);
    return ($vlist);
}


#returns offset value for the given key (event type)    
sub get_offset {
    
    # return offset from event/offset pair
    my $key = shift;
    my $offset = $event_offset_table{$key};
    return ($event_offset_table{$key});
}

# returns pid for pkt with given seqnum/srcid/dstid value
sub pid_from_seq_num {
    my ($sid, $did, $seq) = @_;
    
    for (my $n=0; $n < $BUFFER_SIZE; $n++) {
	
        #look for seqnum match
	my @cols = split (" ", $buffer[$n]);
	my @src = split ("{", $cols[$ESRC_FIELD]);
	
	if ($cols[$SEQNUM_FIELD] == $seq && $src[1] == $sid 
	    && $cols[$EDST_FIELD] == $did) {

	    #return pid if match found
	    return $cols[$PID_FIELD];
	}
    }
    
    # didn't find the seqnum/src/dst match
    return (-1);
}

#returns pkt-id for the given offset and current line value
sub pid_from_offset {
    
    #return packet-id using offset value
    my $offset = shift; # offset wrt current line pos
    my $pos = shift;    # cuurent line position
    my $pid;

    # if offset is -ve, read from buffer
    if ($offset < 0) {
	my $line = $buffer[$pos + $offset];
	my @fields = split(" ", $line);
	$pid = $fields[$PID_FIELD];
    
    # if offset is +ve, read from input file
    } else {
	$pid = pid_from_infile ($offset);
    }
    return $pid;
}
	    
sub pid_from_infile {
    my $offset = shift;
    my $i = 0;
    my ($et, $line, @fields, $pid);
    
    while (($i < $offset) && ($line = <INPUT>)) {
	# read line and place in buffer
	@fields = split(" ", $line);
	$et = $fields[$EVENT_FIELD];

	
	# interested only in packet level nam traces
	if ($et eq '+' || $et eq '-' || $et eq 'h' || $et eq 'r' || $et eq 'd') { 
	    push_b (@buffer, $line);
	    $i++;
	}
    }
    
    return ($pid = $fields[$PID_FIELD]);
}



sub push_b {
    
    # pop first line out (from left hand side) 
    # if buffer size greater than BUFFER_SIZE
    my $size = scalar @buffer;
    if ($size >= $BUFFER_SIZE) {
	shift @buffer;
    }
    
    # then push current line into buffer
    return (push (@buffer, $line)-1);
    
}






getOptions();