From 4101d2e388926c414c2542dc0e4330b958935058 Mon Sep 17 00:00:00 2001
From: Brenden Blanco <bblanco@plumgrid.com>
Date: Fri, 21 Aug 2015 16:33:02 -0700
Subject: [PATCH] Add BPF python helpers for reading trace_pipe

Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>
---
 examples/hello_world.py    |  5 +---
 src/python/bpf/__init__.py | 50 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/examples/hello_world.py b/examples/hello_world.py
index e46b2ca4c74e..73a105ed25fa 100755
--- a/examples/hello_world.py
+++ b/examples/hello_world.py
@@ -16,7 +16,4 @@
 """
 b = BPF(text=prog)
 b.attach_kprobe(event="sys_clone", fn_name="hello")
-try:
-    call(["cat", "/sys/kernel/debug/tracing/trace_pipe"])
-except KeyboardInterrupt:
-    pass
+b.trace_print()
diff --git a/src/python/bpf/__init__.py b/src/python/bpf/__init__.py
index 3233b07a451e..9a5831383047 100644
--- a/src/python/bpf/__init__.py
+++ b/src/python/bpf/__init__.py
@@ -15,6 +15,7 @@
 import atexit
 from collections import MutableMapping
 import ctypes as ct
+import fcntl
 import json
 import os
 import sys
@@ -78,11 +79,14 @@
 lib.bpf_prog_load.argtypes = [ct.c_int, ct.c_void_p, ct.c_size_t,
         ct.c_char_p, ct.c_uint, ct.c_char_p, ct.c_uint]
 lib.bpf_attach_kprobe.restype = ct.c_int
-lib.bpf_attach_kprobe.argtypes = [ct.c_int, ct.c_char_p, ct.c_char_p, ct.c_int, ct.c_int, ct.c_int]
+lib.bpf_attach_kprobe.argtypes = [ct.c_int, ct.c_char_p,
+        ct.c_char_p, ct.c_int, ct.c_int, ct.c_int]
 lib.bpf_detach_kprobe.restype = ct.c_int
 lib.bpf_detach_kprobe.argtypes = [ct.c_char_p]
 
 open_kprobes = {}
+tracefile = None
+TRACEFS = "/sys/kernel/debug/tracing"
 
 @atexit.register
 def cleanup_kprobes():
@@ -90,6 +94,8 @@ def cleanup_kprobes():
         os.close(v)
         desc = "-:kprobes/%s" % k
         lib.bpf_detach_kprobe(desc.encode("ascii"))
+    if tracefile:
+        tracefile.close()
 
 class BPF(object):
     SOCKET_FILTER = 1
@@ -389,3 +395,45 @@ def detach_kretprobe(event):
             raise Exception("Failed to detach BPF from kprobe")
         del open_kprobes[ev_name]
 
+    @staticmethod
+    def trace_open(nonblocking=False):
+        """trace_open(nonblocking=False)
+
+        Open the trace_pipe if not already open
+        """
+
+        global tracefile
+        if not tracefile:
+            tracefile = open("%s/trace_pipe" % TRACEFS)
+            if nonblocking:
+                fd = trace.fileno()
+                fl = fcntl.fcntl(fd, fcntl.F_GETFL)
+                fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
+        return tracefile
+
+    @staticmethod
+    def trace_readline(nonblocking=True):
+        """trace_readline(nonblocking=True)
+
+        Read from the kernel debug trace pipe and return one line
+        If nonblocking is False, this will block until ctrl-C is pressed.
+        """
+
+        trace = BPF.trace_open(nonblocking)
+
+        line = None
+        try:
+            line = trace.readline(128).rstrip()
+        except BlockingIOError:
+            pass
+        return line
+
+    @staticmethod
+    def trace_print():
+        try:
+            while True:
+                line = BPF.trace_readline(nonblocking=False)
+                print(line)
+                sys.stdout.flush()
+        except KeyboardInterrupt:
+            exit()