#!/usr/bin/env python
# @lint-avoid-python-3-compatibility-imports
#
# tcprtt    Summarize TCP RTT as a histogram. For Linux, uses BCC, eBPF.
#
# USAGE: tcprtt [-h] [-T] [-D] [-m] [-i INTERVAL] [-d DURATION]
#           [-p LPORT] [-P RPORT] [-a LADDR] [-A RADDR] [-b] [-B] [-e]
#           [-4 | -6]
#
# Copyright (c) 2020 zhenwei pi
# Licensed under the Apache License, Version 2.0 (the "License")
#
# 23-AUG-2020  zhenwei pi  Created this.

from __future__ import print_function
from bcc import BPF
from time import sleep, strftime
from socket import inet_ntop, inet_pton, AF_INET, AF_INET6
import socket, struct
import argparse
import ctypes

# arguments
examples = """examples:
    ./tcprtt            # summarize TCP RTT
    ./tcprtt -i 1 -d 10 # print 1 second summaries, 10 times
    ./tcprtt -m -T      # summarize in millisecond, and timestamps
    ./tcprtt -p         # filter for local port
    ./tcprtt -P         # filter for remote port
    ./tcprtt -a         # filter for local address
    ./tcprtt -A         # filter for remote address
    ./tcprtt -b         # show sockets histogram by local address
    ./tcprtt -B         # show sockets histogram by remote address
    ./tcprtt -D         # show debug bpf text
    ./tcprtt -e         # show extension summary(average)
    ./tcprtt -4         # trace only IPv4 family
    ./tcprtt -6         # trace only IPv6 family
"""
parser = argparse.ArgumentParser(
    description="Summarize TCP RTT as a histogram",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-i", "--interval",
    help="summary interval, seconds")
parser.add_argument("-d", "--duration", type=int, default=99999,
    help="total duration of trace, seconds")
parser.add_argument("-T", "--timestamp", action="store_true",
    help="include timestamp on output")
parser.add_argument("-m", "--milliseconds", action="store_true",
    help="millisecond histogram")
parser.add_argument("-p", "--lport",
    help="filter for local port")
parser.add_argument("-P", "--rport",
    help="filter for remote port")
parser.add_argument("-a", "--laddr",
    help="filter for local address")
parser.add_argument("-A", "--raddr",
    help="filter for remote address")
parser.add_argument("-b", "--byladdr", action="store_true",
    help="show sockets histogram by local address")
parser.add_argument("-B", "--byraddr", action="store_true",
    help="show sockets histogram by remote address")
parser.add_argument("-e", "--extension", action="store_true",
    help="show extension summary(average)")
parser.add_argument("-D", "--debug", action="store_true",
    help="print BPF program before starting (for debugging purposes)")
group = parser.add_mutually_exclusive_group()
group.add_argument("-4", "--ipv4", action="store_true",
    help="trace IPv4 family only")
group.add_argument("-6", "--ipv6", action="store_true",
    help="trace IPv6 family only")
parser.add_argument("--ebpf", action="store_true",
    help=argparse.SUPPRESS)
args = parser.parse_args()
if not args.interval:
    args.interval = args.duration

# define BPF program
bpf_text = """
#include <uapi/linux/ptrace.h>
#include <linux/tcp.h>
#include <net/sock.h>
#include <net/inet_sock.h>
#include <bcc/proto.h>

typedef struct sock_key {
    u64 addr;
    u64 slot;
} sock_key_t;

typedef struct sock_latenty {
    u64 latency;
    u64 count;
} sock_latency_t;

BPF_HISTOGRAM(hist_srtt, sock_key_t);
BPF_HASH(latency, u64, sock_latency_t);

int trace_tcp_rcv(struct pt_regs *ctx, struct sock *sk, struct sk_buff *skb)
{
    struct tcp_sock *ts = (struct tcp_sock *)sk;
    u32 srtt = ts->srtt_us >> 3;
    const struct inet_sock *inet = (struct inet_sock *)sk;

    /* filters */
    u16 sport = 0;
    u16 dport = 0;
    u32 saddr = 0;
    u32 daddr = 0;
    __u8 saddr6[16];
    __u8 daddr6[16];
    u16 family = 0;

    /* for histogram */
    sock_key_t key;

    /* for avg latency, if no saddr/daddr specified, use 0(addr) as key */
    u64 addr = 0;

    bpf_probe_read_kernel(&sport, sizeof(sport), (void *)&inet->inet_sport);
    bpf_probe_read_kernel(&dport, sizeof(dport), (void *)&inet->inet_dport);
    bpf_probe_read_kernel(&family, sizeof(family), (void *)&sk->__sk_common.skc_family);
    if (family == AF_INET6) {
        bpf_probe_read_kernel(&saddr6, sizeof(saddr6),
                              (void *)&sk->__sk_common.skc_v6_rcv_saddr.s6_addr);
        bpf_probe_read_kernel(&daddr6, sizeof(daddr6),
                              (void *)&sk->__sk_common.skc_v6_daddr.s6_addr);
    } else {
        bpf_probe_read_kernel(&saddr, sizeof(saddr), (void *)&inet->inet_saddr);
        bpf_probe_read_kernel(&daddr, sizeof(daddr), (void *)&inet->inet_daddr);
    }

    LPORTFILTER
    RPORTFILTER
    LADDRFILTER
    RADDRFILTER
    FAMILYFILTER

    FACTOR

    STORE_HIST
    key.slot = bpf_log2l(srtt);
    hist_srtt.atomic_increment(key);

    STORE_LATENCY

    return 0;
}
"""

# filter for local port
if args.lport:
    bpf_text = bpf_text.replace('LPORTFILTER',
        """if (ntohs(sport) != %d)
        return 0;""" % int(args.lport))
else:
    bpf_text = bpf_text.replace('LPORTFILTER', '')

# filter for remote port
if args.rport:
    bpf_text = bpf_text.replace('RPORTFILTER',
        """if (ntohs(dport) != %d)
        return 0;""" % int(args.rport))
else:
    bpf_text = bpf_text.replace('RPORTFILTER', '')

def addrfilter(addr, src_or_dest):
    try:
        naddr = socket.inet_pton(AF_INET, addr)
    except:
        naddr = socket.inet_pton(AF_INET6, addr)
        s = ('\\' + struct.unpack("=16s", naddr)[0].hex('\\')).replace('\\', '\\x')
        filter = "if (memcmp(%s6, \"%s\", 16)) return 0;" % (src_or_dest, s)
    else:
        filter = "if (%s != %d) return 0;" % (src_or_dest, struct.unpack("=I", naddr)[0])
    return filter

# filter for local address
if args.laddr:
    bpf_text = bpf_text.replace('LADDRFILTER', addrfilter(args.laddr, 'saddr'))
else:
    bpf_text = bpf_text.replace('LADDRFILTER', '')

# filter for remote address
if args.raddr:
    bpf_text = bpf_text.replace('RADDRFILTER', addrfilter(args.raddr, 'daddr'))
else:
    bpf_text = bpf_text.replace('RADDRFILTER', '')
if args.ipv4:
    bpf_text = bpf_text.replace('FAMILYFILTER',
        'if (family != AF_INET) { return 0; }')
elif args.ipv6:
    bpf_text = bpf_text.replace('FAMILYFILTER',
        'if (family != AF_INET6) { return 0; }')
else:
    bpf_text = bpf_text.replace('FAMILYFILTER', '')
# show msecs or usecs[default]
if args.milliseconds:
    bpf_text = bpf_text.replace('FACTOR', 'srtt /= 1000;')
    label = "msecs"
else:
    bpf_text = bpf_text.replace('FACTOR', '')
    label = "usecs"

print_header = "srtt"
# show byladdr/byraddr histogram
if args.byladdr:
    bpf_text = bpf_text.replace('STORE_HIST', 'key.addr = addr = saddr;')
    print_header = "Local Address"
elif args.byraddr:
    bpf_text = bpf_text.replace('STORE_HIST', 'key.addr = addr = daddr;')
    print_header = "Remote Addres"
else:
    bpf_text = bpf_text.replace('STORE_HIST', 'key.addr = addr = 0;')
    print_header = "All Addresses"

if args.extension:
    bpf_text = bpf_text.replace('STORE_LATENCY', """
    sock_latency_t newlat = {0};
    sock_latency_t *lat;
    lat = latency.lookup(&addr);
    if (!lat) {
        newlat.latency += srtt;
        newlat.count += 1;
        latency.update(&addr, &newlat);
    } else {
        lat->latency +=srtt;
        lat->count += 1;
    }
    """)
else:
    bpf_text = bpf_text.replace('STORE_LATENCY', '')

# debug/dump ebpf enable or not
if args.debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

# load BPF program
b = BPF(text=bpf_text)
b.attach_kprobe(event="tcp_rcv_established", fn_name="trace_tcp_rcv")

print("Tracing TCP RTT... Hit Ctrl-C to end.")

def print_section(addr):
    addrstr = "*******"
    if (addr):
        addrstr = inet_ntop(AF_INET, struct.pack("I", addr))

    avglat = ""
    if args.extension:
        lats = b.get_table("latency")
        lat = lats[ctypes.c_ulong(addr)]
        avglat = " [AVG %d]" % (lat.latency / lat.count)

    return addrstr + avglat

# output
exiting = 0 if args.interval else 1
dist = b.get_table("hist_srtt")
lathash = b.get_table("latency")
seconds = 0
while (1):
    try:
        sleep(int(args.interval))
        seconds = seconds + int(args.interval)
    except KeyboardInterrupt:
        exiting = 1

    print()
    if args.timestamp:
        print("%-8s\n" % strftime("%H:%M:%S"), end="")

    dist.print_log2_hist(label, section_header=print_header, section_print_fn=print_section)
    dist.clear()
    lathash.clear()

    if exiting or seconds >= args.duration:
        exit()