Merge remote-tracking branch 'refs/remotes/kali/feature/vpn_exit_icon' into develop

2 files changed, 32 insertions, 5 deletions

diff --git a/src/leap/bitmask/services/eip/eipconfig.py b/src/leap/bitmask/services/eip/eipconfig.py
index 09a3d257..e7419b22 100644
--- a/src/leap/bitmask/services/eip/eipconfig.py
+++ b/src/leap/bitmask/services/eip/eipconfig.py
@@ -110,7 +110,7 @@ class VPNGatewaySelector(object):
 
     def get_gateways_list(self):
         """
-        Returns the existing gateways, sorted by timezone proximity.
+        Return the existing gateways, sorted by timezone proximity.
 
         :rtype: list of tuples (location, ip)
                 (str, IPv4Address or IPv6Address object)
@@ -148,16 +148,36 @@ class VPNGatewaySelector(object):
 
     def get_gateways(self):
         """
-        Returns the 4 best gateways, sorted by timezone proximity.
+        Return the 4 best gateways, sorted by timezone proximity.
 
         :rtype: list of IPv4Address or IPv6Address object.
         """
         gateways = [ip for location, ip in self.get_gateways_list()][:4]
         return gateways
 
+    def get_gateways_country_code(self):
+        """
+        Return a dict with ipaddress -> country code mapping.
+
+        :rtype: dict
+        """
+        country_codes = {}
+
+        locations = self._eipconfig.get_locations()
+        gateways = self._eipconfig.get_gateways()
+
+        for idx, gateway in enumerate(gateways):
+            gateway_location = gateway.get('location')
+
+            ip = self._eipconfig.get_gateway_ip(idx)
+            if gateway_location is not None:
+                ccode = locations[gateway['location']]['country_code']
+                country_codes[ip] = ccode
+        return country_codes
+
     def _get_timezone_distance(self, offset):
         '''
-        Returns the distance between the local timezone and
+        Return the distance between the local timezone and
         the one with offset 'offset'.
 
         :param offset: the distance of a timezone to GMT.
@@ -179,7 +199,7 @@ class VPNGatewaySelector(object):
 
     def _get_local_offset(self):
         '''
-        Returns the distance between GMT and the local timezone.
+        Return the distance between GMT and the local timezone.
 
         :rtype: int
         '''
diff --git a/src/leap/bitmask/services/eip/vpnlauncher.py b/src/leap/bitmask/services/eip/vpnlauncher.py
index dcb48e8a..9629afae 100644
--- a/src/leap/bitmask/services/eip/vpnlauncher.py
+++ b/src/leap/bitmask/services/eip/vpnlauncher.py
@@ -25,6 +25,7 @@ import stat
 from abc import ABCMeta, abstractmethod
 from functools import partial
 
+from leap.bitmask.config import flags
 from leap.bitmask.config.leapsettings import LeapSettings
 from leap.bitmask.config.providerconfig import ProviderConfig
 from leap.bitmask.platform_init import IS_LINUX
@@ -122,9 +123,9 @@ class VPNLauncher(object):
         leap_settings = LeapSettings()
         domain = providerconfig.get_domain()
         gateway_conf = leap_settings.get_selected_gateway(domain)
+        gateway_selector = VPNGatewaySelector(eipconfig)
 
         if gateway_conf == leap_settings.GATEWAY_AUTOMATIC:
-            gateway_selector = VPNGatewaySelector(eipconfig)
             gateways = gateway_selector.get_gateways()
         else:
             gateways = [gateway_conf]
@@ -133,6 +134,12 @@ class VPNLauncher(object):
             logger.error('No gateway was found!')
             raise VPNLauncherException('No gateway was found!')
 
+        # this only works for selecting the first gateway, as we're
+        # currently doing.
+        ccodes = gateway_selector.get_gateways_country_code()
+        gateway_ccode = ccodes[gateways[0]]
+        flags.CURRENT_VPN_COUNTRY = gateway_ccode
+
         logger.debug("Using gateways ips: {0}".format(', '.join(gateways)))
         return gateways