Merge pull request #54 from smartystreets/eric/x-forwarded-for

Eric/x forwarded for

Author: RyanLCox1

smartystreets_python_sdk/client_builder.py

@@ -23,6 +23,7 @@ def __init__(self, signer):
         self.max_timeout = 10
         self.url_prefix = None
         self.proxy = None
+        self.ip = None
         self.debug = None
         self.header = None
         self.licenses = []
@@ -114,6 +115,15 @@ def with_custom_header(self, custom_header):
         """
         self.header = custom_header
         return self
+    
+    def with_x_forwarded_for(self, ip):
+        """
+        Add and X-Forwarded-For header when necessary.
+        :param ip: Input the desired ip for the X-Forwarded-For header
+        :return: Returns self to accommodate method chaining
+        """
+        self.ip = ip
+        return self
 
     def with_debug(self):
         """
@@ -170,7 +180,7 @@ def build_sender(self):
         if self.http_sender is not None:
             return self.http_sender
 
-        sender = smarty.RequestsSender(self.max_timeout, self.proxy)
+        sender = smarty.RequestsSender(self.max_timeout, self.proxy, self.ip)
         sender.debug = self.debug
 
         sender = smarty.StatusCodeSender(sender)

smartystreets_python_sdk/requests_sender.py

@@ -5,14 +5,16 @@
 
 
 class RequestsSender:
-    def __init__(self, max_timeout=None, proxy=None):
+    def __init__(self, max_timeout=None, proxy=None, ip=None):
         self.session = Session()
         self.max_timeout = max_timeout or 10
         self.proxy = proxy
         self.debug = None
+        self.ip = ip
 
     def send(self, smarty_request):
-        request = build_request(smarty_request)
+        ip = self.ip
+        request = build_request(smarty_request, ip)
         prepped_request = self.session.prepare_request(request)
         prepped_proxies = self.build_proxies()
         if self.debug:
@@ -43,13 +45,15 @@ def build_proxies(self):
             return {'http': proxy_string, 'https': proxy_string}
 
 
-def build_request(smarty_request):
+def build_request(smarty_request, ip=None):
     try:
         request = Request(url=smarty_request.url_prefix, params=smarty_request.parameters)
         request.headers['User-Agent'] = "smartystreets (sdk:python@{})".format(version.__version__)
         request.headers['Content-Type'] = smarty_request.content_type
         if smarty_request.referer:
             request.headers['Referer'] = smarty_request.referer
+        if ip != None:
+            request.headers['X-Forwarded-For'] = ip
         if smarty_request.payload:
             request.data = smarty_request.payload
             request.method = 'POST'

test/x_forwarded_for_test.py

@@ -0,0 +1,54 @@
+import smartystreets_python_sdk as smarty
+import unittest
+from mock import patch
+
+
+def mocked_session_send(request, **kwargs):
+    class MockResponse:
+        def __init__(self, payload, status_code):
+            self.text = payload
+            self.status_code = status_code
+            self.headers = None
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, type, value, traceback):
+            pass
+
+        def json(self):
+            return self.text
+
+    mockresponse = MockResponse("This is the test payload.", 200)
+
+    return mockresponse
+
+
+class TestCustomHeaderSender(unittest.TestCase):
+
+    def test_populates_header(self):
+        sender = smarty.RequestsSender(ip = '0.0.0.0')
+
+        self.assertEqual(sender.ip, '0.0.0.0')
+
+    @patch('requests.Session.send', side_effect=mocked_session_send)
+    def test_x_forwarded_for_header_set(self, mock_send):
+        sender = smarty.RequestsSender(ip = '0.0.0.0')
+        smartyrequest = smarty.Request()
+        smartyrequest.url_prefix = "http://localhost"
+        smartyrequest.payload = "This is the test content."
+
+        request = smarty.requests_sender.build_request(smartyrequest, '0.0.0.0')
+
+        self.assertEqual('0.0.0.0', request.headers['X-Forwarded-For'])
+
+    @patch('requests.Session.send', side_effect=mocked_session_send)
+    def test_custom_headers_used(self, mock_send):
+        sender = smarty.RequestsSender(ip = '0.0.0.0')
+        smartyrequest = smarty.Request()
+        smartyrequest.url_prefix = "http://localhost"
+        smartyrequest.payload = "This is the test content."
+
+        request = smarty.requests_sender.build_request(smartyrequest, '0.0.0.0')
+
+        self.assertEqual('0.0.0.0', request.headers['X-Forwarded-For'])