diff options
Diffstat (limited to 'examples')
| -rwxr-xr-x | examples/mitmproxywrapper.py | 86 | 
1 files changed, 70 insertions, 16 deletions
diff --git a/examples/mitmproxywrapper.py b/examples/mitmproxywrapper.py index 48963896..27ffd609 100755 --- a/examples/mitmproxywrapper.py +++ b/examples/mitmproxywrapper.py @@ -10,12 +10,15 @@  import subprocess  import re  import argparse +import contextlib +import os +import sys  class Wrapper(object): -    def __init__(self, port): +    def __init__(self, port, extra_arguments=None):          self.port = port -        self.primary_service_name = self.find_primary_service_name() +        self.extra_arguments = extra_arguments      def run_networksetup_command(self, *arguments):          return subprocess.check_output(['sudo', 'networksetup'] + list(arguments)) @@ -50,40 +53,91 @@ 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(): +            cmd = ['mitmproxy', '-p', str(self.port)] +            if self.extra_arguments: +                cmd.extend(self.extra_arguments) +            cmd.extend(['--palette', 'light']) +            subprocess.check_call(cmd) + +    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 ensure_superuser(cls): +        if os.getuid() != 0: +            print 'Relaunching with sudo...' +            os.execv('/usr/bin/sudo', ['/usr/bin/sudo'] + sys.argv)      @classmethod      def main(cls): -        parser = argparse.ArgumentParser(description='Helper tool for OS X proxy configuration and mitmproxy') +        parser = argparse.ArgumentParser( +            description='Helper tool for OS X proxy configuration and mitmproxy.', +            epilog='Any additional arguments will be passed on unchanged to 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() +        args, extra_arguments = parser.parse_known_args() -        wrapper = cls(port=args.port) +        wrapper = cls(port=args.port, extra_arguments=extra_arguments) +                  if args.toggle:              wrapper.toggle_proxy() +#         elif args.honeyproxy: +#             wrapper.wrap_honeyproxy()          else:              wrapper.wrap_mitmproxy() +  if __name__ == '__main__': +    Wrapper.ensure_superuser()      Wrapper.main()  | 
