From a7a27102677a7799591f65c883c5e614bbbe3e9b Mon Sep 17 00:00:00 2001
From: Hitoshi Irino <irino@sfc.wide.ad.jp>
Date: Fri, 23 Feb 2024 13:13:33 +0900
Subject: [PATCH] Fix V9DataFlowSet to avoid following struct.unpack error:

Exception in thread Thread-2:
Traceback (most recent call last):
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/python-netflow-v9-softflowd/netflow/collector.py", line 129, in run
    export = parse_packet(pkt.data, templates)
  File "/python-netflow-v9-softflowd/netflow/utils.py", line 96, in parse_packet
    return V9ExportPacket(data, templates["netflow"])
  File "/python-netflow-v9-softflowd/netflow/v9.py", line 558, in __init__
    dfs = V9DataFlowSet(data[offset:], matched_template)
  File "/python-netflow-v9-softflowd/netflow/v9.py", line 231, in __init__
    unpacked_values = struct.unpack(struct_format, data[offset:offset + struct_len])
---
 netflow/v9.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/netflow/v9.py b/netflow/v9.py
index aee63db..14cb78f 100644
--- a/netflow/v9.py
+++ b/netflow/v9.py
@@ -205,7 +205,7 @@ def __init__(self, data, template):
         offset = 4
 
         # As the field lengths are variable V9 has padding to next 32 Bit
-        padding_size = 4 - (self.length % 4)  # 4 Byte
+        padding_size = 4 - (self.length % 4) if self.length % 4 else 0 # 4 Byte
 
         # For performance reasons, we use struct.unpack to get individual values. Here
         # we prepare the format string for parsing it. The format string is based on the template fields and their
@@ -225,10 +225,16 @@ def __init__(self, data, template):
                 struct_format += '%ds' % flen
             struct_len += flen
 
-        while offset <= (self.length - padding_size):
+        while offset + struct_len <= (self.length - padding_size):
             # Here we actually unpack the values, the struct format string is used in every data record
             # iteration, until the final offset reaches the end of the whole data stream
-            unpacked_values = struct.unpack(struct_format, data[offset:offset + struct_len])
+            try:
+                unpacked_values = struct.unpack(struct_format, data[offset:offset + struct_len])
+            except:
+                print("struct.unpack is failed. (offset: {}, struct_len: {}, length: {}, padding_size: {})"
+                      .format(offset, struct_len, self.length, padding_size))
+                offset += struct_len
+                continue
 
             new_record = V9DataRecord()
             for field, value in zip(template.fields, unpacked_values):