1#!/usr/bin/env python
2#
3# solisten      Trace TCP listen events
4#               For Linux, uses BCC, eBPF. Embedded C.
5#
6# USAGE: solisten.py [-h] [-p PID] [--show-netns]
7#
8# This is provided as a basic example of TCP connection & socket tracing.
9# It could be useful in scenarios where load balancers needs to be updated
10# dynamically as application is fully initialized.
11#
12# All IPv4 listen attempts are traced, even if they ultimately fail or the
13# the listening program is not willing to accept().
14#
15# Copyright (c) 2016 Jean-Tiare Le Bigot.
16# Licensed under the Apache License, Version 2.0 (the "License")
17#
18# 04-Mar-2016	Jean-Tiare Le Bigot	Created this.
19
20import os
21from socket import inet_ntop, AF_INET, AF_INET6, SOCK_STREAM, SOCK_DGRAM
22from struct import pack
23import argparse
24from bcc import BPF
25import ctypes as ct
26from bcc.utils import printb
27
28# Arguments
29examples = """Examples:
30    ./solisten.py              # Stream socket listen
31    ./solisten.py -p 1234      # Stream socket listen for specified PID only
32    ./solisten.py --netns 4242 # " for the specified network namespace ID only
33    ./solisten.py --show-netns # Show network ns ID (useful for containers)
34"""
35
36parser = argparse.ArgumentParser(
37    description="Stream sockets listen",
38    formatter_class=argparse.RawDescriptionHelpFormatter,
39    epilog=examples)
40parser.add_argument("--show-netns", action="store_true",
41    help="show network namespace")
42parser.add_argument("-p", "--pid", default=0, type=int,
43    help="trace this PID only")
44parser.add_argument("-n", "--netns", default=0, type=int,
45    help="trace this Network Namespace only")
46parser.add_argument("--ebpf", action="store_true",
47    help=argparse.SUPPRESS)
48
49
50# BPF Program
51bpf_text = """
52#include <net/net_namespace.h>
53#include <bcc/proto.h>
54#pragma clang diagnostic push
55#pragma clang diagnostic ignored "-Wenum-conversion"
56#include <net/inet_sock.h>
57#pragma clang diagnostic pop
58
59// Common structure for UDP/TCP IPv4/IPv6
60struct listen_evt_t {
61    u64 ts_us;
62    u64 pid_tgid;
63    u64 backlog;
64    u64 netns;
65    u64 proto;    // familiy << 16 | type
66    u64 lport;    // use only 16 bits
67    u64 laddr[2]; // IPv4: store in laddr[0]
68    char task[TASK_COMM_LEN];
69};
70BPF_PERF_OUTPUT(listen_evt);
71
72// Send an event for each IPv4 listen with PID, bound address and port
73int kprobe__inet_listen(struct pt_regs *ctx, struct socket *sock, int backlog)
74{
75        // cast types. Intermediate cast not needed, kept for readability
76        struct sock *sk = sock->sk;
77        struct inet_sock *inet = (struct inet_sock *)sk;
78
79        // Built event for userland
80        struct listen_evt_t evt = {
81            .ts_us = bpf_ktime_get_ns() / 1000,
82            .backlog = backlog,
83        };
84
85        // Get process comm. Needs LLVM >= 3.7.1
86        // see https://github.com/iovisor/bcc/issues/393
87        bpf_get_current_comm(evt.task, TASK_COMM_LEN);
88
89        // Get socket IP family
90        u16 family = sk->__sk_common.skc_family;
91        evt.proto = family << 16 | SOCK_STREAM;
92
93        // Get PID
94        evt.pid_tgid = bpf_get_current_pid_tgid();
95
96        ##FILTER_PID##
97
98        // Get port
99        evt.lport = inet->inet_sport;
100        evt.lport = ntohs(evt.lport);
101
102        // Get network namespace id, if kernel supports it
103#ifdef CONFIG_NET_NS
104        evt.netns = sk->__sk_common.skc_net.net->ns.inum;
105#else
106        evt.netns = 0;
107#endif
108
109        ##FILTER_NETNS##
110
111        // Get IP
112        if (family == AF_INET) {
113            evt.laddr[0] = inet->inet_rcv_saddr;
114        } else if (family == AF_INET6) {
115            bpf_probe_read(evt.laddr, sizeof(evt.laddr),
116                           sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
117        }
118
119        // Send event to userland
120        listen_evt.perf_submit(ctx, &evt, sizeof(evt));
121
122        return 0;
123};
124"""
125
126# event data
127TASK_COMM_LEN = 16      # linux/sched.h
128class ListenEvt(ct.Structure):
129    _fields_ = [
130        ("ts_us", ct.c_ulonglong),
131        ("pid_tgid", ct.c_ulonglong),
132        ("backlog", ct.c_ulonglong),
133        ("netns", ct.c_ulonglong),
134        ("proto", ct.c_ulonglong),
135        ("lport", ct.c_ulonglong),
136        ("laddr", ct.c_ulonglong * 2),
137        ("task", ct.c_char * TASK_COMM_LEN)
138    ]
139
140    # TODO: properties to unpack protocol / ip / pid / tgid ...
141
142# Format output
143def event_printer(show_netns):
144    def print_event(cpu, data, size):
145        # Decode event
146        event = ct.cast(data, ct.POINTER(ListenEvt)).contents
147
148        pid = event.pid_tgid & 0xffffffff
149        proto_family = event.proto & 0xff
150        proto_type = event.proto >> 16 & 0xff
151
152        if proto_family == SOCK_STREAM:
153            protocol = "TCP"
154        elif proto_family == SOCK_DGRAM:
155            protocol = "UDP"
156        else:
157            protocol = "UNK"
158
159        address = ""
160        if proto_type == AF_INET:
161            protocol += "v4"
162            address = inet_ntop(AF_INET, pack("I", event.laddr[0]))
163        elif proto_type == AF_INET6:
164            address = inet_ntop(AF_INET6, event.laddr)
165            protocol += "v6"
166
167        # Display
168        if show_netns:
169            printb(b"%-6d %-12.12s %-12s %-6s %-8s %-5s %-39s" % (
170                pid, event.task, event.netns, protocol, event.backlog,
171                event.lport, address,
172            ))
173        else:
174            printb(b"%-6d %-12.12s %-6s %-8s %-5s %-39s" % (
175                pid, event.task, protocol, event.backlog,
176                event.lport, address,
177            ))
178
179    return print_event
180
181if __name__ == "__main__":
182    # Parse arguments
183    args = parser.parse_args()
184
185    pid_filter = ""
186    netns_filter = ""
187
188    if args.pid:
189        pid_filter = "if (evt.pid_tgid != %d) return 0;" % args.pid
190    if args.netns:
191        netns_filter = "if (evt.netns != %d) return 0;" % args.netns
192
193    bpf_text = bpf_text.replace("##FILTER_PID##", pid_filter)
194    bpf_text = bpf_text.replace("##FILTER_NETNS##", netns_filter)
195
196    if args.ebpf:
197        print(bpf_text)
198        exit()
199
200    # Initialize BPF
201    b = BPF(text=bpf_text)
202    b["listen_evt"].open_perf_buffer(event_printer(args.show_netns))
203
204    # Print headers
205    if args.show_netns:
206        print("%-6s %-12s %-12s %-6s %-8s %-5s %-39s" %
207              ("PID", "COMM", "NETNS", "PROTO", "BACKLOG", "PORT", "ADDR"))
208    else:
209        print("%-6s %-12s %-6s %-8s %-5s %-39s" %
210              ("PID", "COMM", "PROTO", "BACKLOG", "PORT", "ADDR"))
211
212    # Read events
213    while 1:
214        try:
215            b.perf_buffer_poll()
216        except KeyboardInterrupt:
217            exit()
218