From 2f285ca6764d0d44446bfd8f517805559ba098ee Mon Sep 17 00:00:00 2001
From: Allen Luce <aluce@adonnetwork.com>
Date: Thu, 7 Sep 2017 16:00:24 -0700
Subject: [PATCH] Add stream following capability

The -J option will print packets following a match but only those that
are in the same TCP stream as the match.

If another packet matches during the following, the number counter is
reset. This helps mitigate problems with a match not having enough
packets after it in the same stream (which could mean ngrep prints
nothing forever after the match). This functionality can be defeated
with the -Q option.
---
 ngrep.8 | 12 +++++++++++-
 ngrep.c | 52 +++++++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/ngrep.8 b/ngrep.8
index 1750b83..c238f08 100644
--- a/ngrep.8
+++ b/ngrep.8
@@ -12,7 +12,7 @@ ngrep \- network grep
 
 .SH SYNOPSIS
 
-.B ngrep <-hNXViwqpevxlDtTRM> <-IO
+.B ngrep <-hNXViwqpevxlDtTRMQ> <-IO
 .I pcap_dump
 .B > < -n
 .I num
@@ -20,6 +20,8 @@ ngrep \- network grep
 .I dev
 .B > < -A
 .I num
+.B > < -J
+.I num
 .B > < -s
 .I snaplen
 .B > < -S
@@ -184,6 +186,14 @@ this option to force ngrep to listen on interface \fIdev\fP.
 .IP "-A num"
 Dump \fInum\fP packets of trailing context after matching a packet.
 
+.IP "-J num"
+Like -A, but print only those packets in the same TCP stream as the
+matched packet.
+
+.IP -Q
+When a new matching packet comes in but the number of packets given to
+-J hasn't been reached yet, start with this new match and follow it instead.
+ 
 .IP "\fI match expression\fP"
 A match expression is either an extended regular expression, or if the
 \fI-X\fP option is specified, a string signifying a hexadecimal value.
diff --git a/ngrep.c b/ngrep.c
index 3df9389..6688d19 100644
--- a/ngrep.c
+++ b/ngrep.c
@@ -106,7 +106,7 @@
 
 uint32_t snaplen = 65535, limitlen = 65535, promisc = 1, to = 100;
 uint32_t match_after = 0, keep_matching = 0, matches = 0, max_matches = 0;
-uint32_t seen_frames = 0;
+uint32_t seen_frames = 0, follow_stream = 0, reset_count = 0;
 
 #if USE_TCPKILL
 uint32_t tcpkill_active = 0;
@@ -186,6 +186,12 @@ void (*print_time)() = NULL, (*dump_delay)() = dump_delay_proc_init;
 
 uint32_t ws_row, ws_col = 80, ws_col_forced = 0;
 
+/*
+ * Stream following functionality
+ */
+const char *follow_ip_src = NULL, *follow_ip_dst = NULL;
+uint16_t follow_sport, follow_dport;
+
 
 int main(int argc, char **argv) {
     int32_t c;
@@ -210,8 +216,7 @@ int main(int argc, char **argv) {
         setlocale(LC_CTYPE, locale);
     }
 #endif
-
-    while ((c = getopt(argc, argv, "LNhXViwqpevxlDtTRMK:Cs:n:c:d:A:I:O:S:P:F:W:")) != EOF) {
+    while ((c = getopt(argc, argv, "LNhXViwqpevxlDtTRMQK:Cs:n:c:d:A:I:J:O:S:P:F:W:")) != EOF) {
         switch (c) {
             case 'W': {
                 if (!strcasecmp(optarg, "normal"))
@@ -250,6 +255,15 @@ int main(int argc, char **argv) {
                 if (match_after < UINT32_MAX)
                     match_after++;
                 break;
+            case 'J':
+                match_after = _atoui32(optarg);
+                if (match_after < UINT32_MAX)
+                    match_after++;
+                follow_stream = 1;
+                break;
+    	    case 'Q':
+                reset_count = 1;
+                break;
 #if defined(_WIN32)
             case 'L':
                 win32_listdevices();
@@ -913,9 +927,35 @@ void dump_packet(struct pcap_pkthdr *h, u_char *p, uint8_t proto, unsigned char
     if (len > limitlen)
         len = limitlen;
 
-    if ((len > 0 && match_func(data, len, &match_index, &match_size) == invert_match) && !keep_matching)
+    int notmatched = len > 0 && match_func(data, len, &match_index, &match_size) == invert_match;
+    if (notmatched && !keep_matching)
         return;
 
+    if (!notmatched && reset_count)
+        follow_ip_src = NULL;
+    
+    if (keep_matching && follow_stream) {
+        if (follow_ip_src == NULL) {
+            follow_ip_src = ip_src;
+            follow_ip_dst = ip_dst;
+            follow_sport = sport;
+            follow_dport = dport;
+        }
+        if (!(
+              (strcmp(ip_src, follow_ip_src) == 0 && // Forward stream
+               strcmp(ip_dst, follow_ip_dst) == 0 &&
+               sport == follow_sport &&
+               dport == follow_dport) ||
+              (!strcmp(ip_src, follow_ip_dst) == 0 && // Reverse stream
+               !strcmp(ip_dst, follow_ip_src) == 0 &&
+               sport == follow_dport &&
+               dport == follow_sport)
+              )) {
+            keep_matching++;
+            return;
+        }
+    }
+    
     if (!live_read && want_delay)
         dump_delay(h);
 
@@ -1408,7 +1448,7 @@ void usage(void) {
 #if defined(_WIN32)
            "L"
 #endif
-           "hNXViwqpevxlDtTRM> <-IO pcap_dump> <-n num> <-d dev> <-A num>\n"
+           "hNXViwqpevxlDtTRM> <-IO pcap_dump> <-n num> <-d dev> <-A num> <-J num>\n"
            "             <-s snaplen> <-S limitlen> <-W normal|byline|single|none> <-c cols>\n"
            "             <-P char> <-F file>"
 #if USE_TCPKILL
@@ -1437,6 +1477,8 @@ void usage(void) {
            "   -O  is dump matched packets in pcap format to pcap_dump\n"
            "   -n  is look at only num packets\n"
            "   -A  is dump num packets after a match\n"
+           "   -J  is dump num packets in TCP stream after a match\n"
+           "   -Q  is reset -J follow count after another match\n"
            "   -s  is set the bpf caplen\n"
            "   -S  is set the limitlen on matched packets\n"
            "   -W  is set the dump format (normal, byline, single, none)\n"