Encapsulate token-based auth in a class.

4 files changed, 111 insertions, 57 deletions

diff --git a/src/leap/soledad/auth.py b/src/leap/soledad/auth.py
index a099c1a6..1d8f1a42 100644
--- a/src/leap/soledad/auth.py
+++ b/src/leap/soledad/auth.py
@@ -27,32 +27,44 @@ they can do token-based auth requests to the Soledad server.
 from u1db.remote.http_client import HTTPClientBase
 
 
-def set_token_credentials(self, uuid, token):
+class TokenBasedAuth(object):
     """
-    Store given credentials so we can sign the request later.
-
-    @param uuid: The user's uuid.
-    @type uuid: str
-    @param token: The authentication token.
-    @type token: str
+    Encapsulate token-auth methods for classes that inherit from
+    u1db.remote.http_client.HTTPClient.
     """
-    self._creds = {'token': (uuid, token)}
 
+    def set_token_credentials(self, uuid, token):
+        """
+        Store given credentials so we can sign the request later.
 
-def _sign_request(self, method, url_query, params):
-    """
-    Return an authorization header to be included in the HTTP request.
-
-    @param method: The HTTP method.
-    @type method: str
-    @param url_query: The URL query string.
-    @type url_query: str
-    @param params: A list with encoded query parameters.
-    @type param: list
-    """
-    if 'token' in self._creds:
-        uuid, token = self._creds['token']
-        auth = '%s:%s' % (uuid, token)
-        return [('Authorization', 'Token %s' % auth.encode('base64')[:-1])]
-    else:
-        return HTTPClientBase._sign_request(self, method, url_query, params)
+        @param uuid: The user's uuid.
+        @type uuid: str
+        @param token: The authentication token.
+        @type token: str
+        """
+        self._creds = {'token': (uuid, token)}
+
+
+    def _sign_request(self, method, url_query, params):
+        """
+        Return an authorization header to be included in the HTTP request, in
+        the form:
+
+            [('Authorization', 'Token <base64 encoded creds')]
+
+        @param method: The HTTP method.
+        @type method: str
+        @param url_query: The URL query string.
+        @type url_query: str
+        @param params: A list with encoded query parameters.
+        @type param: list
+
+        @return: The Authorization header.
+        @rtype: list of tuple
+        """
+        if 'token' in self._creds:
+            uuid, token = self._creds['token']
+            auth = '%s:%s' % (uuid, token)
+            return [('Authorization', 'Token %s' % auth.encode('base64')[:-1])]
+        else:
+            return HTTPClientBase._sign_request(self, method, url_query, params)
diff --git a/src/leap/soledad/backends/leap_backend.py b/src/leap/soledad/backends/leap_backend.py
index 46c787a9..2585379a 100644
--- a/src/leap/soledad/backends/leap_backend.py
+++ b/src/leap/soledad/backends/leap_backend.py
@@ -35,10 +35,7 @@ from u1db.remote.http_target import HTTPSyncTarget
 
 from leap.common.keymanager import KeyManager
 from leap.common.check import leap_assert
-from leap.soledad.auth import (
-    set_token_credentials,
-    _sign_request,
-)
+from leap.soledad.auth import TokenBasedAuth
 
 #
 # Exceptions
@@ -249,7 +246,7 @@ class LeapDocument(Document):
 # LeapSyncTarget
 #
 
-class LeapSyncTarget(HTTPSyncTarget):
+class LeapSyncTarget(HTTPSyncTarget, TokenBasedAuth):
     """
     A SyncTarget that encrypts data before sending and decrypts data after
     receiving.
@@ -259,9 +256,32 @@ class LeapSyncTarget(HTTPSyncTarget):
     # Token auth methods.
     #
 
-    set_token_credentials = set_token_credentials
+    def set_token_credentials(self, uuid, token):
+        """
+        Store given credentials so we can sign the request later.
+
+        @param uuid: The user's uuid.
+        @type uuid: str
+        @param token: The authentication token.
+        @type token: str
+        """
+        TokenBasedAuth.set_token_credentials(self, uuid, token)
 
-    _sign_request = _sign_request
+    def _sign_request(self, method, url_query, params):
+        """
+        Return an authorization header to be included in the HTTP request.
+
+        @param method: The HTTP method.
+        @type method: str
+        @param url_query: The URL query string.
+        @type url_query: str
+        @param params: A list with encoded query parameters.
+        @type param: list
+
+        @return: The Authorization header.
+        @rtype: list of tuple
+        """
+        return TokenBasedAuth._sign_request(self, method, url_query, params)
 
     #
     # Modified HTTPSyncTarget methods.
diff --git a/src/leap/soledad/shared_db.py b/src/leap/soledad/shared_db.py
index 419d8017..3929e828 100644
--- a/src/leap/soledad/shared_db.py
+++ b/src/leap/soledad/shared_db.py
@@ -29,10 +29,7 @@ except ImportError:
 from u1db.remote import http_database, http_client
 
 
-from leap.soledad.auth import (
-    set_token_credentials,
-    _sign_request,
-)
+from leap.soledad.auth import TokenBasedAuth
 
 SOLEDAD_CERT = None
 
@@ -78,7 +75,7 @@ class Unauthorized(Exception):
     """
 
 
-class SoledadSharedDatabase(http_database.HTTPDatabase):
+class SoledadSharedDatabase(http_database.HTTPDatabase, TokenBasedAuth):
     """
     This is a shared recovery database that enables users to store their
     encryption secrets in the server and retrieve them afterwards.
@@ -90,9 +87,32 @@ class SoledadSharedDatabase(http_database.HTTPDatabase):
     # Token auth methods.
     #
 
-    set_token_credentials = set_token_credentials
+    def set_token_credentials(self, uuid, token):
+        """
+        Store given credentials so we can sign the request later.
+
+        @param uuid: The user's uuid.
+        @type uuid: str
+        @param token: The authentication token.
+        @type token: str
+        """
+        TokenBasedAuth.set_token_credentials(self, uuid, token)
 
-    _sign_request = _sign_request
+    def _sign_request(self, method, url_query, params):
+        """
+        Return an authorization header to be included in the HTTP request.
+
+        @param method: The HTTP method.
+        @type method: str
+        @param url_query: The URL query string.
+        @type url_query: str
+        @param params: A list with encoded query parameters.
+        @type param: list
+
+        @return: The Authorization header.
+        @rtype: list of tuple
+        """
+        return TokenBasedAuth._sign_request(self, method, url_query, params)
 
     #
     # Modified HTTPDatabase methods.
diff --git a/src/leap/soledad/tests/test_leap_backend.py b/src/leap/soledad/tests/test_leap_backend.py
index 2053bb33..b0e0aaec 100644
--- a/src/leap/soledad/tests/test_leap_backend.py
+++ b/src/leap/soledad/tests/test_leap_backend.py
@@ -95,14 +95,20 @@ LEAP_SCENARIOS = [
 
 
 def make_token_http_database_for_test(test, replica_uid):
-    http_db = test_backends.make_http_database_for_test(test, replica_uid, 'test')
-    http_db.set_token_credentials = auth.set_token_credentials
+    test.startServer()
+    test.request_state._create_database(replica_uid)
 
-    def _sign_request(method, url_query, params):
-        return auth._sign_request(http_db, method, url_query, params)
+    class _HTTPDatabaseWithToken(
+            http_database.HTTPDatabase, auth.TokenBasedAuth):
 
-    http_db._sign_request = _sign_request
-    http_db.set_token_credentials(http_db, 'user-uuid', 'auth-token')
+        def set_token_credentials(self, uuid, token):
+            auth.TokenBasedAuth.set_token_credentials(self, uuid, token)
+
+        def _sign_request(self, method, url_query, params):
+            return auth.TokenBasedAuth._sign_request(self, method, url_query, params)
+
+    http_db = _HTTPDatabaseWithToken(test.getURL('test'))
+    http_db.set_token_credentials('user-uuid', 'auth-token')
     return http_db
 
 
@@ -113,12 +119,6 @@ def copy_token_http_database_for_test(test, db):
     # CORRUPT USER DATA. USE SYNC INSTEAD, OR WE WILL SEND NINJA TO YOUR
     # HOUSE.
     http_db = test.request_state._copy_database(db)
-    http_db.set_token_credentials = auth.set_token_credentials
-
-    def _sign_request(method, url_query, params):
-        return auth._sign_request(http_db, method, url_query, params)
-
-    http_db._sign_request = _sign_request
     http_db.set_token_credentials(http_db, 'user-uuid', 'auth-token')
     return http_db
 
@@ -154,13 +154,14 @@ class TestLeapClientBase(test_http_client.TestHTTPClientBase):
     def getClientWithToken(self, **kwds):
         self.startServer()
 
-        class _HTTPClientWithToken(http_client.HTTPClientBase):
+        class _HTTPClientWithToken(
+                http_client.HTTPClientBase, auth.TokenBasedAuth):
 
             def set_token_credentials(self, uuid, token):
-                auth.set_token_credentials(self, uuid, token)
+                auth.TokenBasedAuth.set_token_credentials(self, uuid, token)
 
             def _sign_request(self, method, url_query, params):
-                return auth._sign_request(self, method, url_query, params)
+                return auth.TokenBasedAuth._sign_request(self, method, url_query, params)
 
         return _HTTPClientWithToken(self.getURL('dbase'), **kwds)
 
@@ -546,16 +547,17 @@ def token_leap_https_sync_target(test, host, path):
 # The following tests come from `u1db.tests.test_http_database`.
 #-----------------------------------------------------------------------------
 
-class _HTTPDatabase(http_database.HTTPDatabase):
+class _HTTPDatabase(http_database.HTTPDatabase, auth.TokenBasedAuth):
     """
     Wraps our token auth implementation.
     """
 
     def set_token_credentials(self, uuid, token):
-        auth.set_token_credentials(self, uuid, token)
+        auth.TokenBasedAuth.set_token_credentials(self, uuid, token)
 
     def _sign_request(self, method, url_query, params):
-        return auth._sign_request(self, method, url_query, params)
+        return auth.TokenBasedAuth._sign_request(
+            self, method, url_query, params)
 
 
 class TestHTTPDatabaseWithCreds(test_http_database.TestHTTPDatabaseCtrWithCreds):