diff --git a/smartystreets_python_sdk/client_builder.py b/smartystreets_python_sdk/client_builder.py index 5e91304..e0ba02a 100644 --- a/smartystreets_python_sdk/client_builder.py +++ b/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) diff --git a/smartystreets_python_sdk/requests_sender.py b/smartystreets_python_sdk/requests_sender.py index 85c195a..23e5495 100644 --- a/smartystreets_python_sdk/requests_sender.py +++ b/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' diff --git a/test/x_forwarded_for_test.py b/test/x_forwarded_for_test.py new file mode 100644 index 0000000..d19d3ff --- /dev/null +++ b/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']) \ No newline at end of file