From 81d57f3c4ce33a2ce0848c39d701a4df459dbcc5 Mon Sep 17 00:00:00 2001
From: Dominik Pataky <software+commits@dpataky.eu>
Date: Sat, 1 Aug 2020 10:46:35 +0200
Subject: [PATCH] Handle SIGINT and SIGTERM in yielding listener

Signals INT and TERM were not correctly handled in the 'while True' loop
of the yielding listener function. Now, the loop breaks as expected,
terminating the listener thread and the application.
---
 netflow/collector.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/netflow/collector.py b/netflow/collector.py
index fd71bd4..0cc9538 100644
--- a/netflow/collector.py
+++ b/netflow/collector.py
@@ -12,6 +12,7 @@ import gzip
 import json
 import logging
 import queue
+import signal
 import socket
 import socketserver
 import threading
@@ -113,6 +114,7 @@ class ThreadedNetFlowListener(threading.Thread):
     def run(self):
         # Process packets from the queue
         try:
+            # TODO: use per-client templates
             templates = {"netflow": {}, "ipfix": {}}
             to_retry = []
             while not self._shutdown.is_set():
@@ -171,11 +173,20 @@ class ThreadedNetFlowListener(threading.Thread):
 def get_export_packets(host: str, port: int) -> ParsedPacket:
     """A threaded generator that will yield ExportPacket objects until it is killed
     """
+    def handle_signal(s, f):
+        logger.debug("Received signal {}, raising StopIteration".format(s))
+        raise StopIteration
+    signal.signal(signal.SIGTERM, handle_signal)
+    signal.signal(signal.SIGINT, handle_signal)
+
     listener = ThreadedNetFlowListener(host, port)
     listener.start()
+
     try:
         while True:
             yield listener.get()
+    except StopIteration:
+        pass
     finally:
         listener.stop()
         listener.join()