Merge pull request #4 from mkoistinen/multiple-proxy-type-support

Dynamically build menu-items based on configured proxies

ToggleProxy.py

@@ -7,53 +7,70 @@ import  commands, re
 
 class ToggleProxy(NSObject):
 
+    # This is a dictionary of the proxy-types we support, each with a
+    # dictionary of some unique attributes for each, namely:
+    #
+    #   'pref'          : This is a constant defining which preference itemmarks if this proxy is enabled
+    #   'title'         : This is what will appear in the menu
+    #   'action'        : This is the method that will be called if the user toggles this proxies menuitem
+    #   'keyEquivalent' : Self-explanatory, but unused
+    #   'menuitem'      : This will store the menu item for this proxy once it is created
+
+    proxies = {
+        'ftp'  : { 'pref': kSCPropNetProxiesFTPEnable,   'title': 'FTP Proxy',   'action': 'toggleFtpProxy:',   'keyEquivalent': "", 'menuitem': None },
+        'http' : { 'pref': kSCPropNetProxiesHTTPEnable,  'title': 'HTTP Proxy',  'action': 'toggleHttpProxy:',  'keyEquivalent': "", 'menuitem': None },
+        'https': { 'pref': kSCPropNetProxiesHTTPSEnable, 'title': 'HTTPS Proxy', 'action': 'toggleHttpsProxy:', 'keyEquivalent': "", 'menuitem': None },
+        'rtsp' : { 'pref': kSCPropNetProxiesRTSPEnable,  'title': 'RTSP Proxy',  'action': 'toggleRtspProxy:',  'keyEquivalent': "", 'menuitem': None },
+        'socks': { 'pref': kSCPropNetProxiesSOCKSEnable, 'title': 'SOCKS Proxy', 'action': 'toggleSocksProxy:', 'keyEquivalent': "", 'menuitem': None },
+    }
+
     def applicationDidFinishLaunching_(self, notification):
         # load icon files
-        self.icons = {
-            '0-0-0' : NSImage.imageNamed_("icon-0-0-0"),
-            '1-0-0' : NSImage.imageNamed_("icon-1-0-0"),
-            '0-1-0' : NSImage.imageNamed_("icon-0-1-0"),
-            '0-0-1' : NSImage.imageNamed_("icon-0-0-1"),
-            '1-1-0' : NSImage.imageNamed_("icon-1-1-0"),
-            '1-0-1' : NSImage.imageNamed_("icon-1-0-1"),
-            '1-1-1' : NSImage.imageNamed_("icon-1-1-1")
-        }
+        self.active_image   = NSImage.imageNamed_("active")
+        self.inactive_image = NSImage.imageNamed_("inactive")
 
         # make status bar item
         self.statusitem = NSStatusBar.systemStatusBar().statusItemWithLength_(NSVariableStatusItemLength)
         self.statusitem.retain()
         self.statusitem.setHighlightMode_(False)
         self.statusitem.setEnabled_(True)
-        self.statusitem.setImage_(self.icons['0-0-0'])
 
         # insert a menu into the status bar item
         self.menu = NSMenu.alloc().init()
         self.statusitem.setMenu_(self.menu)
 
-        # add items to menu
-        self.httpMenuItem = self.menu.addItemWithTitle_action_keyEquivalent_(
-            "HTTP proxy",
-            "toggleHttpProxy:",
-            "")
-        self.httpsMenuItem = self.menu.addItemWithTitle_action_keyEquivalent_(
-            "HTTPS proxy",
-            "toggleHttpsProxy:",
-            "")
-        self.socksMenuItem = self.menu.addItemWithTitle_action_keyEquivalent_(
-            "SOCKS proxy",
-            "toggleSocksProxy:",
-            "")
-        self.menu.addItem_(NSMenuItem.separatorItem())
-        self.menu.addItemWithTitle_action_keyEquivalent_(
-            "Quit",
-            "quitApp:",
-            "")
-
         # open connection to the dynamic (configuration) store
         self.store = SCDynamicStoreCreate(None, "name.klep.toggleproxy", self.dynamicStoreCallback, None)
 
-        # start working
-        self.loadNetworkServices()
+        proxyRef = SCNetworkServiceCopyProtocol(self.service, kSCNetworkProtocolTypeProxies)
+        prefDict = SCNetworkProtocolGetConfiguration(proxyRef)
+
+        separatorRequired = False
+
+        # For each of the proxies we are concerned with, check to see if any
+        # are configured. If so (even if not enabled), create a menuitem for
+        # that proxy type.
+
+        for proxy in self.proxies.values():
+            enabled = CFDictionaryGetValue(prefDict, proxy['pref'])
+            if enabled is not None:
+                proxy['menuitem'] = self.menu.addItemWithTitle_action_keyEquivalent_(
+                    proxy['title'],
+                    proxy['action'],
+                    proxy['keyEquivalent']
+                )
+                separatorRequired = True
+            else:
+                proxy['menuitem'] = None
+
+        if separatorRequired:
+            self.menu.addItem_(NSMenuItem.separatorItem())
+
+        # Need a way to quit
+        self.menu.addItemWithTitle_action_keyEquivalent_("Quit", "quitApp:", "")
+
+        # Start working
+        # self.loadNetworkServices()
         self.watchForProxyChanges()
         self.updateProxyStatus()
 
@@ -62,20 +79,17 @@ class ToggleProxy(NSObject):
         # get primary interface
         return SCDynamicStoreCopyValue(self.store, 'State:/Network/Global/IPv4')['PrimaryInterface']
 
-    def loadNetworkServices(self):
-        """ Load and store a list of network services indexed by their BSDName """
-        self.services   = {}
-        # Create a dummy preference reference
+    @property
+    def service(self):
+        """ Returns the service relating to self.interface """
         prefs = SCPreferencesCreate(kCFAllocatorDefault, 'PRG', None)
+
         # Fetch the list of services
-        for service in SCNetworkServiceCopyAll(prefs):
-            # This is what we're after, the user-defined service name e.g., "Built-in Ethernet"
-            name      = SCNetworkServiceGetName(service)
-            # Interface reference
-            interface = SCNetworkServiceGetInterface(service)
-            # The BSDName of the interface, E.g., "en1", this will be the index of our list
-            bsdname   = SCNetworkInterfaceGetBSDName(interface)
-            self.services[bsdname] = name
+        for serviceRef in SCNetworkServiceCopyAll(prefs):
+            interface = SCNetworkServiceGetInterface(serviceRef)
+            if self.interface == SCNetworkInterfaceGetBSDName(interface):
+                return serviceRef
+        return None
 
     def watchForProxyChanges(self):
         """ install a watcher for proxy changes """
@@ -92,40 +106,45 @@ class ToggleProxy(NSObject):
     def updateProxyStatus(self):
         """ update proxy status """
         # load proxy dictionary
-        proxydict       = SCDynamicStoreCopyProxies(None)
+        proxydict = SCDynamicStoreCopyProxies(None)
 
         # get status for primary interface
-        status          = proxydict['__SCOPED__'][self.interface]
+        status    = proxydict['__SCOPED__'][self.interface]
+
+        # Are any proxies active now?
+        anyProxyEnabled = False
 
         # update menu items according to their related proxy state
-        self.httpMenuItem.setState_(  status.get('HTTPEnable', False)  and NSOnState or NSOffState )
-        self.httpsMenuItem.setState_( status.get('HTTPSEnable', False) and NSOnState or NSOffState )
-        self.socksMenuItem.setState_( status.get('SOCKSEnable', False) and NSOnState or NSOffState )
-
-        # update icon
-        self.statusitem.setImage_(
-            self.icons['%d-%d-%d' % (
-                status.get('HTTPEnable', False)  and 1 or 0,
-                status.get('HTTPSEnable', False) and 1 or 0,
-                status.get('SOCKSEnable', False) and 1 or 0
-            )]
-        )
+        for proxy in self.proxies.values():
+            if proxy['menuitem']:
+                proxy['menuitem'].setState_(status.get(proxy['pref'], False) and NSOnState or NSOffState)
+                if status.get(proxy['pref'], False):
+                    anyProxyEnabled = True
+
+        # set image
+        self.statusitem.setImage_(anyProxyEnabled and self.active_image or self.inactive_image)
 
     def quitApp_(self, sender):
         NSApp.terminate_(self)
 
+    def toggleFtpProxy_(self, sender):
+        self.toggleProxy(self.proxies['ftp']['menuitem'], 'ftpproxy')
+
     def toggleHttpProxy_(self, sender):
-        self.toggleProxy(self.httpMenuItem, 'webproxy')
+        self.toggleProxy(self.proxies['http']['menuitem'], 'webproxy')
 
     def toggleHttpsProxy_(self, sender):
-        self.toggleProxy(self.httpsMenuItem, 'securewebproxy')
+        self.toggleProxy(self.proxies['https']['menuitem'], 'securewebproxy')
+
+    def toggleRtspProxy_(self, sender):
+        self.toggleProxy(self.proxies['socks']['menuitem'], 'streamingproxy')
 
     def toggleSocksProxy_(self, sender):
-        self.toggleProxy(self.socksMenuItem, 'socksfirewallproxy')
+        self.toggleProxy(self.proxies['socks']['menuitem'], 'socksfirewallproxy')
 
     def toggleProxy(self, item, target):
         """ callback for clicks on menu item """
-        servicename = self.services.get(self.interface)
+        servicename = SCNetworkServiceGetName(self.service)
         if not servicename:
             NSLog("interface '%s' not found in services?" % self.interface)
             return