From 809f599649364670e3e069032668297083905ace Mon Sep 17 00:00:00 2001
From: J <dev195@users.noreply.github.com>
Date: Mon, 29 Apr 2024 11:39:55 -0400
Subject: [PATCH] Updated decode.py with new connection management arguments
 and a way of handling packet timeouts when listening live on an interface

---
 dshell/decode.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/dshell/decode.py b/dshell/decode.py
index 62bd364..8b37be1 100755
--- a/dshell/decode.py
+++ b/dshell/decode.py
@@ -32,6 +32,7 @@
 import tempfile
 import zipfile
 from collections import OrderedDict
+from datetime import timedelta
 from getpass import getpass
 from glob import glob
 from importlib import import_module
@@ -250,6 +251,17 @@ def main(plugin_args=None, **kwargs):
             plugin.out.extra = True
             plugin.out.set_format(plugin.out.format)
 
+        # Set some attributes for ConnectionPlugins
+        if hasattr(plugin, "timeout"):
+            # Set wait time since last packet arrived in a connection before
+            # considering connection closed
+            if t := kwargs.get("conntimeout"):
+                td = timedelta(seconds=int(t))
+                plugin.timeout = td
+            # Set max number of allowed open connections
+            if t := kwargs.get("connmax"):
+                plugin.max_open_connections = int(t)
+
         # Set the BPF filters
         # Each plugin has its own default BPF that will be extended or replaced
         # based on --no-vlan, --ebpf, or --bpf arguments.
@@ -447,8 +459,12 @@ def read_packets(input: str, interface=False, bpf=None, count=None) -> Iterable[
         try:
             header, packet_data = capture.next()
             if header is None and not packet_data:
-                # probably the end of the capture
-                break
+                if not interface:
+                    # probably the end of the capture
+                    break
+                else:
+                    # interface timed out
+                    continue
             if count and frame - 1 >= count:
                 # we've reached the maximum number of packets to process
                 break
@@ -464,6 +480,11 @@ def read_packets(input: str, interface=False, bpf=None, count=None) -> Iterable[
 
             yield packet
 
+        # handle SIGINT gracefully, break read loop and allow shutdown
+        except KeyboardInterrupt:
+            logger.debug("Caught KeyboardInterrupt or SIGINT. Closing capture.")
+            break
+
         except pcapy.PcapError as e:
             estr = str(e)
             eformat = "Error processing '{i}' - {e}"
@@ -557,6 +578,12 @@ def main_command_line():
     parser.add_argument('--unzipdir', type=str, metavar="DIRECTORY",
                       default=tempfile.gettempdir(),
                       help='Directory to use when decompressing input files (.gz, .bz2, and .zip only)')
+    parser.add_argument('--conn-timeout', dest="conntimeout", type=int,
+                      metavar="SECONDS", default=3600,
+                      help="Number of seconds to wait after last packet in a connection before closing it (default: 3600)")
+    parser.add_argument('--conn-max-open', dest='connmax', type=int,
+                      metavar="NUMBER", default=1000,
+                      help="Number of connections to hold in an open state before Dshell begins closing the oldest (default: 1000)")
 
     multiprocess_group = parser.add_argument_group("multiprocessing arguments")
     multiprocess_group.add_argument('-P', '--parallel', dest='multiprocessing', action='store_true',