Change the proxy settings for all connected interfaces, not just the primary one.

1 files changed, 49 insertions, 12 deletions

diff --git a/examples/mitmproxywrapper.py b/examples/mitmproxywrapper.py
index 48963896..751332b7 100755
--- a/examples/mitmproxywrapper.py
+++ b/examples/mitmproxywrapper.py
@@ -10,12 +10,12 @@
 import subprocess
 import re
 import argparse
+import contextlib
 
 class Wrapper(object):
     
     def __init__(self, port):
         self.port = port
-        self.primary_service_name = self.find_primary_service_name()
 
     def run_networksetup_command(self, *arguments):
         return subprocess.check_output(['sudo', 'networksetup'] + list(arguments))
@@ -50,40 +50,77 @@ class Wrapper(object):
         interface, = re.findall(r'PrimaryInterface\s*:\s*(.+)', stdout)
         return interface
 
-    def find_primary_service_name(self):
+    def primary_service_name(self):
         return self.interface_name_to_service_name_map()[self.primary_interace_name()]
 
     def proxy_enabled_for_service(self, service):
         return self.proxy_state_for_service(service)['Enabled'] == 'Yes'
 
     def toggle_proxy(self):
-        if self.proxy_enabled_for_service(self.primary_service_name):
-            self.disable_proxy_for_service(self.primary_service_name)
-        else:
-            self.enable_proxy_for_service(self.primary_service_name)
+        new_state = not self.proxy_enabled_for_service(self.primary_service_name())
+        for service_name in self.connected_service_names():
+            if self.proxy_enabled_for_service(service_name) and not new_state:
+                self.disable_proxy_for_service(service_name)
+            elif not self.proxy_enabled_for_service(service_name) and new_state:
+                self.enable_proxy_for_service(service_name)
+
+    def connected_service_names(self):
+        scutil_script = 'list\n'
+        stdout = self.run_command_with_input('/usr/sbin/scutil', scutil_script)
+        service_ids = re.findall(r'State:/Network/Service/(.+)/IPv4', stdout)
 
-    def wrap_mitmproxy(self):
-        if not self.proxy_enabled_for_service(self.primary_service_name):
-            self.enable_proxy_for_service(self.primary_service_name)
+        service_names = []
+        for service_id in service_ids:
+            scutil_script = 'show Setup:/Network/Service/{}\n'.format(service_id)
+            stdout = self.run_command_with_input('/usr/sbin/scutil', scutil_script)
+            service_name, = re.findall(r'UserDefinedName\s*:\s*(.+)', stdout)
+            service_names.append(service_name)
 
-        subprocess.check_call(['mitmproxy', '-p', str(self.port), '--palette', 'light'])
+        return service_names
 
-        if self.proxy_enabled_for_service(self.primary_service_name):
-            self.disable_proxy_for_service(self.primary_service_name)
+    def wrap_mitmproxy(self):
+        with self.wrap_proxy():
+            subprocess.check_call(['mitmproxy', '-p', str(self.port), '--palette', 'light'])
+
+    def wrap_honeyproxy(self):
+        with self.wrap_proxy():
+            popen = subprocess.Popen('honeyproxy.sh')
+            try:
+                popen.wait()
+            except KeyboardInterrupt:
+                popen.terminate()
+
+    @contextlib.contextmanager
+    def wrap_proxy(self):
+        connected_service_names = self.connected_service_names()
+        for service_name in connected_service_names:
+            if not self.proxy_enabled_for_service(service_name):
+                self.enable_proxy_for_service(service_name)
+        
+        yield
+
+        for service_name in connected_service_names:
+            if self.proxy_enabled_for_service(service_name):
+                self.disable_proxy_for_service(service_name)
 
     @classmethod
     def main(cls):
         parser = argparse.ArgumentParser(description='Helper tool for OS X proxy configuration and mitmproxy')
         parser.add_argument('-t', '--toggle', action='store_true', help='just toggle the proxy configuration')
+#         parser.add_argument('--honeyproxy', action='store_true', help='run honeyproxy instead of mitmproxy')
         parser.add_argument('-p', '--port', type=int, help='override the default port of 8080', default=8080)
         args = parser.parse_args()
 
         wrapper = cls(port=args.port)
+        
         if args.toggle:
             wrapper.toggle_proxy()
+#         elif args.honeyproxy:
+#             wrapper.wrap_honeyproxy()
         else:
             wrapper.wrap_mitmproxy()
 
+
 if __name__ == '__main__':
     Wrapper.main()