From 0360f3134e13c92eeac39c438541d232556d9f3f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tarek=20Ziad=C3=A9?= <tarek@mozilla.com>
Date: Mon, 3 Feb 2020 15:03:28 +0000
Subject: [PATCH] Bug 1612790 - Dump domain stats r=Bebe

This patch allows us to dump per-domain hit counts to add more insights in the
proxy usage.

Differential Revision: https://phabricator.services.mozilla.com/D61432

--HG--
extra : moz-landing-system : lando
---
 .../mozproxy/mozproxy/backends/mitm/mitm.py     |  1 +
 .../scripts/alternate-server-replay-4.0.4.py    | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/testing/mozbase/mozproxy/mozproxy/backends/mitm/mitm.py b/testing/mozbase/mozproxy/mozproxy/backends/mitm/mitm.py
index a954b2c19b0b..a8e0ab57db16 100644
--- a/testing/mozbase/mozproxy/mozproxy/backends/mitm/mitm.py
+++ b/testing/mozbase/mozproxy/mozproxy/backends/mitm/mitm.py
@@ -224,6 +224,7 @@ class Mitmproxy(Playback):
                 args = [
                     "-v",
                     "--set", "upstream_cert=false",
+                    "--set", "upload_dir=" + self.upload_dir,
                     "--set", "websocket=false",
                     "--set", "server_replay_files={}".format(",".join(recording_paths)),
                     "--scripts", script,
diff --git a/testing/mozbase/mozproxy/mozproxy/backends/mitm/scripts/alternate-server-replay-4.0.4.py b/testing/mozbase/mozproxy/mozproxy/backends/mitm/scripts/alternate-server-replay-4.0.4.py
index 76fc3e92559f..a1479d3754ef 100644
--- a/testing/mozbase/mozproxy/mozproxy/backends/mitm/scripts/alternate-server-replay-4.0.4.py
+++ b/testing/mozbase/mozproxy/mozproxy/backends/mitm/scripts/alternate-server-replay-4.0.4.py
@@ -10,6 +10,7 @@ import os
 import json
 import hashlib
 import urllib
+from collections import defaultdict
 
 import typing
 from urllib import parse
@@ -85,6 +86,7 @@ class AlternateServerPlayback:
         ctx.master.addons.remove(ctx.master.addons.get("serverplayback"))
         self.flowmap = {}
         self.configured = False
+        self.netlocs = defaultdict(int)
 
     def load(self, loader):
         loader.add_option(
@@ -93,6 +95,10 @@ class AlternateServerPlayback:
             [],
             "Replay server responses from a saved file.",
         )
+        loader.add_option(
+            "upload_dir", str, "",
+            "Upload directory",
+        )
 
     def load_flows(self, flows):
         """
@@ -171,8 +177,19 @@ class AlternateServerPlayback:
             self.configured = True
             self.load_files(ctx.options.server_replay_files)
 
+    def done(self):
+        if not ctx.options.upload_dir:
+            return
+        netlocs = sorted(self.netlocs.items(), key=lambda item: -item[1])
+        path = os.path.join(ctx.options.upload_dir, "mitm_netlocs.log")
+        with open(path, "w") as f:
+            for url, count in netlocs:
+                f.write("%s: %d\n" % (url, count))
+
     def request(self, f):
         if self.flowmap:
+            parsed_url = urllib.parse.urlparse(parse.unquote(f.request.url))
+            self.netlocs[parsed_url.netloc] += 1
             rflow = self.next_flow(f)
             if rflow:
                 response = rflow.response.copy()