Add capture for GnuTLS and argparse options

Now it also captures gnutls_record_send and gnutls_record_recv calls.
Added options to filter by pid or command name, or just one lib.

tools/sslsniff.py

 #!/usr/bin/python
 #
-# sslsniff  Captures data on SSL_READ or SSL_WRITE functions of OpenSSL
+# sslsniff  Captures data on read/recv or write/send functions of OpenSSL and
+#           GnuTLS
 #           For Linux, uses BCC, eBPF.
 #
+# USAGE: sslsniff.py [-h] [-p PID] [-c COMM] [-o] [-g] [-d]
+#
 # Licensed under the Apache License, Version 2.0 (the "License")
 #
 # 12-Aug-2016    Adrian Lopez   Created this.
 # 13-Aug-2016    Mark Drayton   Fix SSL_Read
+# 17-Aug-2016    Adrian Lopez   Capture GnuTLS and add options
+#
 
 from __future__ import print_function
 import ctypes as ct
 from bcc import BPF
+import argparse
+
+# arguments
+examples = """examples:
+    ./sslsniff              # sniff OpenSSL and GnuTLS functions
+    ./sslsniff -p 181       # sniff PID 181 only
+    ./sslsniff -c curl      # sniff curl command only
+    ./sslsniff --no-openssl # don't show OpenSSL calls
+    ./sslsniff --no-gnutls  # don't show GnuTLS calls
+"""
+parser = argparse.ArgumentParser(
+    description="Sniff SSL data",
+    formatter_class=argparse.RawDescriptionHelpFormatter,
+    epilog=examples)
+parser.add_argument("-p", "--pid", help="sniff this PID only.")
+parser.add_argument("-c", "--comm",
+                    help="sniff only commands matching string.")
+parser.add_argument("-o", "--no-openssl", action="store_false", dest="openssl",
+                    help="do not show OpenSSL calls.")
+parser.add_argument("-g", "--no-gnutls", action="store_false", dest="gnutls",
+                    help="do not show GnuTLS calls.")
+parser.add_argument('-d', '--debug', dest='debug', action='count', default=0,
+                    help='debug mode.')
+args = parser.parse_args()
+
 
 prog = """
 #include <linux/ptrace.h>
@@ -27,9 +57,12 @@ struct probe_SSL_data_t {
 BPF_PERF_OUTPUT(perf_SSL_write);
 
 int probe_SSL_write(struct pt_regs *ctx, void *ssl, void *buf, int num) {
+        u32 pid = bpf_get_current_pid_tgid();
+        FILTER
+
         struct probe_SSL_data_t __data = {0};
         __data.timestamp_ns = bpf_ktime_get_ns();
-        __data.pid = bpf_get_current_pid_tgid();
+        __data.pid = pid;
         __data.len = num;
 
         bpf_get_current_comm(&__data.comm, sizeof(__data.comm));
@@ -48,12 +81,16 @@ BPF_HASH(bufs, u32, u64);
 
 int probe_SSL_read_enter(struct pt_regs *ctx, void *ssl, void *buf, int num) {
         u32 pid = bpf_get_current_pid_tgid();
+        FILTER
+
         bufs.update(&pid, (u64*)&buf);
         return 0;
 }
 
 int probe_SSL_read_exit(struct pt_regs *ctx, void *ssl, void *buf, int num) {
         u32 pid = bpf_get_current_pid_tgid();
+        FILTER
+
         u64 *bufp = bufs.lookup(&pid);
         if (bufp == 0) {
                 return 0;
@@ -77,17 +114,34 @@ int probe_SSL_read_exit(struct pt_regs *ctx, void *ssl, void *buf, int num) {
 }
 """
 
-b = BPF(text=prog)
+if args.pid:
+    prog = prog.replace('FILTER', 'if (pid != %s) { return 0; }' % args.pid)
+else:
+    prog = prog.replace('FILTER', '')
+
+if args.debug:
+    print(prog)
+
 
-# Join to ssl_write
-b.attach_uprobe(name="ssl", sym="SSL_write", fn_name="probe_SSL_write")
+b = BPF(text=prog)
 
-# Join to ssl_read
 # It looks like SSL_read's arguments aren't available in a return probe so you
 # need to stash the buffer address in a map on the function entry and read it
 # on its exit (Mark Drayton)
-b.attach_uprobe(name="ssl", sym="SSL_read", fn_name="probe_SSL_read_enter")
-b.attach_uretprobe(name="ssl", sym="SSL_read", fn_name="probe_SSL_read_exit")
+#
+if args.openssl:
+    b.attach_uprobe(name="ssl", sym="SSL_write", fn_name="probe_SSL_write")
+    b.attach_uprobe(name="ssl", sym="SSL_read", fn_name="probe_SSL_read_enter")
+    b.attach_uretprobe(name="ssl", sym="SSL_read",
+                       fn_name="probe_SSL_read_exit")
+
+if args.gnutls:
+    b.attach_uprobe(name="gnutls", sym="gnutls_record_send",
+                    fn_name="probe_SSL_write")
+    b.attach_uprobe(name="gnutls", sym="gnutls_record_recv",
+                    fn_name="probe_SSL_read_enter")
+    b.attach_uretprobe(name="gnutls", sym="gnutls_record_recv",
+                       fn_name="probe_SSL_read_exit")
 
 # define output data structure in Python
 TASK_COMM_LEN = 16  # linux/sched.h
@@ -114,16 +168,22 @@ start = 0
 
 
 def print_event_write(cpu, data, size):
-    print_event(cpu, data, size, "SSL_WRITE")
+    print_event(cpu, data, size, "WRITE/SEND")
 
 
 def print_event_read(cpu, data, size):
-    print_event(cpu, data, size, "SSL_READ")
+    print_event(cpu, data, size, "READ/RECV")
 
 
 def print_event(cpu, data, size, rw):
     global start
     event = ct.cast(data, ct.POINTER(Data)).contents
+
+    # Filter events by command
+    if args.comm:
+        if not args.comm == event.comm:
+            return
+
     if start == 0:
         start = event.timestamp_ns
     time_s = (float(event.timestamp_ns - start)) / 1000000000
@@ -133,11 +193,11 @@ def print_event(cpu, data, size, rw):
     e_mark = "-" * 5 + " END DATA " + "-" * 5
 
     truncated_bytes = event.len - MAX_BUF_SIZE
-    if truncated_bytes > 0 :
+    if truncated_bytes > 0:
         e_mark = "-" * 5 + " END DATA (TRUNCATED, " + str(truncated_bytes) + \
                 " bytes lost) " + "-" * 5
 
-    print("%-12s %-18.9f %-16s %-6d %-6d\n%s\n%s\n%s\n" % (rw, time_s,
+    print("%-12s %-18.9f %-16s %-6d %-6d\n%s\n%s\n%s\n\n" % (rw, time_s,
                                                              event.comm,
                                                              event.pid,
                                                              event.len,

tools/sslsniff_example.txt

@@ -59,3 +59,27 @@ SSL_READ     0.129967972        curl             12915  1270
     div {
         w
 ----- END DATA (TRUNCATED, 798 bytes lost) -----
+
+
+
+
+USAGE message:
+
+usage: sslsniff.py [-h] [-p PID] [-c COMM] [-o] [-g] [-d]
+
+Sniff SSL data
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -p PID, --pid PID     sniff this PID only.
+  -c COMM, --comm COMM  sniff only commands matching string.
+  -o, --no-openssl      do not show OpenSSL calls.
+  -g, --no-gnutls       do not show GnuTLS calls.
+  -d, --debug           debug mode.
+
+examples:
+    ./sslsniff              # sniff OpenSSL and GnuTLS functions
+    ./sslsniff -p 181       # sniff PID 181 only
+    ./sslsniff -c curl      # sniff curl command only
+    ./sslsniff --no-openssl # don't show OpenSSL calls
+    ./sslsniff --no-gnutls  # don't show GnuTLS calls