#!/usr/bin/env perl
# dmon -- DTrace-backed IP network real-time monitoring utility
# Copyright (C) 2022-2024 Sergey Matveev <stargrave@stargrave.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
=pod

=head1 DESCRIPTION

dmon is DTrace-backed IP network real-time monitoring utility. It can be
treated as a replacement for iftop utility, that does not always work
properly with IPv6.

It generates DTrace script to gather data about send and received IP
packets, or about TCP, UDP, UDP-Lite, SCTP ones. It collects and
aggregates necessary data and sends it back to that Perl program, that
prettifies its output.

=head1 USAGE

There are two modes of operation:

=over

=item IP packets monitoring

You have to specify exact network interface names you wish to take into
account. Pay attention that all packets are captured, even forwarded
ones, that is useful on router.

=item TCP/UDP/UDP-Lite/SCTP packets monitoring

With C<-p> option you have to specify the transport protocol you want to
monitor. Only explicit send and receive calls are gathered, so no
forwarded packets won't be counted.

=back

=head2 Hostnames resolving

By default this program tries to asynchronously resolve all hostname for
IP addresses in background. You can turn this off with C<-H> option.
C<-P> option also disables service name lookup for the port numbers.

=head1 DISPLAY

Each line shows source and destination addresses with the traffic
direction. If transport mode monitoring is active, then each address
also contains the transport-level port. Unless explicitly disabled by
command lines options, ports and IP addresses are tried to be resolved
to human readable names.

Number of transferred bytes is shown nearby. There are three numbers:
bytes passed for the last second, for last 5 seconds and for the 10
seconds. All calculations are done in powers of two, so C<KB> actually
means C<KiB>, but it is shortened to save the terminal space.
Only the payload is considered, so transport/IP headers overhead is
not counted. All entries are ordered by 10sec speed value.

Top section contains number of send/received packets (also as a tuple of
1s/5s/10s) and total amount of send/received traffic with corresponding
peak values.

=head1 AUTHOR

Sergey Matveev L<mailto:stargrave@stargrave.org>

=cut

use strict;
use warnings;

my $VERSION = "0.1.0";

########################################################################

use Getopt::Std;
$Getopt::Std::STANDARD_HELP_VERSION = 1;

sub VERSION_MESSAGE { print "dmon $VERSION\n" }

my $HeightDefault = 40;
my %Opts = (h => $HeightDefault);

sub HELP_MESSAGE {
    print "Usage: $0 [{-d|-D}] [-h HEIGHT] [-H] [-P]
    {-p [tcp|udp|udplite|sctp] | ifname0 ifname1 ...}
-d -- output only generated DTrace script
-D -- take input from stdin, not summoned DTrace
-h -- set output height, by default it is $HeightDefault
-H -- do not lookup host names
-P -- do not lookup service names
-p -- do not collect IP packets, but specified protocol ones
";
}

getopts "dDh:HPp:", \%Opts;
if ((not exists $Opts{p}) and ($#ARGV < 0)) { die HELP_MESSAGE }
$Opts{h}--;

#########################################################################

use File::Temp qw/tempfile/;
my ($dfh, $dfn) = tempfile();
END { unlink $dfn if defined $dfn }
my $Proto = (defined $Opts{p}) ? $Opts{p} : "ip";
print $dfh "
#pragma D option quiet
#pragma D option switchrate=10Hz
";
my $tick1secT = '
    printa("T %@d %@d %@d %@d\n", @rxpkts, @txpkts, @rxbytes, @txbytes);
    clear(@rxpkts);
    clear(@txpkts);
    clear(@rxbytes);
    clear(@txbytes);
';
if ($Proto eq "ip") {
    my $conds = join " || ", map { "args[3]->if_name == \"$_\"" } @ARGV;
    sub ipprobe {
        my ($probe, $dir) = @_;
        print $dfh "
$probe
/$conds/
{
    @" . $dir . "[args[2]->ip_saddr, args[2]->ip_daddr] = sum(args[2]->ip_plength);
    \@${dir}pkts = count();
    \@${dir}bytes = sum(args[2]->ip_plength);
}\n";
    }
    ipprobe "ip:::send", "tx";
    ipprobe "ip:::receive", "rx";
    print $dfh "
profile:::tick-1sec
{
    printa(\"< %s %s %\@d\\n\", \@rx);
    trunc(\@rx);
    printa(\"> %s %s %\@d\\n\", \@tx);
    trunc(\@tx);
    $tick1secT
}\n";
} else {
    sub nonipprobe {
        my ($probe, $dir) = @_;
        print $dfh "
$probe
{
    @" . $dir . "[args[2]->ip_saddr,
        args[4]->${Proto}_sport,
        args[2]->ip_daddr,
        args[4]->${Proto}_dport] = sum(args[2]->ip_plength);
    \@${dir}pkts = count();
    \@${dir}bytes = sum(args[2]->ip_plength);
}\n";
    }
    nonipprobe "${Proto}:::send", "tx";
    nonipprobe "${Proto}:::receive", "rx";
    print $dfh "
profile:::tick-1sec
{
    printa(\"< [%s]:%d [%s]:%d %\@d\\n\", \@rx);
    trunc(\@rx);
    printa(\"> [%s]:%d [%s]:%d %\@d\\n\", \@tx);
    trunc(\@tx);
$tick1secT\n}\n";
}

########################################################################

my @sizePrefix = ("E", "P", "T", "G", "M", "K");
my @countBound = (1e18, 1e15, 1e12, 1e9, 1e6, 1e3);
my @sizeBound = (1<<60, 1<<50, 1<<40, 1<<30, 1<<20, 1<<10);

sub human {
    my ($s, $suffix, $bounds) = @_;
    foreach my $i (0 .. $#{$bounds}) {
        next if $s < @{$bounds}[$i];
        return sprintf "%.2f%s%s", $s / @{$bounds}[$i], $sizePrefix[$i], $suffix;
    }
    return "$s$suffix";
}

sub humanCount { return human shift, "", \@countBound }
sub humanSize { return human shift, "B", \@sizeBound }

########################################################################

use Socket qw(AF_INET AF_INET6 inet_pton);

sub gethost {
    my $host = shift;
    my $family = (index($host, ":") == -1) ? AF_INET : AF_INET6;
    my $packed = inet_pton $family, $host;
    my @addr = gethostbyaddr $packed, $family;
    return (@addr > 0) ? $addr[0] : "[$host]";
}

use threads;
use threads::shared;
my %HostCache :shared;
my $ResolverQueue;

unless (defined $Opts{H}) {
    use Thread::Queue;
    $ResolverQueue = Thread::Queue->new();
    $ResolverQueue->limit = 0;
    threads->create(sub {
        while (my $host = $ResolverQueue->dequeue()) {
            my $name = gethost $host;
            lock %HostCache;
            $HostCache{$host} = $name;
        }
    })->detach();
}

########################################################################

my %ServCache;

sub humanHost {
    my $pair = shift;
    my ($host, $port);
    if ($Proto eq "ip") {
        $host = $pair;
    } else {
        my $i = rindex($pair, ":");
        $host = substr $pair, 1, $i-2;
        $port = substr $pair, $i+1;
    }

    my $name;
    if ($Opts{H}) {
        $name = "[$host]";
    } else {
        {
            lock %HostCache;
            $name = $HostCache{$host};
        }
        unless (defined $name) {
            $ResolverQueue->enqueue($host);
            $name = $host;
        }
    }
    if ($Proto eq "ip") { return $name };

    my $serv;
    if ($Opts{P}) {
        $serv = $port;
    } else {
        $serv = $ServCache{$port};
        unless (defined $serv) {
            $serv = getservbyport $port, $Proto;
            unless (defined $serv) { $serv = $port }
            $ServCache{$port} = $serv;
        }
    }
    return "$name:$serv";
}

########################################################################

sub agg1510 {
    my $sum = $_[0];
    my @r = ($sum);
    foreach (@_[1 .. 4]) { $sum += $_ }
    push @r, int($sum / 5);
    foreach (@_[5 .. 9]) { $sum += $_ }
    push @r, int($sum / 10);
    return @r;
}

my (@allRxPktsAgg, @allTxPktsAgg);
my (@allRxBytesAgg, @allTxBytesAgg);
my @totalBytesAgg;
my ($allRxBytesPeak, $allTxBytesPeak, $totalBytesPeak) = (0, 0, 0);
format STDOUT_TOP =
Rx: @||||||  @||||||  @||||||    @||||||  @||||||  @||||||    peak: @||||||
humanCount($allRxPktsAgg[0]), humanCount($allRxPktsAgg[1]), humanCount($allRxPktsAgg[2]), humanSize($allRxBytesAgg[0]), humanSize($allRxBytesAgg[1]), humanSize($allRxBytesAgg[2]), humanSize($allRxBytesPeak)
Tx: @||||||  @||||||  @||||||    @||||||  @||||||  @||||||    peak: @||||||
humanCount($allTxPktsAgg[0]), humanCount($allTxPktsAgg[1]), humanCount($allTxPktsAgg[2]), humanSize($allTxBytesAgg[0]), humanSize($allTxBytesAgg[1]), humanSize($allTxBytesAgg[2]), humanSize($allTxBytesPeak)
Total:                           @||||||  @||||||  @||||||    peak: @||||||
humanSize($allRxBytesAgg[0] + $allTxBytesAgg[0]), humanSize($allRxBytesAgg[1] + $allTxBytesAgg[1]), humanSize($allRxBytesAgg[2] + $allTxBytesAgg[2]), humanSize($totalBytesPeak)
------------------------------------------------------------------------------------------------------------------------
.

my ($left, $dir, $right, $agg);
format =
@>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>... @| @<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<...  @||||||  @||||||  @||||||
humanHost($left), $dir, humanHost($right), humanSize($agg->[0]), humanSize($agg->[1]), humanSize($agg->[2])
.

########################################################################

close $dfh;
my $dexec;
if ($Opts{d}) {
    open $dexec, "<", "$dfn" or die "$!";
    while (<$dexec>) { print }
    exit;
}
if ($Opts{D}) {
    $dexec = \*STDIN;
} else {
    open $dexec, "-|", "dtrace -s $dfn" or die "$!";
    sleep 1;
}
unlink $dfn;

my %bytes = ("<" => {}, ">" => {});
my %seen = ("<" => {}, ">" => {});
my @allRxPkts = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
my @allTxPkts = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
my @allRxBytes = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
my @allTxBytes = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
while (<$dexec>) {
    next unless /^[<>T]/;
    my @cols = split;
    if ($cols[0] eq "T") {
        unshift @allRxPkts, $cols[1]; pop @allRxPkts;
        unshift @allTxPkts, $cols[2]; pop @allTxPkts;
        unshift @allRxBytes, $cols[3]; pop @allRxBytes;
        unshift @allTxBytes, $cols[4]; pop @allTxBytes;
        if ($cols[3] > $allRxBytesPeak) { $allRxBytesPeak =  $cols[3] }
        if ($cols[4] > $allTxBytesPeak) { $allTxBytesPeak =  $cols[4] }
        if (($cols[3] + $cols[4]) > $totalBytesPeak) {
            $totalBytesPeak = $cols[3] + $cols[4];
        }
        @allRxPktsAgg = agg1510 @allRxPkts;
        @allTxPktsAgg = agg1510 @allTxPkts;
        @allRxBytesAgg = agg1510 @allRxBytes;
        @allTxBytesAgg = agg1510 @allTxBytes;

        my @res;
        foreach my $dir (keys %bytes) {
            foreach my $src (keys %{$bytes{$dir}}) {
                foreach my $dst (keys %{$bytes{$dir}{$src}}) {
                    my $vals = $bytes{$dir}{$src}{$dst};
                    if ((not exists $seen{$dir}{$src}) or
                        (not exists $seen{$dir}{$src}{$dst})) {
                        unshift @{$vals}, 0;
                    }
                    if (@{$vals} > 10) {
                        $bytes{$dir}{$src}{$dst} = [@{$vals}[0 .. 9]];
                    }
                    my @r = ($dir eq "<") ? ($dst, "<=", $src) : ($src, "=>", $dst);
                    push @r, [agg1510 @{$vals}];
                    if ($r[-1][-1] == 0) {
                        delete $bytes{$dir}{$src}{$dst};
                        next;
                    }
                    push @res, \@r;
                }
            }
        }
        @res = sort { $b->[-1][2] <=> $a->[-1][2] } @res;
        printf "\033[H\033[J";
        $- = 0;
        foreach (@res[0 .. $Opts{h}]) {
            last unless defined;
            ($left, $dir, $right, $agg) = @{$_};
            write;
        }
        %seen = ("<" => {}, ">" => {});
        next;
    }
    unless (exists $bytes{$cols[0]}{$cols[1]}) {
        %{$bytes{$cols[0]}{$cols[1]}} = ();
    }
    unless (exists $bytes{$cols[0]}{$cols[1]}{$cols[2]}) {
        @{$bytes{$cols[0]}{$cols[1]}{$cols[2]}} = (0, 0, 0, 0, 0, 0, 0, 0, 0);
    }
    unshift @{$bytes{$cols[0]}{$cols[1]}{$cols[2]}}, $cols[3];
    unless (exists $seen{$cols[0]}{$cols[1]}) {
        %{$seen{$cols[0]}{$cols[1]}} = ();
    }
    $seen{$cols[0]}{$cols[1]}{$cols[2]} = 1;
}