From 590869d07eea77aed3bb74ba562e2dba1414bf22 Mon Sep 17 00:00:00 2001
From: Timothy Noel <tnoelhere@gmail.com>
Date: Mon, 22 Jan 2024 13:38:12 -0500
Subject: [PATCH] fix issue #49

---
 python/ja4ssh.py | 40 ++++++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/python/ja4ssh.py b/python/ja4ssh.py
index 0567114..20783e2 100644
--- a/python/ja4ssh.py
+++ b/python/ja4ssh.py
@@ -42,22 +42,15 @@ def process_extra_parameters(entry, x, direction):
 ## we return 1 whenever a new stats entry is added based on the sample rate
 ## This way the caller can print this packet out
 def update_ssh_entry(entry, x, ssh_sample_count, debug_stream=None):
-    if (entry['count'] % ssh_sample_count) == 0:
-        to_ja4ssh(entry) if entry['count'] != 0 else None
+
+    if entry['count'] == 0 and len(entry['stats']) == 0:
         entry['stats'].append(dict(ja4sh_stats))
 
-        if debug_stream and int(x['stream']) == debug_stream:
-            if entry['count'] != 0:
-                idx = len(entry['stats']) - 1
-                try:
-                    computed = entry[f'JA4SSH.{idx}']
-                    print (f'computed JA4SSH.{idx}: {computed}')
-                except Exception as e:
-                    pass
-            
-    entry['count'] += 1
+    # Only count SSH PSHACK packets
+    if 'ssh' in x['protos']:
+        entry['count'] += 1
+
 
-    # Now we update the payload lengths, and the acks
     e = entry['stats'][-1]
     direction = 'client' if entry['src'] == x['src'] else 'server'
 
@@ -73,12 +66,31 @@ def update_ssh_entry(entry, x, ssh_sample_count, debug_stream=None):
     if 'ssh' in x['protos']:
         process_extra_parameters(entry, x, direction)
 
+    if x['stream'] == debug_stream:
+        print (f"stats[{len(entry['stats'])}]:tcp flag = {x['flags']}, c{e['client_packets']}s{e['server_packets']}_c{e['client_acks']}s{e['server_acks']}")
+
+    if (entry['count'] % ssh_sample_count) == 0:
+        to_ja4ssh(entry) if entry['count'] != 0 else None
+        if (entry['count'] / ssh_sample_count) == len(entry['stats']):
+            print (f'adding new entry at count {entry["count"]}')
+            entry['stats'].append(dict(ja4sh_stats))
+
+        if debug_stream and int(x['stream']) == debug_stream:
+            if entry['count'] != 0:
+                idx = len(entry['stats']) - 1
+                try:
+                    computed = entry[f'JA4SSH.{idx}']
+                    print (f'computed JA4SSH.{idx}: {computed}')
+                except Exception as e:
+                    pass
+
 # computes the JA4SSH from the segment x:
 # The segment has data as specified by ja4sh_stats
 ##
 def to_ja4ssh(x):
-    e = x['stats'][-1]
     idx = len(x['stats'])
+    print (f'calling ja4ssh with idx = {idx}')
+    e = x['stats'][idx-1]
     if e['client_payloads'] or e['server_payloads']:
         mode_client = max(e['client_payloads'], key=e['client_payloads'].count) if e['client_payloads'] else 0
         mode_server = max(e['server_payloads'], key=e['server_payloads'].count) if e['server_payloads'] else 0