diff --git a/chains/links/transport_meta.py b/chains/links/transport_meta.py index c043842..88ef1d1 100644 --- a/chains/links/transport_meta.py +++ b/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,7 +31,7 @@ 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 @@ -39,7 +39,9 @@ def transport_meta_data(self): @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() diff --git a/chains/links/transport_meta_test.py b/chains/links/transport_meta_test.py new file mode 100644 index 0000000..afecebf --- /dev/null +++ b/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'