From e5f6ee68989949fcc77db037ab9a3e0da6e03cde Mon Sep 17 00:00:00 2001
From: Jonathan Cobb <jonathan@kyuss.org>
Date: Wed, 8 Apr 2020 16:18:34 -0400
Subject: [PATCH] monitor mitmdump memory usage, restart if memory goes above
 max allowed

---
 .../roles/mitmproxy/files/bubble_api.py       |  7 +++++++
 .../roles/mitmproxy/files/mitmdump_monitor.sh | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/automation/roles/mitmproxy/files/bubble_api.py b/automation/roles/mitmproxy/files/bubble_api.py
index 33b14fc4..382e129e 100644
--- a/automation/roles/mitmproxy/files/bubble_api.py
+++ b/automation/roles/mitmproxy/files/bubble_api.py
@@ -4,9 +4,16 @@
 import requests
 import traceback
 import sys
+import os
 import datetime
 from bubble_config import bubble_network, bubble_port
 
+# Write python PID to file so that mitmdump_monitor.sh can check for excessive memory usage and restart if needed
+MITMDUMP_PID_FILE_PATH = '/home/mitmproxy/mitmdump.pid'
+MITMDUMP_PID_FILE = open(MITMDUMP_PID_FILE_PATH, "w")
+MITMDUMP_PID_FILE.write("%d" % os.getpid())
+MITMDUMP_PID_FILE.close()
+
 HEADER_USER_AGENT = 'User-Agent'
 HEADER_REFERER = 'Referer'
 
diff --git a/automation/roles/mitmproxy/files/mitmdump_monitor.sh b/automation/roles/mitmproxy/files/mitmdump_monitor.sh
index 49816351..2743d91b 100644
--- a/automation/roles/mitmproxy/files/mitmdump_monitor.sh
+++ b/automation/roles/mitmproxy/files/mitmdump_monitor.sh
@@ -16,6 +16,8 @@ function log {
 
 BUBBLE_MITM_MARKER=/home/bubble/.mitmdump_monitor
 ROOT_KEY_MARKER=/usr/share/bubble/mitmdump_monitor
+MITMDUMP_PID_FILE=/home/mitmproxy/mitmdump.pid
+MAX_MITM_PCT_MEM=18
 
 # Start with MITM proxy turned off
 if [[ ! -f ${BUBBLE_MITM_MARKER} ]] ; then
@@ -57,5 +59,22 @@ while : ; do
       fi
     fi
   fi
+
+  # Check process memory usage, restart mitmdump if memory goes above max % allowed
+  if [[ -f ${MITMDUMP_PID_FILE} && -s ${MITMDUMP_PID_FILE} ]] ; then
+    MITM_PID="$(cat ${MITMDUMP_PID_FILE})"
+    PCT_MEM="$(ps q ${MITM_PID} -o %mem --no-headers | tr -d [[:space:]] | cut -f1 -d. | sed 's/[^0-9]*//g')"
+    # log "Info: mitmdump pid ${MITM_PID} using ${PCT_MEM}% of memory"
+    if [[ ! -z "${PCT_MEM}" ]] ; then
+      if [[ ${PCT_MEM} -ge ${MAX_MITM_PCT_MEM} ]] ; then
+        log "Warn: mitmdump: pid=$(cat ${MITMDUMP_PID_FILE}) memory used > max, restarting: ${PCT_MEM}% >= ${MAX_MITM_PCT_MEM}%"
+        supervisorctl restart mitmdump
+      fi
+    else
+      log "Error: could not determine mitmdump % memory, maybe PID file ${MITMDUMP_PID_FILE} is out of date? pid found was ${MITM_PID}"
+    fi
+  else
+    log "Error: mitmdump PID file ${MITMDUMP_PID_FILE} not found or empty"
+  fi
   sleep 5s
 done