From 8487dcd0b9565657e1e6e89c7d8467d54a7c41ba Mon Sep 17 00:00:00 2001
From: Kali Kaneko <kali@leap.se>
Date: Fri, 11 Aug 2017 13:17:37 -0400
Subject: [feature] allow manual gateway selection for vpn

For now, the way to select a gateway is to add a section in
bitmaskd.cfg:

  [vpn_prefs]
  locations = ["frankfurt", "seattle__wa"]
  countries = ["DE", "US"]

Note that the location indication has priority over country code.

This will be exposed by the UI in release 0.11

- Resolves: #8855
---
 src/leap/bitmask/cli/command.py  |  2 +-
 src/leap/bitmask/vpn/gateways.py | 59 ++++++++++++++++++++++++++++++++++++++--
 src/leap/bitmask/vpn/service.py  | 20 +++++++++++---
 3 files changed, 73 insertions(+), 8 deletions(-)

(limited to 'src')

diff --git a/src/leap/bitmask/cli/command.py b/src/leap/bitmask/cli/command.py
index 1daadc5e..5586d091 100644
--- a/src/leap/bitmask/cli/command.py
+++ b/src/leap/bitmask/cli/command.py
@@ -54,7 +54,7 @@ def default_dict_printer(result):
 
     for key, value in result.items():
         if isinstance(value, list):
-            if isinstance(value[0], list):
+            if value and isinstance(value[0], list):
                 value = map(lambda l: ' '.join(l), value)
                 for item in value:
                     pprint('\t' + item)
diff --git a/src/leap/bitmask/vpn/gateways.py b/src/leap/bitmask/vpn/gateways.py
index 950e37c3..92367c30 100644
--- a/src/leap/bitmask/vpn/gateways.py
+++ b/src/leap/bitmask/vpn/gateways.py
@@ -18,7 +18,7 @@
 """
 Gateway Selection
 """
-
+import copy
 import logging
 import os
 import re
@@ -27,24 +27,49 @@ import time
 from leap.common.config import get_path_prefix
 
 
+def _normalized(label):
+    return label.lower().replace(',', '_').replace(' ', '_')
+
+
 class GatewaySelector(object):
 
     # http://www.timeanddate.com/time/map/
     equivalent_timezones = {13: -11, 14: -10}
 
-    def __init__(self, gateways=None, locations=None, tz_offset=None):
+    def __init__(self, gateways=None, locations=None, tz_offset=None, preferred=None):
         '''
         Constructor for GatewaySelector.
 
+        By default, we will use a Time Zone Heuristic to choose the closest
+        gateway to the user.
+
+        If the user specified something in the 'vpn_prefs' section of
+        bitmaskd.cfg, we will passed a dictionary here with entries for
+        'countries' and 'locations', in the form of a list. The 'locations'
+        entry has preference over the 'countries' one.
+
+        :param gateways: an unordered list of all the available gateways, read
+                         from the eip-config file.
+        :type gateways: list
+        :param locations: a dictionary with the locations, as read from the
+                          eip-config file
+        :type locations: dict
         :param tz_offset: use this offset as a local distance to GMT.
         :type tz_offset: int
+        :param preferred: a dictionary containing the country code (cc) and the
+                          locations (locations) manually preferred by the user.
+        :type preferred: dict
         '''
         if gateways is None:
             gateways = []
         if locations is None:
             locations = {}
+        if preferred is None:
+            preferred = {}
         self.gateways = gateways
         self.locations = locations
+        self.preferred = preferred
+
         self._local_offset = tz_offset
         if tz_offset is None:
             tz_offset = self._get_local_offset()
@@ -90,7 +115,35 @@ class GatewaySelector(object):
         result = []
         for ip, distance, label, country in gateways_timezones:
             result.append((label, ip, country))
-        return result
+
+        filtered = self.apply_user_preferences(result)
+        return filtered 
+
+    def apply_user_preferences(self, options):
+        """
+        We re-sort the pre-sorted list of gateway options, according with the
+        user's preferences.
+
+        Location has preference over the Country Code indication.
+        """
+        applied = []
+        presorted = copy.copy(options)
+        for location in self.preferred.get('loc', []):
+            for index, data in enumerate(presorted):
+                label, ip, country = data
+                if _normalized(label) == _normalized(location):
+                    applied.append((label, ip, country))
+                    presorted.pop(index)
+
+        for cc in self.preferred.get('cc', []):
+            for index, data in enumerate(presorted):
+                label, ip, country = data
+                if _normalized(country) == _normalized(cc):
+                    applied.append((label, ip, country))
+                    presorted.pop(index)
+        if presorted:
+            applied += presorted 
+        return applied
 
     def get_gateways_country_code(self):
         country_codes = {}
diff --git a/src/leap/bitmask/vpn/service.py b/src/leap/bitmask/vpn/service.py
index b0dcba86..d2f952ef 100644
--- a/src/leap/bitmask/vpn/service.py
+++ b/src/leap/bitmask/vpn/service.py
@@ -19,7 +19,7 @@
 """
 VPN service declaration.
 """
-
+import json
 import os
 
 from time import strftime
@@ -133,6 +133,7 @@ class VPNService(HookableService):
         return {'result': 'vpn stopped'}
 
     def do_status(self):
+        # TODO - add the current gateway and CC to the status
         childrenStatus = {
             'vpn': {'status': 'off', 'error': None},
             'firewall': {'status': 'off', 'error': None},
@@ -220,10 +221,21 @@ class VPNService(HookableService):
         bonafide = self.parent.getServiceNamed('bonafide')
         config = yield bonafide.do_provider_read(provider, 'eip')
 
-        sorted_gateways = GatewaySelector(
-            config.gateways, config.locations).select_gateways()
+        try:
+            _cco = self.parent.get_config('vpn_prefs', 'countries', "")
+            pref_cco = json.loads(_cco)
+        except ValueError:
+            pref_cco = []
+        try:
+            _loc = self.parent.get_config('vpn_prefs', 'locations', "")
+            pref_loc = json.loads(_loc)
+        except ValueError:
+            pref_loc = []
 
-        # TODO - add manual gateway selection ability.
+        sorted_gateways = GatewaySelector(
+            config.gateways, config.locations,
+            preferred={'cc': pref_cco, 'loc': pref_loc}
+            ).select_gateways()
 
         extra_flags = config.openvpn_configuration
 
-- 
cgit v1.2.3