From 02b128a12831f4c43fb734dd1e1ca08deb66142c Mon Sep 17 00:00:00 2001
From: signop <39252060+signop@users.noreply.github.com>
Date: Mon, 14 May 2018 14:43:43 -0700
Subject: [PATCH] Add support for specifying source_address to ClientNetwork.
 (#5990)

For certbot/certbot#3489
---
 acme/acme/client.py      | 14 +++++++++++++-
 acme/acme/client_test.py | 25 +++++++++++++++++++++++++
 acme/setup.py            |  1 +
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/acme/acme/client.py b/acme/acme/client.py
index 7807f0ece..d4f5a4787 100644
--- a/acme/acme/client.py
+++ b/acme/acme/client.py
@@ -12,7 +12,9 @@ from six.moves import http_client  # pylint: disable=import-error
 import josepy as jose
 import OpenSSL
 import re
+from requests_toolbelt.adapters.source import SourceAddressAdapter
 import requests
+from requests.adapters import HTTPAdapter
 import sys
 
 from acme import crypto_util
@@ -857,9 +859,12 @@ class ClientNetwork(object):  # pylint: disable=too-many-instance-attributes
     :param bool verify_ssl: Whether to verify certificates on SSL connections.
     :param str user_agent: String to send as User-Agent header.
     :param float timeout: Timeout for requests.
+    :param source_address: Optional source address to bind to when making requests.
+    :type source_address: str or tuple(str, int)
     """
     def __init__(self, key, account=None, alg=jose.RS256, verify_ssl=True,
-                 user_agent='acme-python', timeout=DEFAULT_NETWORK_TIMEOUT):
+                 user_agent='acme-python', timeout=DEFAULT_NETWORK_TIMEOUT,
+                 source_address=None):
         # pylint: disable=too-many-arguments
         self.key = key
         self.account = account
@@ -869,6 +874,13 @@ class ClientNetwork(object):  # pylint: disable=too-many-instance-attributes
         self.user_agent = user_agent
         self.session = requests.Session()
         self._default_timeout = timeout
+        adapter = HTTPAdapter()
+
+        if source_address is not None:
+            adapter = SourceAddressAdapter(source_address)
+
+        self.session.mount("http://", adapter)
+        self.session.mount("https://", adapter)
 
     def __del__(self):
         # Try to close the session, but don't show exceptions to the
diff --git a/acme/acme/client_test.py b/acme/acme/client_test.py
index c17b83210..f3018ed81 100644
--- a/acme/acme/client_test.py
+++ b/acme/acme/client_test.py
@@ -1129,6 +1129,31 @@ class ClientNetworkWithMockedResponseTest(unittest.TestCase):
         self.assertRaises(requests.exceptions.RequestException,
                           self.net.post, 'uri', obj=self.obj)
 
+class ClientNetworkSourceAddressBindingTest(unittest.TestCase):
+    """Tests that if ClientNetwork has a source IP set manually, the underlying library has
+    used the provided source address."""
+
+    def setUp(self):
+        self.source_address = "8.8.8.8"
+
+    def test_source_address_set(self):
+        from acme.client import ClientNetwork
+        net = ClientNetwork(key=None, alg=None, source_address=self.source_address)
+        for adapter in net.session.adapters.values():
+            self.assertTrue(self.source_address in adapter.source_address)
+
+    def test_behavior_assumption(self):
+        """This is a test that guardrails the HTTPAdapter behavior so that if the default for
+        a Session() changes, the assumptions here aren't violated silently."""
+        from acme.client import ClientNetwork
+        # Source address not specified, so the default adapter type should be bound -- this
+        # test should fail if the default adapter type is changed by requests
+        net = ClientNetwork(key=None, alg=None)
+        session = requests.Session()
+        for scheme in session.adapters.keys():
+            client_network_adapter = net.session.adapters.get(scheme)
+            default_adapter = session.adapters.get(scheme)
+            self.assertEqual(client_network_adapter.__class__, default_adapter.__class__)
 
 if __name__ == '__main__':
     unittest.main()  # pragma: no cover
diff --git a/acme/setup.py b/acme/setup.py
index 72ab5919b..e91c36b3d 100644
--- a/acme/setup.py
+++ b/acme/setup.py
@@ -19,6 +19,7 @@ install_requires = [
     'pyrfc3339',
     'pytz',
     'requests[security]>=2.4.1',  # security extras added in 2.4.1
+    'requests-toolbelt>=0.3.0',
     'setuptools',
     'six>=1.9.0',  # needed for python_2_unicode_compatible
 ]