packet lib: Simpler API to assemble a packet

For example:

(Current)
e = ethernet.ethernet(…)
i = ipv4.ipv4(…)
u = udp.udp(…)
pkt = packet.Packet()
pkt.add_protocols(e)
pkt.add_protocols(i)
pkt.add_protocols(u)

(New)
e = ethernet.ethernet(…)
i = ipv4.ipv4(…)
u = udp.udp(…)
pkt = e/i/u

Signed-off-by: Satoshi Kobayashi <satoshi-k@stratosphere.co.jp>
Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

2 files changed, 26 insertions, 0 deletions

diff --git a/ryu/lib/packet/packet.py b/ryu/lib/packet/packet.py
index 785a28be..f21edbe2 100644
--- a/ryu/lib/packet/packet.py
+++ b/ryu/lib/packet/packet.py
@@ -107,6 +107,10 @@ class Packet(object):
         self.protocol_idx += 1
         return p
 
+    def __div__(self, trailer):
+        self.add_protocol(trailer)
+        return self
+
     def __iter__(self):
         return self
 
@@ -131,3 +135,13 @@ class Packet(object):
     def __str__(self):
         return ', '.join(repr(protocol) for protocol in self.protocols)
     __repr__ = __str__  # note: str(list) uses __repr__ for elements
+
+
+# XXX: Hack for preventing recursive import
+def _PacketBase__div__(self, trailer):
+    pkt = Packet()
+    pkt.add_protocol(self)
+    pkt.add_protocol(trailer)
+    return pkt
+
+packet_base.PacketBase.__div__ = _PacketBase__div__
diff --git a/ryu/tests/unit/packet/test_packet.py b/ryu/tests/unit/packet/test_packet.py
index 8e70f461..03e036dc 100644
--- a/ryu/tests/unit/packet/test_packet.py
+++ b/ryu/tests/unit/packet/test_packet.py
@@ -41,6 +41,8 @@ class TestPacket(unittest.TestCase):
     dst_ip = '192.168.128.10'
     src_ip = '192.168.122.20'
     dst_ip_bin = addrconv.ipv4.text_to_bin(dst_ip)
+    src_port = 50001
+    dst_port = 50002
     src_ip_bin = addrconv.ipv4.text_to_bin(src_ip)
     payload = '\x06\x06\x47\x50\x00\x00\x00\x00' \
         + '\xcd\xc5\x00\x00\x00\x00\x00\x00' \
@@ -712,3 +714,13 @@ class TestPacket(unittest.TestCase):
 
         eq_(pkt_str, str(pkt))
         eq_(pkt_str, repr(pkt))
+
+    def test_div_api(self):
+        e = ethernet.ethernet(self.dst_mac, self.src_mac, ether.ETH_TYPE_IP)
+        i = ipv4.ipv4()
+        u = udp.udp(self.src_port, self.dst_port)
+        pkt = e/i/u
+        ok_(isinstance(pkt, packet.Packet))
+        ok_(isinstance(pkt.protocols[0], ethernet.ethernet))
+        ok_(isinstance(pkt.protocols[1], ipv4.ipv4))
+        ok_(isinstance(pkt.protocols[2], udp.udp))