Allow custom authorization mechanisms.

docs/examples/custom_auth.py

@@ -0,0 +1,74 @@
+from __future__ import print_function
+from twisted.internet.endpoints import serverFromString
+from twisted.internet.task import react
+from twisted.internet.defer import inlineCallbacks, maybeDeferred
+from twisted.web.resource import Resource
+from twisted.web.server import Site
+from twisted.web.http_headers import Headers
+from twisted.web.http import FORBIDDEN
+
+import treq
+
+
+class SillyAuthResource(Resource):
+    """
+    A resource that uses a silly, header-based authentication
+    mechanism.
+    """
+    isLeaf = True
+
+    def __init__(self, header, secret):
+        self._header = header
+        self._secret = secret
+
+    def render_GET(self, request):
+        headers = request.requestHeaders
+        request_secret = headers.getRawHeaders(self._header, [b''])[0]
+        if request_secret != self._secret:
+            request.setResponseCode(FORBIDDEN)
+            return b"No good."
+        return b"It's good!"
+
+
+class SillyAuth(object):
+    """
+    I implement a silly, header-based authentication mechanism.
+    """
+
+    def __init__(self, header, secret, agent):
+        self._header = header
+        self._secret = secret
+        self._agent = agent
+
+    def request(self, method, uri, headers=None, bodyProducer=None):
+        if headers is None:
+            headers = Headers({})
+        headers.setRawHeaders(self._header, [self._secret])
+        return self._agent.request(method, uri, headers, bodyProducer)
+
+
+@inlineCallbacks
+def main(reactor, *args):
+    header = b'x-silly-auth'
+    secret = b'secret'
+
+    auth_resource = SillyAuthResource(header=header, secret=secret)
+
+    endpoint = serverFromString(reactor, "tcp:8080")
+    listener = yield endpoint.listen(Site(auth_resource))
+
+    def sillyAuthCallable(agent):
+        return SillyAuth(header, secret, agent)
+
+    response = yield treq.get(
+        'http://localhost:8080/',
+        auth=sillyAuthCallable,
+    )
+
+    content = yield response.content()
+    print(content)
+
+    yield maybeDeferred(listener.stopListening)
+
+
+react(main, [])

docs/howto.rst

@@ -54,6 +54,18 @@ Full example: :download:`basic_auth.py <examples/basic_auth.py>`
 
 .. _RFC 2617: http://www.ietf.org/rfc/rfc2617.txt
 
+The ``auth`` keyword argument also supports custom authorization
+implementations.  ``auth`` should be a callable that accepts as its
+only argument the currently configured ``IAgent`` and should return an
+``IAgent`` that implements the desired authorization mechanism.
+
+.. literalinclude:: examples/custom_auth.py
+    :linenos:
+    :lines: 10-33
+
+Full example: :download:`basic_auth.py <examples/custom_auth.py>`
+
+
 Redirects
 ---------
 

treq/auth.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 from twisted.web.http_headers import Headers
+from six.moves.urllib.parse import urlparse
 import base64
 
 
@@ -26,6 +27,31 @@ def request(self, method, uri, headers=None, bodyProducer=None):
             method, uri, headers=headers, bodyProducer=bodyProducer)
 
 
+class _PinToFirstHostAgent(object):
+    """
+    An {twisted.web.iweb.IAgent} implementing object that takes two
+    agents, using the first as when the current request's host name
+    matches first request's host name, and the second when it does
+    not.
+    """
+
+    def __init__(self, first_agent, second_agent):
+        self._first_agent = first_agent
+        self._second_agent = second_agent
+        self._first_host = None
+
+    def request(self, method, uri, headers=None, bodyProducer=None):
+        hostname = urlparse(uri).hostname
+        if self._first_host in (None, hostname):
+            self._first_host = hostname
+            agent = self._first_agent
+        else:
+            agent = self._second_agent
+
+        return agent.request(
+            method, uri, headers=headers, bodyProducer=bodyProducer)
+
+
 def add_basic_auth(agent, username, password):
     creds = base64.b64encode(
         '{0}:{1}'.format(username, password).encode('ascii'))
@@ -37,5 +63,7 @@ def add_basic_auth(agent, username, password):
 def add_auth(agent, auth_config):
     if isinstance(auth_config, tuple):
         return add_basic_auth(agent, auth_config[0], auth_config[1])
+    elif callable(auth_config):
+        return auth_config(agent)
 
     raise UnknownAuthConfig(auth_config)

treq/client.py

@@ -30,7 +30,7 @@
 from twisted.python.components import registerAdapter
 
 from treq._utils import default_reactor
-from treq.auth import add_auth
+from treq.auth import add_auth, _PinToFirstHostAgent
 from treq import multipart
 from treq.response import _Response
 from requests.cookies import cookiejar_from_dict, merge_cookies
@@ -203,6 +203,11 @@ def request(self, method, url, **kwargs):
         cookies = merge_cookies(self._cookiejar, cookies)
         wrapped_agent = CookieAgent(self._agent, cookies)
 
+        auth = kwargs.get('auth')
+        if auth:
+            auth_agent = add_auth(wrapped_agent, auth)
+            wrapped_agent = _PinToFirstHostAgent(auth_agent, wrapped_agent)
+
         if kwargs.get('allow_redirects', True):
             if kwargs.get('browser_like_redirects', False):
                 wrapped_agent = BrowserLikeRedirectAgent(wrapped_agent)
@@ -212,10 +217,6 @@ def request(self, method, url, **kwargs):
         wrapped_agent = ContentDecoderAgent(wrapped_agent,
                                             [(b'gzip', GzipDecoder)])
 
-        auth = kwargs.get('auth')
-        if auth:
-            wrapped_agent = add_auth(wrapped_agent, auth)
-
         d = wrapped_agent.request(
             method, url, headers=headers,
             bodyProducer=bodyProducer)

treq/test/test_auth.py

@@ -2,9 +2,15 @@
 
 from twisted.web.client import Agent
 from twisted.web.http_headers import Headers
+from twisted.web.iweb import IAgent
 
 from treq.test.util import TestCase
-from treq.auth import _RequestHeaderSettingAgent, add_auth, UnknownAuthConfig
+from treq.auth import (_RequestHeaderSettingAgent,
+                       _PinToFirstHostAgent,
+                       add_auth,
+                       UnknownAuthConfig)
+
+from zope.interface import implementer
 
 
 class RequestHeaderSettingAgentTests(TestCase):
@@ -42,6 +48,61 @@ def test_overrides_per_request_headers(self):
         )
 
 
+class PinToFirstHostAgentTests(TestCase):
+    def setUp(self):
+        self.first_agent = mock.Mock(Agent)
+        self.second_agent = mock.Mock(Agent)
+        self.pinning_agent = _PinToFirstHostAgent(self.first_agent,
+                                                  self.second_agent)
+
+    def test_first_request_uses_first_agent(self):
+        self.pinning_agent.request("method", "http://www.something.com/")
+
+        self.first_agent.request.assert_called_once_with(
+            "method", "http://www.something.com/",
+            headers=None, bodyProducer=None)
+        self.assertFalse(self.second_agent.request.called)
+
+    def test_request_matching_first_uses_first_agent(self):
+        self.pinning_agent.request("method", "http://www.something.com/a")
+        self.pinning_agent.request("method", "http://www.something.com/b")
+
+        self.assertEqual(self.first_agent.request.call_args_list, [
+            mock.call("method", "http://www.something.com/a",
+                      headers=None, bodyProducer=None),
+            mock.call("method", "http://www.something.com/b",
+                      headers=None, bodyProducer=None)
+        ])
+        self.assertFalse(self.second_agent.request.called)
+
+    def test_request_not_matching_first_uses_second_agent(self):
+        self.pinning_agent.request("method", "http://www.something.com/a")
+        self.pinning_agent.request("method", "http://www.other.com/a")
+
+        self.first_agent.request.assert_called_once_with(
+            "method", "http://www.something.com/a",
+            headers=None, bodyProducer=None)
+
+        self.second_agent.request.assert_called_once_with(
+            "method", "http://www.other.com/a",
+            headers=None, bodyProducer=None)
+
+    def test_some_request_matching_first_uses_second_agent(self):
+        self.pinning_agent.request("method", "http://www.something.com/a")
+        self.pinning_agent.request("method", "http://www.other.com/a")
+        self.pinning_agent.request("method", "http://www.something.com/b")
+
+        self.assertEqual(self.first_agent.request.call_args_list, [
+            mock.call("method", "http://www.something.com/a",
+                      headers=None, bodyProducer=None),
+            mock.call("method", "http://www.something.com/b",
+                      headers=None, bodyProducer=None)
+        ])
+        self.second_agent.request.assert_called_once_with(
+            "method", "http://www.other.com/a",
+            headers=None, bodyProducer=None)
+
+
 class AddAuthTests(TestCase):
     def setUp(self):
         self.rhsa_patcher = mock.patch('treq.auth._RequestHeaderSettingAgent')
@@ -70,6 +131,22 @@ def test_add_basic_auth_huge(self):
             agent,
             Headers({b'authorization': [auth]}))
 
+    def test_auth_callable(self):
+        agent = mock.Mock()
+
+        @implementer(IAgent)
+        class AuthorizingAgent(object):
+
+            def __init__(self, agent):
+                self.wrapped_agent = agent
+
+            def request(self, method, uri, headers=None, bodyProducer=None):
+                """Not called by this test"""
+
+        wrapping_agent = add_auth(agent, AuthorizingAgent)
+        self.assertIsInstance(wrapping_agent, AuthorizingAgent)
+        self.assertIs(wrapping_agent.wrapped_agent, agent)
+
     def test_add_unknown_auth(self):
         agent = mock.Mock()
-        self.assertRaises(UnknownAuthConfig, add_auth, agent, mock.Mock())
+        self.assertRaises(UnknownAuthConfig, add_auth, agent, object())