Guard transport metadata against raw byte payloads

chains/links/transport_meta.py

@@ -3,7 +3,7 @@
 
 # Local imports
 from chains.links import link
-from chains.utils import file_utils, log_utils, data_utils
+from chains.utils import file_utils, data_utils
 
 
 class TransportMeta(link.Link):
@@ -31,15 +31,17 @@ def transport_meta_data(self):
                 item['transport'] = data_utils.make_dict(trans_data)
                 item['transport']['type'] = trans_type
                 item['transport']['flags'] = self._readable_flags(item['transport'])
-                item['transport']['data'] = trans_data['data']
+                item['transport']['data'] = getattr(trans_data, 'data', None)
 
             # All done
             yield item
 
     @staticmethod
     def _get_transport_type(transport):
         """Give the transport as a string or None if not one"""
-        return transport.__class__.__name__  if transport.__class__.__name__ != 'str' else None
+        if not transport or isinstance(transport, (bytes, bytearray, str)):
+            return None
+        return transport.__class__.__name__
 
     @staticmethod
     def _readable_flags(transport):
@@ -70,6 +72,7 @@ def _readable_flags(transport):
             _flag_list.append('psh')
         return _flag_list
 
+
 def test():
     """Test for TransportMeta class"""
     import pprint
@@ -94,5 +97,6 @@ def test():
     for item in tmeta.output_stream:
         pprint.pprint(item)
 
+
 if __name__ == '__main__':
     test()

chains/links/transport_meta_test.py

@@ -0,0 +1,31 @@
+import dpkt
+
+from chains.links.transport_meta import TransportMeta
+
+
+def _packet(transport):
+    return {
+        'packet': {'data': transport},
+        'transport': None,
+    }
+
+
+def test_raw_bytes_are_left_unclassified():
+    meta = TransportMeta()
+    meta.input_stream = [_packet(b'\x00\x01\x02')]
+
+    result = next(meta.output_stream)
+
+    assert result['transport'] is None
+
+
+def test_tcp_transport_metadata_is_preserved():
+    tcp = dpkt.tcp.TCP(sport=12345, dport=80, flags=dpkt.tcp.TH_SYN, data=b'hello')
+    meta = TransportMeta()
+    meta.input_stream = [_packet(tcp)]
+
+    result = next(meta.output_stream)
+
+    assert result['transport']['type'] == 'TCP'
+    assert result['transport']['flags'] == ['syn']
+    assert result['transport']['data'] == b'hello'