From dd2d92b0f93df120c2b5bf53677f21fc22530e61 Mon Sep 17 00:00:00 2001
From: Timothy Schroeder <tschroeder@zendesk.com>
Date: Tue, 2 Aug 2016 10:09:56 -0700
Subject: [PATCH] Allow using a proxy.

---
 pgoapi/pgoapi.py  | 14 ++++++++++----
 pgoapi/rpc_api.py | 10 +++++++---
 requirements.txt  |  3 ++-
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/pgoapi/pgoapi.py b/pgoapi/pgoapi.py
index 479406a7..2e9d998b 100755
--- a/pgoapi/pgoapi.py
+++ b/pgoapi/pgoapi.py
@@ -28,7 +28,6 @@
 import re
 import six
 import logging
-import requests
 
 from . import __title__, __version__, __copyright__
 from pgoapi.rpc_api import RpcApi
@@ -54,6 +53,7 @@ def __init__(self, provider=None, oauth2_refresh_token=None, username=None, pass
             self.set_authentication(provider, oauth2_refresh_token, username, password)
 
         self.set_api_endpoint("pgorelease.nianticlabs.com/plfe")
+        self._proxy = None
 
         self._position_lat = position_lat
         self._position_lng = position_lng
@@ -93,6 +93,9 @@ def set_position(self, lat, lng, alt):
         self._position_lng = lng
         self._position_alt = alt
 
+    def set_proxy(self, proxy_config):
+        self._proxy = proxy_config
+
     def get_api_endpoint(self):
         return self._api_endpoint
 
@@ -106,7 +109,8 @@ def get_auth_provider(self):
         return self._auth_provider
 
     def create_request(self):
-        request = PGoApiRequest(self, self._position_lat, self._position_lng, self._position_alt)
+        request = PGoApiRequest(self, self._position_lat, self._position_lng,
+                                self._position_alt, self._proxy)
         return request
 
     def activate_signature(self, lib_path):
@@ -177,7 +181,7 @@ def login(self, provider, username, password, lat=None, lng=None, alt=None, app_
 
 
 class PGoApiRequest:
-    def __init__(self, parent, position_lat, position_lng, position_alt):
+    def __init__(self, parent, position_lat, position_lng, position_alt, proxy=None):
         self.log = logging.getLogger(__name__)
 
         self.__parent__ = parent
@@ -190,6 +194,8 @@ def __init__(self, parent, position_lat, position_lng, position_alt):
         self._position_lng = position_lng
         self._position_alt = position_alt
 
+        self._proxy = proxy
+
         self._req_method_list = []
 
     def call(self):
@@ -203,7 +209,7 @@ def call(self):
             self.log.info('Not logged in')
             return NotLoggedInException()
 
-        request = RpcApi(self._auth_provider)
+        request = RpcApi(self._auth_provider, self._proxy)
 
         lib_path = self.__parent__.get_signature_lib()
         if lib_path is not None:
diff --git a/pgoapi/rpc_api.py b/pgoapi/rpc_api.py
index 2e12da9f..50f066b2 100755
--- a/pgoapi/rpc_api.py
+++ b/pgoapi/rpc_api.py
@@ -58,11 +58,15 @@ class RpcApi:
     RPC_ID = 0
     START_TIME = 0
 
-    def __init__(self, auth_provider):
+    def __init__(self, auth_provider, proxy_config=None):
 
         self.log = logging.getLogger(__name__)
 
         self._session = requests.session()
+
+        if proxy_config is not None:
+            self._session.proxies = proxy_config
+
         self._session.headers.update({'User-Agent': 'Niantic App'})
         self._session.verify = True
 
@@ -113,8 +117,8 @@ def _make_rpc(self, endpoint, request_proto_plain):
         request_proto_serialized = request_proto_plain.SerializeToString()
         try:
             http_response = self._session.post(endpoint, data=request_proto_serialized)
-        except requests.exceptions.ConnectionError as e:
-            raise ServerBusyOrOfflineException(e)
+        except requests.exceptions.ConnectionError:
+            raise ServerBusyOrOfflineException
 
         return http_response
 
diff --git a/requirements.txt b/requirements.txt
index 24078864..f9bd3286 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,4 +4,5 @@ requests==2.10.0
 s2sphere==0.2.4
 gpsoauth==0.3.0
 six
-xxhash
\ No newline at end of file
+xxhash
+requests[socks]