analyze_clients: use output from count-packets-per-ip

This is order of magnitude faster than dpkt implementation.

pcap/count-packets-per-ip.lua

@@ -93,7 +93,7 @@ local duration_s = (now_ms - chunk_since_ms) / 1e3
 log:info(string.format("duration of input PCAP (s): %.3f", duration_s))
 log:info(string.format("number of clients: %d", n_clients))
 
-csv_output:write("ip,ip_since_ms,ip_until_ms,packets,ip_chunk_qps\n")
+csv_output:write('"ip","ip_since_ms","ip_until_ms","packets","ip_chunk_qps"\n')
 for ip, data in pairs(clients) do
 	csv_output:write('"')
 	csv_output:write(ip)

tools/analyze_clients.py

@@ -2,20 +2,14 @@
 
 # pylint: disable=wrong-import-order,wrong-import-position
 import argparse
-from collections import defaultdict
-import ipaddress
+import csv
 import logging
 import os
 import statistics
 import sys
 import traceback
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Union
 
-import dns
-import dns.exception
-import dns.message
-
-import dpkt
 # Force matplotlib to use a different backend to handle machines without a display
 import matplotlib
 matplotlib.use('Agg')
@@ -24,28 +18,9 @@ from matplotlib.lines import Line2D  # noqa
 import matplotlib.pyplot as plt  # noqa
 
 
-IP = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
 SCALE_MAGIC = 10000
 
 
-class MockRaw:
-    def __init__(self, data: bytes) -> None:
-        try:
-            self.data = dpkt.ip.IP(data)
-        except dpkt.UnpackError:
-            self.data = dpkt.ip6.IP6(data)
-
-
-LINK_TYPES = {
-    dpkt.pcap.DLT_EN10MB: dpkt.ethernet.Ethernet,
-    dpkt.pcap.DLT_LINUX_SLL: dpkt.sll.SLL,
-    # dpkt.pcap.DLT_RAW can have 3 different values in dpkt - Gotta Catch Em All!
-    12: MockRaw,
-    14: MockRaw,
-    101: MockRaw,
-}
-
-
 COLORS = [matplotlib.colors.to_rgba(c)
           for c in plt.rcParams['axes.prop_cycle'].by_key()['color']]
 
@@ -67,78 +42,15 @@ def init_plot(title):
     return ax
 
 
-def create_filter(ips: Optional[List[IP]] = None) -> str:
-    cap_filter = 'udp dst port 53'
-    if ips:
-        hosts = ['host {}'.format(ip) for ip in ips]
-        cap_filter += ' and ({})'.format(' or '.join(hosts))
-    return cap_filter
-
-
 def count_client_queries(
             filename: str,
-            ips: Optional[List[IP]] = None,
-            since: float = 0,
-            until: float = float('+inf'),
-            include_malformed: bool = False
-        ) -> Dict[bytes, int]:
-    clients = defaultdict(int)  # type: Dict[bytes, int]
-
-    with open(filename, 'rb') as fin:
-        pcap = dpkt.pcap.Reader(fin)
-        filter_ = create_filter(ips)
-        logging.debug('using filter: "%s"', filter_)
-        pcap.setfilter(filter_)
-        if pcap.datalink() not in LINK_TYPES:
-            logging.critical("Unsupported PCAP linktype: %d", pcap.datalink())
-            sys.exit(1)
-
-        parse = LINK_TYPES[pcap.datalink()]
-        start = None
-        end = None
-
-        for ts, pkt in pcap:
-            if start is None:
-                start = ts + since
-                end = ts + until
+        ) -> Dict[str, int]:
+    with open(filename, newline='') as csvfile:
+        reader = csv.DictReader(csvfile, delimiter=',', quotechar='"', quoting=csv.QUOTE_NONNUMERIC)
+        return {row['ip']: int(row['packets']) for row in reader}
 
-            if ts < start:
-                continue
-            if ts >= end:
-                break
 
-            link = parse(pkt)
-
-            ip = link.data
-            if not isinstance(ip, (dpkt.ip.IP, dpkt.ip6.IP6)):
-                continue
-
-            udp = ip.data  # NOTE: ip packet musn't be fragmented
-            if not isinstance(udp, dpkt.udp.UDP):
-                continue
-
-            payload = udp.data
-            if not isinstance(payload, bytes):
-                continue
-
-            if len(payload) < 3:
-                continue  # small garbage isn't supported
-            if payload[2] & 0x80:
-                continue  # QR=1 -> response
-
-            if not include_malformed:
-                try:
-                    dns.message.from_wire(payload)
-                except dns.exception.FormError:
-                    continue
-
-            # do mapping from original ip to client; new client otherwise
-            clients[ip.src] += 1
-
-    return clients
-
-
-def plot_client_query_scatter(ax, clients: Dict[bytes, int], color):
+def plot_client_query_scatter(ax, clients: Dict[str, int], color):
     data = clients.values()
 
     x = []
@@ -148,6 +60,7 @@ def plot_client_query_scatter(ax, clients: Dict[bytes, int], color):
     lmax = 10
     while lmax <= max(data):
         samples = list(n for n in data if lmin <= n < lmax)
+        sanity_nsamples += len(samples)
         x.append(statistics.mean(samples))
         y.append(len(samples) / len(data) * 100)
         s.append(sum(samples))
@@ -176,42 +89,22 @@ def main():
 
     parser = argparse.ArgumentParser(
         description='Analyze query distribution among clients in input pcap')
-    parser.add_argument('pcap', nargs='+', help='PCAP(s) to analyze (output from pellet.py)')
+    parser.add_argument('csv', nargs='+',
+                        help='CSV(s) to visualize (output from count-packets-per-ip.lua)')
     parser.add_argument('-o', '--output', type=str, default='clients.svg',
                         help='output filename (default: clients.svg)')
-    parser.add_argument(
-        '-r', '--resolvers', type=str, nargs='*',
-        help='only use data flowing to specified addresses (IP/IPv6)')
-    parser.add_argument('--since', type=float, default=0,
-                        help='Omit data before this time (secs since test start)')
-    parser.add_argument('--until', type=float, default=float('+inf'),
-                        help='Omit data after this time (secs since test start)')
-    parser.add_argument('-m', '--include-malformed', action='store_true',
-                        help='include malformed packets')
     args = parser.parse_args()
 
-    ips = []
-    if args.resolvers:
-        for resolver in args.resolvers:
-            try:
-                ip = ipaddress.ip_address(resolver)
-            except ValueError as exc:
-                logging.critical('--resolvers: %s', exc)
-                sys.exit(1)
-            else:
-                ips.append(ip)
-
     ax = init_plot("Query distribution among clients")
     handles = []
     lines = []
     labels = []
 
-    for color, pcap in zip(COLORS, args.pcap):
-        label = os.path.basename(pcap)
+    for color, csv_inf in zip(COLORS, args.csv):
+        label = os.path.basename(csv_inf)
         logging.info('Processing: %s', label)
         try:
-            clients = count_client_queries(pcap, ips, args.since, args.until,
-                                           args.include_malformed)
+            clients_qps = count_client_queries(csv_inf)
         except FileNotFoundError as exc:
             logging.critical('%s', exc)
             sys.exit(1)
@@ -222,7 +115,7 @@ def main():
         else:
             labels.append(label)
             lines.append(Line2D([0], [0], color=color, lw=4))
-            handles.append(plot_client_query_scatter(ax, clients, color))
+            handles.append(plot_client_query_scatter(ax, clients_qps, color))
 
     ax.legend(lines, labels, loc="lower left")
     plt.savefig(args.output, dpi=300)