Merge pull request #27 from guedou/matplotlib

Matplotlib support

scapy/arch/__init__.py

@@ -13,11 +13,20 @@ from scapy.error import *
 import scapy.config
 
 try:
-    import Gnuplot
-    GNUPLOT=1
+    from matplotlib import get_backend as matplotlib_get_backend
+    import matplotlib.pyplot as plt
+    MATPLOTLIB = 1
+    if "inline" in matplotlib_get_backend():
+        MATPLOTLIB_INLINED = 1
+    else:
+        MATPLOTLIB_INLINED = 0
+    MATPLOTLIB_DEFAULT_PLOT_KARGS = { "color": "r", "marker": "+", "ls": "" }
 except ImportError:
-    log_loading.info("Can't import python gnuplot wrapper . Won't be able to plot.")
-    GNUPLOT=0
+    plt = None
+    MATPLOTLIB = 0
+    MATPLOTLIB_INLINED = 0
+    MATPLOTLIB_DEFAULT_PLOT_KARGS = dict()
+    log_loading.info("Can't import matplotlib. Won't be able to plot.")
 
 try:
     import pyx

scapy/layers/inet.py

@@ -22,6 +22,7 @@ from scapy.automaton import Automaton,ATMT
 
 import scapy.as_resolvers
 
+from scapy.arch import plt, MATPLOTLIB_INLINED, MATPLOTLIB_DEFAULT_PLOT_KARGS
 
 ####################
 ## IP Tools class ##
@@ -915,21 +916,49 @@ def defragment(plist):
 ### Add timeskew_graph() method to PacketList
 def _packetlist_timeskew_graph(self, ip, **kargs):
     """Tries to graph the timeskew between the timestamps and real time for a given ip"""
+
+    # Filter TCP segments which source address is 'ip'
     res = map(lambda x: self._elt2pkt(x), self.res)
     b = filter(lambda x:x.haslayer(IP) and x.getlayer(IP).src == ip and x.haslayer(TCP), res)
+
+    # Build a list of tuples (creation_time, replied_timestamp)
     c = []
     for p in b:
         opts = p.getlayer(TCP).options
         for o in opts:
             if o[0] == "Timestamp":
                 c.append((p.time,o[1][0]))
+
+    # Stop if the list is empty
     if not c:
         warning("No timestamps found in packet list")
-        return
-    d = map(lambda (x,y): (x%2000,((x-c[0][0])-((y-c[0][1])/1000.0))),c)
-    g = Gnuplot.Gnuplot()
-    g.plot(Gnuplot.Data(d,**kargs))
-    return g
+        return []
+
+    # Prepare the data that will be plotted
+    first_creation_time = c[0][0]
+    first_replied_timestamp = c[0][1]
+
+    def _wrap_data(ts_tuple, wrap_seconds=2000):
+        """Wrap the list of tuples."""
+
+        ct,rt = ts_tuple # (creation_time, replied_timestamp)
+        X = ct % wrap_seconds
+        Y = ((ct-first_creation_time) - ((rt-first_replied_timestamp)/1000.0))
+
+        return X, Y
+
+    data = map(_wrap_data, c)
+
+    # Mimic the default gnuplot output
+    if kargs == {}:
+        kargs = MATPLOTLIB_DEFAULT_PLOT_KARGS
+    lines = plt.plot(data, **kargs)
+
+    # Call show() if matplotlib is not inlined
+    if not MATPLOTLIB_INLINED:
+        plt.show()
+
+    return lines
 
 PacketList.timeskew_graph = new.instancemethod(_packetlist_timeskew_graph, None, PacketList)
 

scapy/plist.py

@@ -15,9 +15,7 @@ from collections import defaultdict
 
 from utils import do_graph,hexdump,make_table,make_lined_table,make_tex_table,get_temp_file
 
-import arch
-if arch.GNUPLOT:
-    Gnuplot=arch.Gnuplot
+from scapy.arch import plt, MATPLOTLIB_INLINED, MATPLOTLIB_DEFAULT_PLOT_KARGS
 
 
 
@@ -137,48 +135,97 @@ lfilter: truth function to apply to each packet to decide whether it will be dis
         return make_tex_table(self.res, *args, **kargs)
 
     def plot(self, f, lfilter=None, **kargs):
-        """Applies a function to each packet to get a value that will be plotted with GnuPlot. A gnuplot object is returned
-        lfilter: a truth function that decides whether a packet must be ploted"""
-        g=Gnuplot.Gnuplot()
+        """Applies a function to each packet to get a value that will be plotted
+        with matplotlib. A list of matplotlib.lines.Line2D is returned.
+
+        lfilter: a truth function that decides whether a packet must be ploted
+        """
+
+        # Get the list of packets
         l = self.res
+
+        # Apply the filter
         if lfilter is not None:
             l = filter(lfilter, l)
+
+        # Apply the function f to the packets
         l = map(f, l)
-        g.plot(Gnuplot.Data(l, **kargs))
-        return g
+
+        # Mimic the default gnuplot output
+        if kargs == {}:
+            kargs = MATPLOTLIB_DEFAULT_PLOT_KARGS
+        lines = plt.plot(l, **kargs)
+
+        # Call show() if matplotlib is not inlined
+        if not MATPLOTLIB_INLINED:
+            plt.show()
+
+        return lines
 
     def diffplot(self, f, delay=1, lfilter=None, **kargs):
         """diffplot(f, delay=1, lfilter=None)
-        Applies a function to couples (l[i],l[i+delay])"""
-        g = Gnuplot.Gnuplot()
+        Applies a function to couples (l[i],l[i+delay])
+
+        A list of matplotlib.lines.Line2D is returned.
+        """
+
+        # Get the list of packets
         l = self.res
+
+        # Apply the filter
         if lfilter is not None:
             l = filter(lfilter, l)
+
+        # Apply the function f to compute the difference
         l = map(f, l[:-delay],l[delay:])
-        g.plot(Gnuplot.Data(l, **kargs))
-        return g
+
+        # Mimic the default gnuplot output
+        if kargs == {}:
+            kargs = MATPLOTLIB_DEFAULT_PLOT_KARGS
+        lines = plt.plot(l, **kargs)
+
+        # Call show() if matplotlib is not inlined
+        if not MATPLOTLIB_INLINED:
+            plt.show()
+
+        return lines
 
     def multiplot(self, f, lfilter=None, **kargs):
-        """Uses a function that returns a label and a value for this label, then plots all the values label by label"""
-        g=Gnuplot.Gnuplot()
+        """Uses a function that returns a label and a value for this label, then
+        plots all the values label by label.
+
+        A list of matplotlib.lines.Line2D is returned.
+        """
+
+        # Get the list of packets
         l = self.res
+
+        # Apply the filter
         if lfilter is not None:
             l = filter(lfilter, l)
 
-        d={}
-        for e in l:
-            k,v = f(e)
-            if k in d:
-                d[k].append(v)
-            else:
-                d[k] = [v]
-        data=[]
-        for k in d:
-            data.append(Gnuplot.Data(d[k], title=k, **kargs))
+        # Apply the function f to the packets
+        d_x = {}
+        d_y = {}
+        for k, v in map(f, l):
+            x,y = v
+            d_x[k] = d_x.get(k, []) + [x]
+            d_y[k] = d_y.get(k, []) + [y]
+
+        # Mimic the default gnuplot output
+        if kargs == {}:
+            kargs = MATPLOTLIB_DEFAULT_PLOT_KARGS
+
+        lines = []
+        for k in d_x:
+            lines += plt.plot(d_x[k], d_y[k], label=k, **kargs)
+        plt.legend(loc="center right", bbox_to_anchor=(1.5, 0.5))
 
-        g.plot(*data)
-        return g
+        # Call show() if matplotlib is not inlined
+        if not MATPLOTLIB_INLINED:
+            plt.show()
 
+        return lines
 
     def rawhexdump(self):
         """Prints an hexadecimal dump of each packet in the list"""