use unpack_from

Author: Koos85

lib/ipflow/flowset.py

@@ -6,13 +6,14 @@
 
 def on_flowset_template(
     line: bytes,
+    pos: int,
+    pos_end: int,
     source: str,
     source_id: int,
     source_uptime: int,
 ):
-    pos = 0
-    while pos < len(line):
-        template_id, field_count = struct.unpack('>HH', line[pos:pos+4])
+    while pos < pos_end:
+        template_id, field_count = struct.unpack_from('>HH', line, pos)
         # rfc 3954 5.1
         # NetFlow Collectors SHOULD use the combination of the source IP
         # address and the Source ID field to separate different export
@@ -24,7 +25,7 @@ def on_flowset_template(
             continue
 
         fields = [
-            Field(*struct.unpack('>HH', line[i:i+4]))
+            Field(*struct.unpack_from('>HH', line, i))
             for i in range(pos + 4, pos + 4 + field_count * 4, 4)
         ]
 
@@ -40,14 +41,15 @@ def on_flowset_template(
 
 def on_flowset(
     line: bytes,
+    pos: int,
+    pos_end: int,
     flowset_id: int,
     source: str,
     source_id: int,
 ):
     key = source, source_id, flowset_id
     template = flowset_templates.get(key)
     if template:
-        fmt = template._fmt
-        pad = len(line) % template.length
-        for values in struct.iter_unpack(fmt, line[:-pad] if pad else line):
-            yield Flow(template, values)
+        # assume 3 padding
+        for i in range(pos, pos_end - 3, template.length):
+            yield Flow(template, template._fmt.unpack_from(line, i))

lib/ipflow/parser.py

@@ -15,31 +15,30 @@ def on_packet(line: bytes, source: str):
         unix_secs,
         sequence_number,
         source_id,
-    ) = struct.unpack(HEADER_FMT, line[:HEADER_SIZE])
+    ) = struct.unpack_from(HEADER_FMT, line)
 
     # we don't use header.count because it includes the Options Template
     # FlowSets (FlowSet ID 1), which we ignore
     pos = HEADER_SIZE
     while pos + 4 < len(line):
-        flowset_id, length = struct.unpack('>HH', line[pos:pos+4])
+        flowset_id, length = struct.unpack_from('>HH', line, pos)
 
         # prevent endless loop
         if length == 0:
             logging.error('failed to parse packet')
             break
 
-        flowset = line[pos+4:pos+length]
-        pos += length
-
         if flowset_id == 0:
             try:
-                on_flowset_template(flowset, source, source_id, sysuptime)
+                on_flowset_template(line, pos + 4, pos + length, source,
+                                    source_id, sysuptime)
             except Exception:
                 logging.error('failed to parse FlowSet template')
                 break
         elif flowset_id > 255:
             try:
-                for flow in on_flowset(flowset, flowset_id, source, source_id):
+                for flow in on_flowset(line, pos + 4, pos + length, flowset_id,
+                                       source, source_id):
                     yield flow
             except Exception:
                 logging.warning('failed to parse FlowSet')

lib/ipflow/template.py

@@ -1,3 +1,4 @@
+import struct
 from typing import List
 from .field import Field
 
@@ -13,7 +14,7 @@ class DataTemplate:
 
     def __init__(self, fmt: str, length: int, fields: List[Field],
                  index: List[int], source_uptime: int):
-        self._fmt = fmt
+        self._fmt = struct.Struct(fmt)
         self.fields = fields
         self.index = index
         self.length = length

lib/subscription.py

@@ -5,7 +5,7 @@
 
 
 class Subscription(NamedTuple):
-    address: Union[IPv4Address, IPv4Address]
+    address: Union[IPv4Address, IPv6Address]
     result: List[Flow]
     timestamp: int