From d8747a7b810421d4a1c5aa88e85012558890e577 Mon Sep 17 00:00:00 2001
From: Sander Vrijders <sander.vrijders@ugent.be>
Date: Mon, 26 Mar 2018 12:11:32 +0200
Subject: storyboard: Add capturing of traffic

This adds the functionality to capture traffic on certain interfaces
so that it can be inspected with tools like wireshark. A user needs to
pass a start and end time and a node and DIF. Rumba will then
determine the correct interface to capture on.

Implements #41
---
 rumba/model.py         |  3 ++-
 rumba/ssh_support.py   |  8 +++++---
 rumba/storyboard.py    | 34 ++++++++++++++++++++++++++++++++++
 rumba/testbeds/qemu.py |  2 ++
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/rumba/model.py b/rumba/model.py
index e12bddd..0ec301a 100644
--- a/rumba/model.py
+++ b/rumba/model.py
@@ -243,6 +243,7 @@ class Node(object):
         self.ssh_config = SSHConfig(name)
         self.ipcps = []
         self.policies = dict()
+        self.has_tcpdump = False
         if policies is None:
             policies = dict()
         for dif in self.difs:
@@ -898,4 +899,4 @@ class Executor:
 
     def fetch_files(self, node, paths, destination, sudo=False):
         for path in paths:
-            self.fetch_file(node, path, destination, sudo)
\ No newline at end of file
+            self.fetch_file(node, path, destination, sudo)
diff --git a/rumba/ssh_support.py b/rumba/ssh_support.py
index b1492e7..a9dff28 100644
--- a/rumba/ssh_support.py
+++ b/rumba/ssh_support.py
@@ -336,8 +336,10 @@ def copy_files_from_testbed(testbed, ssh_config, paths,
         destination = destination + '/'
 
     if sudo:
-        execute_command(testbed, ssh_config,
-                        'sudo chmod a+rw %s' % (" ".join(paths)))
+        cmd = 'chmod a+rw %s' % (" ".join(paths))
+        if ssh_config.username != 'root':
+            cmd = "sudo %s" % command
+        execute_command(testbed, ssh_config, cmd)
 
     if ssh_config.client is None:
         client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
@@ -436,4 +438,4 @@ def aptitude_install(testbed, node, packages):
             "while ! " + sudo("apt-get update") + "; do sleep 1; done",
             "while ! " + sudo(package_install) + "; do sleep 1; done"]
 
-    execute_proxy_commands(testbed, node.ssh_config, cmds, time_out=None)
\ No newline at end of file
+    execute_proxy_commands(testbed, node.ssh_config, cmds, time_out=None)
diff --git a/rumba/storyboard.py b/rumba/storyboard.py
index ea0c51e..61dd8cc 100644
--- a/rumba/storyboard.py
+++ b/rumba/storyboard.py
@@ -34,6 +34,7 @@ import math
 import os
 import random
 import time
+import uuid
 
 import rumba.model as model
 import rumba.ssh_support as ssh_support
@@ -821,6 +822,39 @@ class StoryBoard(SBEntity):
         buffer = StringIO(string)
         self.parse_script(buffer, clean)
 
+    def capture_traffic(self, start, end, node, dif):
+        """
+        Captures the traffic of an interface on a certain node.
+
+        :param start: The time to start capturing.
+        :param end: The time to stop capturing.
+        :param node: The node to capture on.
+        :param dif: The Shim Ethernet DIF of the node, Rumba
+                    automatically resolves the correct interface.
+        """
+        for ipcp in dif.ipcps:
+            if ipcp.node is not node:
+                continue
+            # In case tcpdump is not present, this assumes a testbed
+            # with Ubuntu/Debian just like the rest of installation
+            if not node.has_tcpdump:
+                ssh_support.aptitude_install(self.experiment.testbed,
+                                             node, ["tcpdump"])
+                node.has_tcpdump = True
+
+            # Create random string
+            pcap_file = node.name + '_' + dif.name + '_' + \
+                        str(uuid.uuid4())[0:4] + ".pcap"
+
+            tcpd_client = Client(ap="tcpdump", options="-i %s -w %s" \
+                                 % (ipcp.ifname, pcap_file))
+            duration = end - start
+            cb = functools.partial(node.fetch_file, pcap_file,
+                                   self.experiment.log_dir, sudo=True)
+            action = functools.partial(self.run_client, tcpd_client,
+                                       duration=duration, node=node,
+                                       callback = cb)
+            self._script.add_event(Event(action, ev_time=start))
 
 class Event(object):
 
diff --git a/rumba/testbeds/qemu.py b/rumba/testbeds/qemu.py
index f3daefd..df80b7e 100644
--- a/rumba/testbeds/qemu.py
+++ b/rumba/testbeds/qemu.py
@@ -230,6 +230,8 @@ class Testbed(mod.Testbed):
         #  End of shim/node parsing block
         ##
         #
+        for node in experiment.nodes:
+            node.has_tcpdump = True
 
         def executor(list_of_commands):
             for cmd in list_of_commands:
-- 
cgit v1.2.3