[tor-commits] [ooni-probe/master] Properly implement timeout in txscapy

[isis at torproject.org]

commit b267f3bd882fb5da6aa20c9a106ff3a10630e4f9
Author: Arturo Filastò <arturo at filasto.net>
Date:   Sat Sep 29 11:39:02 2012 +0000

    Properly implement timeout in txscapy
    * Fix the example scapy test and scapy template
---
 nettests/example_scapyt.py |    3 ++-
 ooni/lib/txscapy.py        |   35 ++++++++++++++++++++++++++++-------
 ooni/templates/scapyt.py   |    2 +-
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/nettests/example_scapyt.py b/nettests/example_scapyt.py
index 2e73150..c7f342d 100644
--- a/nettests/example_scapyt.py
+++ b/nettests/example_scapyt.py
@@ -10,5 +10,6 @@ class ExampleScapy(scapyt.ScapyTest):
     author = "Arturo Filastò"
     version = 0.1
 
-    inputs = [IP(dst='8.8.8.8')/UDP(), IP()/TCP()]
+    inputs = [IP(dst="8.8.8.8")/TCP(dport=31337),
+            IP(dst="ooni.nu")/TCP(dport=31337)]
 
diff --git a/ooni/lib/txscapy.py b/ooni/lib/txscapy.py
index f67daa3..ec605e5 100644
--- a/ooni/lib/txscapy.py
+++ b/ooni/lib/txscapy.py
@@ -36,8 +36,8 @@ from scapy.all import RawPcapWriter, MTU, BasePacketList, conf
 class PcapWriter(RawPcapWriter):
     def __init__(self, filename, linktype=None, gz=False, endianness="",
                  append=False, sync=False):
-        RawPcapWriter.__init__(self, filename, linktype=None, gz=False,
-                               endianness="", append=False, sync=False)
+        RawPcapWriter.__init__(self, filename, linktype=linktype, gz=gz,
+                               endianness=endianness, append=append, sync=sync)
         fdesc.setNonBlocking(self.f)
 
     def _write_header(self, pkt):
@@ -90,7 +90,7 @@ class Scapy(object):
     recv = False
 
     def __init__(self, pkts=None, maxPacketSize=8192, reactor=None, filter=None,
-            iface=None, nofilter=None, pcapfile=None, *arg, **kw):
+            iface=None, nofilter=None, pcapfile=None, timeout=None, *arg, **kw):
         if self.debug:
             log.startLogging(sys.stdout)
 
@@ -102,7 +102,11 @@ class Scapy(object):
 
         if pkts:
             self._buildPacketQueues(pkts)
-            self._buildSocket()
+            try:
+                self._buildSocket()
+            except Exception, e:
+                log.err("Unable to build socket. Are you root?")
+                sys.exit()
 
         self.cthreads = 0
         self.mthreads = 80
@@ -122,6 +126,9 @@ class Scapy(object):
         if pcapfile:
             self.pcapwriter = PcapWriter(pcapfile)
 
+        if timeout and self.recv:
+            pass
+
     def _buildSocket(self, filter=None, iface=None, nofilter=None):
         self.socket = ScapySocket(filter, iface, nofilter)
         if self.recv:
@@ -213,6 +220,7 @@ class Scapy(object):
             return
         if self.pcapwriter and not self.write_only_answers:
             self.pcapwriter.write(pkt)
+
         self.processPacket(pkt)
 
         h = pkt.hashret()
@@ -251,7 +259,13 @@ class Scapy(object):
         """
         self.socket.send(pkt)
 
-    def sr(self, pkts, filter=None, iface=None, nofilter=0, *args, **kw):
+    def timeout(self, *arg, **kw):
+        if not self.done:
+            self._reactor.callLater(self.timeoutSeconds, self.timeout, None)
+        else:
+            self.deferred.callback(None)
+
+    def sr(self, pkts, filter=None, iface=None, nofilter=0, timeout=None, *args, **kw):
         """
         Wraps the scapy sr function.
 
@@ -271,6 +285,9 @@ class Scapy(object):
         @param iface:    listen answers only on the given interface
         """
         self.recv = True
+        if timeout:
+            self.timeoutSeconds = timeout
+            self._reactor.callLater(timeout, self.timeout, None)
         self._sendrcv(pkts, filter=filter, iface=iface, nofilter=nofilter)
 
     def send(self, pkts, filter=None, iface=None, nofilter=0, *args, **kw):
@@ -338,12 +355,16 @@ class Scapy(object):
         self.threadpool.stop()
         self.running = False
 
+ at defer.inlineCallbacks
 def txsr(*args, **kw):
     tr = Scapy(*args, **kw)
     tr.sr(*args, **kw)
-    return tr.deferred
+    yield tr.deferred
+    tr.finalClose()
 
+ at defer.inlineCallbacks
 def txsend(*arg, **kw):
     tr = Scapy(*arg, **kw)
     tr.send(*arg, **kw)
-    return tr.deferred
+    yield tr.deferred
+    tr.finalClose()
diff --git a/ooni/templates/scapyt.py b/ooni/templates/scapyt.py
index f762ae5..ac32ad8 100644
--- a/ooni/templates/scapyt.py
+++ b/ooni/templates/scapyt.py
@@ -31,7 +31,7 @@ class ScapyTest(TestCase):
 
     receive = True
     timeout = 1
-    pcapfile = 'scapytest.pcap'
+    pcapfile = '/home/x/code/networking/ooni-probe/scapytest.pcap'
     input = IP()/TCP()
     reactor = None
     def setUp(self):