#!/usr/bin/perl

# Program: iotrace
# Purpose: Act like strace for Linux enough to log STDIN STDOUT STDERR I/O.
# This is used for man-in-the-middle logging but will not alter any packets.

use strict;
use warnings;
use Getopt::Long qw(GetOptions);
use IPC::Open3 qw(open3);
use IO::Select;
use IO::Handle;

Getopt::Long::Configure("require_order");
Getopt::Long::Configure("bundling");
GetOptions
    "o=s"   => \(my $output_log_file),
    "v+"    => \(my $verbose = 0),
    "x+"    => \(my $heX_ify = 0),
    "t+"    => \(my $timing = 0),
    "f+"    => \(my $follow_fork = 0),
    "q+"    => \(my $quiet = 0), # Ignored
    "s=i"   => \(my $size_of_strings), # Ignored
    "e=s"   => \(my $events),  # Ignored
    or die "Usage> $0 -o <output_log> CMD [ARGS]\n";

my @run = @ARGV or die "$0: Missing command?\n";
my $pid = 0;
our $has_hires = eval { require Time::HiRes; 1 };
our $patience_idle = 8000; # Seconds to wait for I/O while process is still running
our $patience_kill_mute = $has_hires ? 5.2 : 6; # Seconds to wait for "mute" process to terminate. If all its handles are closed, then force SIGTERM and SIGKILL if still running after waiting this long.
our $heartbeat_grind_mute = 0.2; # Interval Seconds between SIGCHLD check to test if "mute" child died yet.
our $implicit_close_before_chld = 0.08; # Maximum Seconds to wait for SIGCHLD after all handles have been closed in order to consider it an implicit close instead of explicit close.
our $implicit_chld_before_close = $has_hires ? 0.25 : 1; # Maximum Seconds after receiving SIGCHLD to consider a close to be implicit. Any longer is considered an explicit close.
our $patience_zombie_breather = $has_hires ? 5.6 : 7; # Maximum Seconds of uninterrupted silence after receiving SIGCHLD to wait for the process to close all handles. Any longer, then all open pipes are implicitly slapped closed to sufficate the zombie process. Then it is reaped to log the correct exit status, and the descent process should receive a SIGPIPE if ever attempting to write to its output handle in the future.

if ($timing > 1 and not $has_hires) {
    # If requesting to log high precision, but without HiRes, then just change "-tt" to "-t" to show integer time
    $timing = 1;
}

sub now { $has_hires ? Time::HiRes::time() : time }

sub t {
    my $h = "";
    if ($follow_fork == 1) { $h .= "$pid "; }
    if ($timing) {
        my $now = $timing > 1 ? now : time;
        my @t = localtime $now;
        $h .= sprintf "%02d:%02d:%02d", $t[2], $t[1], $t[0];
        $h .= sprintf ".%06d", 1000000*($now - int($now)) if $timing > 1;
        $h .= " ";
    }
    return $h;
}

# Escape strings like strace does
sub e {
    my $chars = shift;
    if ($heX_ify > 1) {
        # -xx: Super Hex Encode everything
        $chars =~ s/([\s\S])/sprintf "\\x%02x", ord($1)/eg;
    }
    else {
        # Both \\ and \" really need to be escaped
        # But add other helpful chars to make it easier for Perl to read too.
        $chars =~ s/([\\\"\'\$\@])/\\$1/g;
        # Special backslash escape chars for easy legibility
        $chars =~ s/\t/\\t/g;
        $chars =~ s/\r/\\r/g;
        $chars =~ s/\n/\\n/g;
        if ($heX_ify) {
            # -x: Hex Encode only non-ascii chars
            $chars =~ s/([^\ -\~])/sprintf "\\x%02x", ord $1/eg;
        }
        else {
            # Default is octal encoding non-ascii only
            $chars =~ s/([^\ -\~])/sprintf "\\%03o", ord $1/eg;
        }
   }
   return qq{"$chars"};
}

my $full = $run[0];
if (eval { require File::Which; 1; } and $full =~ m{^([\w\-\.]+)$}) {
    $full = File::Which::which($1);
}
if (!$full or $full =~ m{^/} && !-x $full) {
    die "$run[0]: No such file or directory\n";
}

my ($in, $out, $err);
# open3 can't vivify STDERR from undef for some reason
if ($err = IO::Handle->new) {
    my @r = @run;
    $r[0] = $full if $full;
    # Launch target program
    $pid = open3 $in, $out, $err, @r or die "$r[0]: fork exec failure: $!\n";
}

my $fh_log;
if (defined $output_log_file) {
    $output_log_file .= ".$pid" if $follow_fork > 1;
    open $fh_log, ">", $output_log_file or die "$output_log_file: open failure: $!\n";
}
else {
    # XXX - Is it ok to spew all the trace lines out to STDERR if no -o option provided?
    open $fh_log, ">&STDERR";
}
$fh_log->autoflush(1);

print $fh_log t."execve(".e($full).", [".join(', ', map { e($_) } @run).'], '.(
    $verbose ? (
        '['.join(', ', map { e("$_=$ENV{$_}") } sort keys %ENV).']'
    ) :
    \%ENV.' /* '.(scalar keys %ENV).' vars */'
).") = 0\n";

my $child_died = 0;
$SIG{CHLD} = sub { $child_died = now; };

# Map each handle to its corresponding handle
my $proxy = {
    fileno($in)     => *STDIN,
    fileno($out)    => *STDOUT,
    fileno($err)    => *STDERR,

    fileno(*STDIN)  => $in,
    fileno(*STDOUT) => $out,
    fileno(*STDERR) => $err,
};

my $implicitly_closed = {};
my $io = IO::Select->new(values %$proxy);
my $writers = IO::Select->new($in, \*STDOUT, \*STDERR);
# Loop while still open handles or during brief implicit close detection
while ($io->handles or keys %$implicitly_closed) {
    my $maximum_timeout =
        keys %$implicitly_closed ? $implicit_close_before_chld :
        $child_died ? $patience_zombie_breather :
        $patience_idle;
    my @ready = $io->count ? $io->can_read($maximum_timeout) : do {select undef,undef,undef, $maximum_timeout; ()};
    foreach my $fh (@ready) {
        my $fn = fileno($fh) // next;
        my $pr = $proxy->{$fn} or die "Fileno $fn: Impossible Implementation Crash! $!\n";;
        # Find original fileno (STDIN=0, STDOUT=1, STDERR=2):
        my $real_fileno = $fn < 3 ? $fn : fileno($pr);
        if ($writers->exists($fh)) {
            # $fh should only be written to, so if it's suddenly "READABLE",
            # that means it woke up due to the pipe being close()d on the other end.
            # Never attempt to actually read from a "writers" handle.
            # Only log explicit close(). Don't bother if it's probably just an implicit close() upon exit.
            print $fh_log t."close($real_fileno) = 0\n" if !$child_died;
            # Close both sides since there's nowhere for any data to go anymore:
            $io->remove($fh);
            $io->remove($pr);
            $writers->remove($fh);
            close $fh;
            close $pr;
            next;
        }
        my $bytes = sysread($fh, (my $buffer), 16384);
        # Only STDIN (fileno = 0) is "read", otherwise "write"
        my $op = $real_fileno ? "write" : "read";
        print $fh_log t."$op($real_fileno, ".e($buffer).", $bytes) = $bytes\n" if $bytes or $op eq "read";
        if ($bytes) {
            # Forward non-empty packet to the proxy file handle
            syswrite($pr, $buffer, $bytes);
        }
        else {
            # Getting ZERO bytes always means the file handle just closed.
            # Quit listening to this anymore:
            $io->remove($fh);
            # And then close the file handle for real:
            close($fh);

            # And quit listening on the corresponding handle too:
            $io->remove($pr);

            if ($real_fileno == 0 # Immediately close corresponding handle for STDIN immediately in order to signal target program its input stream has ended.
                or $child_died && now > $child_died + $implicit_chld_before_close) { # Or if it exited long enough ago, then this must have been an explicit output handle close from a backgrounded process, and it's safe to log, (even if it was implicitly closed by the grandchild or descendent process).
                close($pr);

                # If it was STDIN (fd 0), then the invoker probably closed it before the target program. But just log the close acting like the target program called close itself (even if STDIN is being ignored by the target program).
                print $fh_log t."close($real_fileno) = 0\n";
            }
            else {
                # An output handle (STDOUT fd 1 or STDERR fd 2) was just closed, but it might have just been an implicit close since the process is still running or exited too recently.
                # So don't log it yet. Just flag it for now. Then keep a close eye on the $child_died timer to determine which way it was.
                # The goal is to mimic the same behavior as the target process whether to do an implicit or explicit close.
                $implicitly_closed->{$real_fileno} = $pr;
            }
        }
    }
    if (!@ready and keys %$implicitly_closed) {
        # Implicit detection timeout exceeded. No more waiting allowed.
        # All $implicitly_closed file handles must be imminently closed!
        if (!$io->count and $child_died) {
            # Must exit immediately in order to implicitly close any $implicitly_closed handles
            last;
        }
        # Otherwise, must explicitly close them all now and log it, then continue the select loop.
        foreach my $fileno (keys %$implicitly_closed) {
            close(delete $implicitly_closed->{$fileno});
            print $fh_log t."close($fileno) = 0\n";
        }
    }
    if ($child_died) {
        if (now > $child_died + $patience_zombie_breather) {
            # Ran out of patience. Just leave to implicitly close all handles.
            last;
        }
        else {
            # Zombie handle said something, so restart the timer, and start the waiting all over again..
            $child_died = now;
        }
    }
}

if (!$child_died) {
    # If process still running, then all handles must have already been closed. So just wait a little bit for termination.
    my $patient = $patience_kill_mute + now;
    select undef,undef,undef,$heartbeat_grind_mute while !$child_died and $patient > now;

    # If STILL running, then offer a little bit more help to terminate
    !$child_died and kill TERM => $pid and sleep 1 and kill KILL => $pid;
}

# Block waiting for process to exit
waitpid( $pid, 0 );
my $signal = $? & 127;
my $child_exit_status = $? >> 8;
if ($signal) {
    my $sig_name = eval { require Config; %Config::Config and [split / /, $Config::Config{sig_name}]->[$signal] } || $signal;
    print $fh_log t."--- GOT SIG$sig_name ($signal) ---\n";
}
print $fh_log t."+++ exited with $child_exit_status +++\n";
close $fh_log;
exit $child_exit_status;
