#!/usr/bin/env perl
# dmon -- DTrace-backed TCP/IP real-time monitoring utility
# Copyright (C) 2022 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/>.

use strict;
use warnings;

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

if ($#ARGV < 0) {
    die "Usage: $0 [tcp|ifname0 ifname1 ...]\n";
};
use File::Temp qw/tempfile/;
my ($dfh, $dfn) = tempfile();
END { unlink $dfn };
print $dfh "
#pragma D option quiet
#pragma D option switchrate=10Hz

profile:::tick-1sec
{
    printf(\"T\\n\");
}\n";
if ($ARGV[0] eq "tcp") {
    sub tcpprobe {
        my ($probe, $dir) = @_;
        print $dfh "
$probe
{
    @" . $dir . "[args[2]->ip_saddr,
        args[4]->tcp_sport,
        args[2]->ip_daddr,
        args[4]->tcp_dport] = sum(args[2]->ip_plength);
    " . $dir . "pkts++;
}\n";
    };
    tcpprobe "tcp:::send", "tx";
    tcpprobe "tcp:::receive", "rx";
    print $dfh '
profile:::tick-1sec
{
    printa("< [%s]:%d [%s]:%d %@d\n", @rx);
    clear(@rx);
    printa("> [%s]:%d [%s]:%d %@d\n", @tx);
    clear(@tx);
    printf("I %d\n", rxpkts);
    printf("O %d\n", txpkts);
    rxpkts = 0;
    txpkts = 0;
    printf("T\n");
}';
} else {
    sub ipprobe {
        my ($probe, $dir) = @_;
        my $conds = join " || ", map { "args[3]->if_name == \"$_\"" } @ARGV;
        print $dfh "
$probe
/$conds/
{
    @" . $dir . "[args[2]->ip_saddr,
        args[2]->ip_daddr] = sum(args[2]->ip_plength);
    " . $dir . "pkts++;
}\n";
    };
    ipprobe "ip:::send", "tx";
    ipprobe "ip:::receive", "rx";
    print $dfh '
profile:::tick-1sec
{
    printa("< %s %s %@d\n", @rx);
    clear(@rx);
    printa("> %s %s %@d\n", @tx);
    clear(@tx);
    printf("I %d\n", rxpkts);
    printf("O %d\n", txpkts);
    rxpkts = 0;
    txpkts = 0;
    printf("T\n");
}';
};

my $height = (defined $ENV{DMON_HEIGHT}) ? $ENV{DMON_HEIGHT} : 40;
$height--;

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

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, "iB", \@sizeBound }
#
########################################################################

my (%tx, %rx);
my ($pktsRx, $pktsTx, $bytesRx, $bytesTx) = (0, 0, 0, 0);
format STDOUT_TOP =
@>>>>>>>>>> / @<<<<<<<<<<                  @>>>>>>>>>>>> / @<<<<<<<<<<<<
humanCount($pktsRx), humanCount($pktsTx), humanSize($bytesRx), humanSize($bytesTx)
------------------------------------------------------------------------
.

my ($left, $dir, $right, $size);
format STDOUT =
@<<<<<<<<<<<<<<<<<<<<<<<<<< @| @>>>>>>>>>>>>>>>>>>>>>>>>>> @>>>>>>>>>>>>
$left, $dir, $right, humanSize($size)
.

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

my $dexec;
open $dexec, "-|", "dtrace -s $dfn" or die "$!";
sleep 1;
unlink $dfn;

my $xx;
while (<$dexec>) {
    next unless /^[<>TIO]/;
    my @cols = split;
    if ($cols[0] eq "T") {
        my @res;
        foreach my $src (keys %rx) {
            foreach my $dst (keys %{$rx{$src}}) {
                push @res, [$dst, "<", $src, $rx{$src}{$dst}];
            };
        };
        foreach my $src (keys %tx) {
            foreach my $dst (keys %{$tx{$src}}) {
                push @res, [$src, ">", $dst, $tx{$src}{$dst}];
            };
        };
        @res = sort { @{$b}[-1] <=> @{$a}[-1] } @res;
        printf "\033[H\033[J";
        $- = 0;
        foreach (@res[0..$height]) {
            last unless defined;
            ($left, $dir, $right, $size) = @{$_};
            write STDOUT;
        };
        (%rx, %tx) = (), ();
        ($pktsRx, $pktsTx, $bytesRx, $bytesTx) = (0, 0, 0, 0);
        next;
    };
    if ($cols[0] eq "I") {
        $pktsRx += $cols[1];
        next;
    };
    if ($cols[0] eq "O") {
        $pktsTx += $cols[1];
        next;
    };
    if ($cols[0] eq "<") {
        $xx = \%rx;
        $pktsRx++;
        $bytesRx += $cols[3];
    } else {
        $xx = \%tx;
        $pktsTx++;
        $bytesTx += $cols[3];
    };
    if (defined ${$xx}{$cols[1]}) {
        ${$xx}{$cols[1]}{$cols[2]} += $cols[3];
    } else {
        ${$xx}{$cols[1]} = {$cols[2] => $cols[3]};
    };
};