#!/usr/bin/perl -w
use strict;

use Net::Pcap;
use Socket;
use IO::File;
use Cwd 'abs_path';
use Errno 'EMFILE';

$|=1 if ! -f STDOUT;

my $live=0;
if($ARGV[0] && $ARGV[0] eq '-l') {
  shift;
  $live=1;
}

@ARGV>=3 or die
  "Usage: $0 {inputfile|-l interface} outdir backend [arguments for backend]\n";
-d $ARGV[1] or mkdir($ARGV[1], 0777) or die "$ARGV[1]: $!\n";
$ARGV[1] = abs_path($ARGV[1]); # the backend can chdir

my ($err, $cap);
if($live) {
  $cap=Net::Pcap::open_live($ARGV[0], 2000, 0, 1000, \$err);
} else {
  $cap=Net::Pcap::open_offline($ARGV[0], \$err);
}
$cap or die "$err\n";
-d $ARGV[1] or mkdir($ARGV[1], 0777) or die "$ARGV[1]: $!\n";

# The "backend" file must define tcpstream_init and tcpstream_callback. The
# init function is given the extra command line arguments once at the start
# of processing, and the callback function is called at the completion of
# every connection. Args: i file, m file, o file
use FindBin;
use lib $FindBin::Bin;
require $ARGV[2];
tcpstream_init(@ARGV[3..$#ARGV]);

#my @fin=('','F');
#my @syn=('','S');
#my @rst=('','R');
#my @psh=('','P');
#my @ack=('','A');
#my @urg=('','U');

sub cleanstr_noquote
{
  my $s=$_[0];
  my $l=length($s);
  my $ret='';
  my $i;
  for($i=0;$i<$l;++$i) {
    my $c=substr($s, $i, 1);
    my $cc=ord($c);
    if($c eq "\t") {
      $ret .= "\\t";
    } elsif($c eq "\r") {
      $ret .= "\\r";
    } elsif($c ne "\n" && ($cc < 32 || $cc > 126)) {
      my $nextc=substr($s, $i+1, 1);
      if($nextc eq '' || $nextc =~ /\d/) {
        $ret .= "\\".sprintf("%03o", $cc);
      } else {
        $ret .= "\\".sprintf("%o", $cc);
      }
    } elsif($c eq "\\") {
      $ret .= "\\".$c;
    } else {
      $ret .= $c;
    }
  }
  return $ret;
}

sub newfile;

use constant INITIATOR => 0;
use constant OTHER => 1;
use constant BOTH => 2;

my $num=1;
my %conn;
# These next 4 variables are per-connection, and should be allocated as part
# of a connection object when the initial SYN is encountered.
my %hdr;
while(defined(my $pkt=Net::Pcap::next($cap, \%hdr))) {
  my $ethertype=substr($pkt, 12, 2);
  if(unpack("n", $ethertype) != 0x800) {
    #warn "Skipping non-IP packet #$num\n";
    next;
  }

  my $ippkt=substr($pkt, 14);
  my $iplen=unpack("n", substr($ippkt, 2, 2));
  $ippkt=substr($ippkt, 0, $iplen);
  my $iphlen=4*(unpack("C", $ippkt) & 0xf);
  my $ipsrc=inet_ntoa(substr($ippkt, 12, 4));
  my $ipdst=inet_ntoa(substr($ippkt, 16, 4));
  my $proto=unpack("C", substr($ippkt, 9, 1));
  if($proto != 6) {
    #warn "Skipping non-TCP packet #$num\n";
    next;
  }
  my $ipflags=unpack("C", substr($ippkt, 6, 1))>>5;
  my $fragoff=unpack("C", substr($ippkt, 6, 1)) & 0x1f;
  if(($ipflags & 0x101) || $fragoff) {
    die "Can't handle fragmented packet #$num\n";
  }

  my $tcppkt=substr($ippkt, $iphlen);
  my ($srcport, $dstport, $seq, $ackval)=unpack("nnNN", $tcppkt);
  my $flags=unpack("C", substr($tcppkt, 13, 1));
  my ($fin, $syn, $rst, $psh, $ack, $urg) =
    map { $_==($flags & $_)?1:0 } map { 1<<$_ } 0..5;
  #$flags="$fin[$fin]$syn[$syn]$rst[$rst]$psh[$psh]$ack[$ack]$urg[$urg]";
  my $conndesc="${ipsrc}_${srcport}_${ipdst}_${dstport}";
  my $rconndesc="${ipdst}_${dstport}_${ipsrc}_${srcport}";
  if(!$conn{$conndesc}) {
    if(!$syn || $ack) {
      #warn "First packet in $conndesc is not SYN\n";
      next;
    }
    $conn{$conndesc} = $conn{$rconndesc} = {
      lastsyn => $hdr{tv_sec},
      initiator => "$ipsrc:$srcport",
      base => "${conndesc}_$hdr{tv_sec}_$hdr{tv_usec}",
      expected_seq => [],
      sacked => [], # Actually all out-of-order packets, whether SACKed or not
      out => [],
      fin => []
    };
    my $base=$conn{$conndesc}{base};
    #-f "$ARGV[1]/${base}_i" and die "${base}_i already exists";
    #-f "$ARGV[1]/${base}_o" and die "${base}_o already exists";
    #-f "$ARGV[1]/${base}_m" and die "${base}_m already exists";
    $conn{$conndesc}{out}[INITIATOR] = newfile("$ARGV[1]/${base}_i")
      or die "${base}_i: $!\n";
    $conn{$conndesc}{out}[OTHER] = newfile("$ARGV[1]/${base}_o")
      or die "${base}_o: $!\n";
    $conn{$conndesc}{out}[BOTH] = newfile("$ARGV[1]/${base}_m")
      or die "${base}_m: $!\n";
  }
  my $conn=$conn{$conndesc};
  $conn->{lastsyn}=$hdr{tv_sec} if $syn;
  my $role=($conn->{initiator} eq "$ipsrc:$srcport")?INITIATOR:OTHER;
  ($conndesc,$rconndesc)=($rconndesc,$conndesc) if $role==OTHER;
  my $dataoff=4*(unpack("C", substr($tcppkt, 12, 1)) >> 4);
  my $data=substr($tcppkt, $dataoff);
  my $toprint='';
  if(!defined($conn->{expected_seq}[$role])) {
    if($syn) {
      $conn->{expected_seq}[$role]=$seq+1;
      $conn->{expected_seq}[$role] &= 0xffffffff;
    } elsif(!$rst) {
      warn "weirdness at #$num\n";
      next;
    }
  } else {
    # Remove stuff we've already seen
    my $diff=unpack "L", pack "l", $conn->{expected_seq}[$role] - $seq;
    if($diff < 1048576) {
      #warn "#$num is a (partial?) repeat of previously ACKable data\n";
      substr($data, 0, $diff)='';
      $seq+=$diff;
    }
    if($seq == $conn->{expected_seq}[$role]) {
      $toprint.=$data;
      $conn->{expected_seq}[$role]+=length($data);
      $conn->{expected_seq}[$role] &= 0xffffffff;
      if($fin && !$conn->{fin}[$role]) {
        $conn->{fin}[$role]=1;
        ++$conn->{expected_seq}[$role];
        $conn->{expected_seq}[$role] &= 0xffffffff;
      }
      my $done=0;
      while(!$done) {
        $done=1;
        while(exists $conn->{sacked}[$role]{$conn->{expected_seq}[$role]}) {
          $done=0;
          my $sack=delete $conn->{sacked}[$role]{$conn->{expected_seq}[$role]};
          $toprint .= $sack->{data};
          $conn->{fin}[$role]=1 if $sack->{fin};
          $conn->{expected_seq}[$role]+=length($sack->{data});
          $conn->{expected_seq}[$role] &= 0xffffffff;
        }
        # Search the sacked data for anything we don't need anymore.
        keys %{$conn->{sacked}[$role]};
        while(my ($sackseq, $sack)=each %{$conn->{sacked}[$role]}) {
          my $diff=unpack "L", pack "l", $conn->{expected_seq}[$role] - $sackseq;
          if($diff < 1048576) {
            substr($sack->{data}, 0, $diff)='';
            delete $conn->{sacked}[$role]{$sackseq};
            $conn->{sacked}[$role]{$sackseq+$diff}=$sack;
            # Get out of the each() loop because the hash has been modified.
            $done=0;
            last;
          }
        }
      }
    } else {
      $diff=unpack "L", pack "l", $seq - $conn->{expected_seq}[$role];
      if($diff < 1048576) {
        #warn "#$num is SACKable, with seq=$seq expected=$conn->{expected_seq}[$role] fin=$fin\n";
        $conn->{sacked}[$role]{$seq+($fin?1:0)}={ data=>$data, fin=>$fin };
      }
    }
  }

  #my $hex=join ' ', map { sprintf "%02x", ord $_ } split //, $data;
  #print "Read packet #$num from $ipsrc:$srcport to $ipdst:$dstport, flags=$flags data=$hex\n";

  if(length($toprint)) {
    $conn->{out}[$role]->print($toprint);
    $conn->{out}[BOTH]->print("\c[[1m") if $role==INITIATOR;
    $conn->{out}[BOTH]->print(cleanstr_noquote($toprint));
    $conn->{out}[BOTH]->print("\c[[m") if $role==INITIATOR;
  }
  if($rst ||
     (!$fin && $ack &&
      $conn->{fin}[INITIATOR] &&
      $conn->{fin}[OTHER] &&
      $ackval == $conn->{expected_seq}[1-$role])) {
    # TODO: technically, the test above doesn't actually prove that both FINs
    # have been ACKed. In practice, it may not matter.
    if(!$rst) {
      for(values(%{$conn->{sacked}[INITIATOR]}),
          values(%{$conn->{sacked}[OTHER]})) {
        next if !length($_->{data});
        use Data::Dumper;
        die "Leftover data!\n".Dumper($conn);
      }
    }
    my $suff='';
    if($rst) {
      my $base=$conn->{base};
      rename("$ARGV[1]/${base}_i", "$ARGV[1]/${base}_i.reset") or die;
      rename("$ARGV[1]/${base}_o", "$ARGV[1]/${base}_o.reset") or die;
      rename("$ARGV[1]/${base}_m", "$ARGV[1]/${base}_m.reset") or die;
      $suff='.reset';
    }
    $conn->{out}[INITIATOR]->close() or die "${conndesc}_i: $!\n";
    $conn->{out}[OTHER]->close() or die "${conndesc}_o: $!\n";
    $conn->{out}[BOTH]->close() or die "${conndesc}_m: $!\n";
    delete @conn{$conndesc,$rconndesc};
    tcpstream_callback("$ARGV[1]/$conn->{base}_i$suff",
                       "$ARGV[1]/$conn->{base}_m$suff",
                       "$ARGV[1]/$conn->{base}_o$suff");
  }
} continue {
  ++$num;
}
keys %conn;
while(my ($conndesc,$conn)=each %conn) {
  rename("$ARGV[1]/$conn->{base}_i", "$ARGV[1]/$conn->{base}_i.incomplete") or next;
  rename("$ARGV[1]/$conn->{base}_o", "$ARGV[1]/$conn->{base}_o.incomplete");
  rename("$ARGV[1]/$conn->{base}_m", "$ARGV[1]/$conn->{base}_m.incomplete");
  $conn->{out}[INITIATOR]->close();
  $conn->{out}[OTHER]->close();
  $conn->{out}[BOTH]->close();
  tcpstream_callback("$ARGV[1]/$conn->{base}_i.incomplete",
                     "$ARGV[1]/$conn->{base}_m.incomplete",
                     "$ARGV[1]/$conn->{base}_o.incomplete");
}

# Open a file for writing, and if EMFILE occurs, close some half-open
# connections and try again. The output files for half-open connections
# must be empty, so don't bother with the callback. Just delete them.
sub newfile
{
  my $ret=IO::File->new(">$_[0]");
  return $ret if $ret || $! != EMFILE;

  print STDERR "Closing stale filehandles...\n";

  my $now=$hdr{tv_sec}; # The most recently read packet

  my $done=0;
  while(!$done) {
    $done=1;
    keys %conn;
    while(my ($conndesc,$conn)=each %conn) {
      next if defined($conn->{expected_seq}[INITIATOR]) && # Always true?
              defined($conn->{expected_seq}[OTHER]);
      #print STDERR "Considering $conndesc...\n";
      next if $conn->{lastsyn} > $now-90;
      print STDERR "$conndesc is stale.\n";
      unlink("$ARGV[1]/$conn->{base}_i") or die "$ARGV[1]/$conn->{base}_i: $!\n";
      unlink("$ARGV[1]/$conn->{base}_o") or die "$ARGV[1]/$conn->{base}_o: $!\n";
      unlink("$ARGV[1]/$conn->{base}_m") or die "$ARGV[1]/$conn->{base}_m: $!\n";
      $conn->{out}[INITIATOR]->close() or die "$conn->{base}_i: $!\n";
      $conn->{out}[OTHER]->close() or die "$conn->{base}_o: $!\n";
      $conn->{out}[BOTH]->close() or die "$conn->{base}_m: $!\n";
      $conndesc =~ /^([^_]*)_([^_]*)_([^_]*)_([^_]*)$/ or die "Internal bug [$conndesc]";
      my $rconndesc="${3}_${4}_${1}_${2}";
      delete @conn{$conndesc,$rconndesc};
      # We must restart the each() loop because the keys have changed
      $done=0;
      last;
    }
  }

  return IO::File->new(">$_[0]");
}