etcd.auth.AuthClient

This extension affords create/read/update without cluttering the basic
etcd.Client implementation.

The model is reworked for a cleaner API: user's roles can be assigned via
list/tuple, permissions are moddeled like a dictionary.

Adding coverage goal to buildout to verify testing progress.

4 files changed, 578 insertions, 0 deletions

diff --git a/buildout.cfg b/buildout.cfg
index cba64c5..4de9036 100644
--- a/buildout.cfg
+++ b/buildout.cfg
@@ -2,6 +2,7 @@
 parts = python
       sphinxbuilder
       test
+      coverage
 develop = .
 eggs = 
      urllib3==1.7.1
@@ -18,6 +19,13 @@ recipe = pbp.recipe.noserunner
 eggs = ${python:eggs}
      mock
 
+[coverage]
+recipe = pbp.recipe.noserunner
+eggs = ${test:eggs}
+     coverage
+defaults = --with-coverage
+     --cover-package=etcd
+
 [sphinxbuilder]
 recipe = collective.recipe.sphinxbuilder
 source = ${buildout:directory}/docs-source
diff --git a/src/etcd/auth.py b/src/etcd/auth.py
new file mode 100644
index 0000000..0f8a2dd
--- /dev/null
+++ b/src/etcd/auth.py
@@ -0,0 +1,369 @@
+import json
+
+import logging
+
+try:
+    # Python 3
+    from http.client import HTTPException
+except ImportError:
+    # Python 2
+    from httplib import HTTPException
+import socket
+import urllib3
+
+from .client import Client
+import etcd
+
+_log = logging.getLogger(__name__)
+
+
+class AuthClient(Client):
+    """
+    Extended etcd client that supports authentication primitives added in 2.1.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(AuthClient, self).__init__(*args, **kwargs)
+
+    def create_user(self, username, password, roles=[], role_action='roles'):
+        """
+        Add a user.
+
+        Args:
+            username (str): Username to create.
+            password (str): Password for username.
+            roles (list): List of roles as strings.
+
+        Returns:
+            EtcdUser
+
+        Raises:
+            etcd.EtcdException: If user can't be created.
+        """
+        try:
+            uri = self.version_prefix + '/auth/users/' + username
+            params = {'user': username}
+            if password:
+                params['password'] = password
+            if roles:
+                params[role_action] = roles
+
+            response = self.json_api_execute(uri, self._MPUT, params=params)
+            res = json.loads(response.data.decode('utf-8'))
+            return EtcdUser(self, res)
+        except Exception as e:
+            _log.error("Failed to create user in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not create user")
+
+    def get_user(self, username):
+        """
+        Look up a user.
+
+        Args:
+            username (str): Username to lookup.
+
+        Returns:
+            EtcdUser
+
+        Raises:
+            etcd.EtcdException: If user can't be found.
+        """
+        try:
+            uri = self.version_prefix + '/auth/users/' + username
+            response = self.api_execute(uri, self._MGET)
+            res = json.loads(response.data.decode('utf-8'))
+            return EtcdUser(self, res)
+        except Exception as e:
+            _log.error("Failed to fetch user in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not fetch user")
+
+    @property
+    def usernames(self):
+        """List user names."""
+        try:
+            uri = self.version_prefix + '/auth/users'
+            response = self.api_execute(uri, self._MGET)
+            res = json.loads(response.data.decode('utf-8'))
+            return res['users']
+        except Exception as e:
+            _log.error("Failed to list users in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not list users")
+
+    @property
+    def users(self):
+        """List users in detail."""
+        return [self.get_user(x) for x in self.usernames]
+
+    def create_role(self, role_name):
+        """
+        Create a role.
+
+        Args:
+            role_name (str): Name of role
+
+        Returns:
+            EtcdRole
+        """
+        return self.modify_role(role_name)
+
+    def get_role(self, role_name):
+        """
+        Look up a role.
+
+        Args:
+            role_name (str): Name of role.
+
+        Returns:
+            EtcdRole
+        """
+        try:
+            uri = self.version_prefix + '/auth/roles/' + role_name
+            response = self.api_execute(uri, self._MGET)
+            res = json.loads(response.data.decode('utf-8'))
+            return EtcdRole(self, res)
+        except Exception as e:
+            _log.error("Failed to fetch user in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not fetch users")
+
+    @property
+    def role_names(self):
+        """List role names."""
+        try:
+            uri = self.version_prefix + '/auth/roles'
+            response = self.api_execute(uri, self._MGET)
+            res = json.loads(response.data.decode('utf-8'))
+            return res['roles']
+        except Exception as e:
+            _log.error("Failed to list roles in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not list roles")
+
+    @property
+    def roles(self):
+        """List roles in detail."""
+        return [self.get_role(x) for x in self.role_names]
+
+    def toggle_auth(self, auth_enabled=True):
+        """
+        Toggle authentication.
+
+        Args:
+            auth_enabled (bool): Should auth be enabled or disabled
+        """
+        try:
+            uri = self.version_prefix + '/auth/enable'
+            action = auth_enabled and self._MPUT or self._MDELETE
+
+            self.api_execute(uri, action)
+        except Exception as e:
+            _log.error("Failed enable authentication in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not toggle authentication")
+
+    def modify_role(self, role_name, permissions=None, perm_key=None):
+        """Modifies role."""
+        try:
+            uri = self.version_prefix + '/auth/roles/' + role_name
+            params = {
+                'role': role_name,
+            }
+            if permissions:
+                params[perm_key] = {
+                    'kv': {
+                        'read': [k for k, v in permissions.items() if
+                                 'R' in v.upper()],
+                        'write': [k for k, v in permissions.items() if
+                                  'W' in v.upper()]
+                    }
+                }
+            response = self.json_api_execute(uri, self._MPUT, params=params)
+            res = json.loads(response.data.decode('utf-8'))
+            return EtcdRole(self, res)
+        except Exception as e:
+            _log.error("Failed to modify role in %s%s: %r",
+                       self._base_uri, self.version_prefix, e)
+            raise etcd.EtcdException("Could not modify role")
+
+    def json_api_execute(self, path, method, params=None, timeout=None):
+        """ Executes the query. """
+
+        some_request_failed = False
+        response = False
+
+        if timeout is None:
+            timeout = self.read_timeout
+
+        if timeout == 0:
+            timeout = None
+
+        if not path.startswith('/'):
+            raise ValueError('Path does not start with /')
+
+        while not response:
+            try:
+                url = self._base_uri + path
+                json_payload = json.dumps(params)
+                headers = self._get_headers()
+                headers['Content-Type'] = 'application/json'
+                response = self.http.urlopen(method,
+                                             url,
+                                             body=json_payload,
+                                             timeout=timeout,
+                                             redirect=self.allow_redirect,
+                                             headers=headers,
+                                             preload_content=False)
+            # urllib3 doesn't wrap all httplib exceptions and earlier versions
+            # don't wrap socket errors either.
+            except (urllib3.exceptions.HTTPError,
+                    HTTPException,
+                    socket.error) as e:
+                _log.error("Request to server %s failed: %r",
+                           self._base_uri, e)
+                if self._allow_reconnect:
+                    _log.info("Reconnection allowed, looking for another "
+                              "server.")
+                    # _next_server() raises EtcdException if there are no
+                    # machines left to try, breaking out of the loop.
+                    self._base_uri = self._next_server()
+                    some_request_failed = True
+                else:
+                    _log.debug("Reconnection disabled, giving up.")
+                    raise etcd.EtcdConnectionFailed(
+                        "Connection to etcd failed due to %r" % e)
+            except:
+                _log.exception("Unexpected request failure, re-raising.")
+                raise
+
+            else:
+                # Check the cluster ID hasn't changed under us.  We use
+                # preload_content=False above so we can read the headers
+                # before we wait for the content of a long poll.
+                cluster_id = response.getheader("x-etcd-cluster-id")
+                id_changed = (self.expected_cluster_id
+                              and cluster_id is not None and
+                              cluster_id != self.expected_cluster_id)
+                # Update the ID so we only raise the exception once.
+                old_expected_cluster_id = self.expected_cluster_id
+                self.expected_cluster_id = cluster_id
+                if id_changed:
+                    # Defensive: clear the pool so that we connect afresh next
+                    # time.
+                    self.http.clear()
+                    raise etcd.EtcdClusterIdChanged(
+                        'The UUID of the cluster changed from {} to '
+                        '{}.'.format(old_expected_cluster_id, cluster_id))
+
+        if some_request_failed:
+            if not self._use_proxies:
+                # The cluster may have changed since last invocation
+                self._machines_cache = self.machines
+            self._machines_cache.remove(self._base_uri)
+        return self._handle_server_response(response)
+
+
+class EtcdUser(object):
+    def __init__(self, auth_client, json_user):
+        self.client = auth_client
+        self.name = json_user.get('user')
+        self._roles = json_user.get('roles') or []
+
+    @property
+    def password(self):
+        """Empty property for password."""
+        return None
+
+    @password.setter
+    def password(self, new_password):
+        """Change user's password."""
+        self.client.create_user(self.name, new_password)
+
+    @property
+    def roles(self):
+        return tuple(self._roles)
+
+    @roles.setter
+    def roles(self, roles):
+        existing_roles = set(self._roles)
+        new_roles = set(roles)
+
+        if existing_roles == new_roles:
+            _log.debug('User %s already belongs to %s', self.name, self._roles)
+            return
+
+        to_revoke = existing_roles - new_roles
+        to_grant = new_roles - existing_roles
+
+        if to_revoke:
+            self.client.create_user(self.name, None, roles=list(to_revoke),
+                                    role_action='revoke')
+        if to_grant:
+            self.client.create_user(self.name, None, roles=list(to_grant),
+                                    role_action='grant')
+        self._roles = new_roles
+
+
+class EtcdRole(object):
+    def __init__(self, auth_client, role_json):
+        self.client = auth_client
+        self.name = role_json.get('role')
+        self.permissions = RolePermissionsDict(self, role_json)
+
+
+class RolePermissionsDict(dict):
+    _PERMISSIONS = {'R', 'W'}
+
+    def __init__(self, etcd_role, role_json, *args, **kwargs):
+        super(RolePermissionsDict, self).__init__(*args, **kwargs)
+        self.role = etcd_role
+        permissions = role_json.get('permissions')
+        if permissions and 'kv' in permissions:
+            self.__add_permissions(permissions, 'read', 'R')
+            self.__add_permissions(permissions, 'write', 'W')
+
+    def __add_permissions(self, permissions, label, symbol):
+        if label in permissions['kv'] and permissions['kv'][label]:
+            for path in permissions['kv'][label]:
+                existing_perms = dict.get(self, path)
+                if existing_perms:
+                    dict.__setitem__(self, path,
+                                     existing_perms + symbol)
+                else:
+                    dict.__setitem__(self, path, symbol)
+
+    def __setitem__(self, key, value):
+        if not value:
+            raise ValueError('Permissions may only be (R)ead or (W)ite')
+        perms = set(x.upper() for x in value)
+        if not perms <= RolePermissionsDict._PERMISSIONS:
+            raise ValueError('Permissions may only be (R)ead or (W)ite')
+
+        role_name = self.role.name
+        perm_dict = {key: value}
+        existing_value = dict.get(self, key)
+
+        if existing_value:
+            existing_perms = set(x.upper() for x in existing_value)
+            if perms != existing_perms:
+                to_grant = perms - existing_perms
+                to_revoke = existing_perms - perms
+
+                if to_revoke:
+                    perm_dict = {key: ''.join(to_revoke)}
+                    self.role.client.modify_role(role_name, perm_dict, 'revoke')
+                if to_grant:
+                    perm_dict = {key: ''.join(to_grant)}
+                    self.role.client.modify_role(role_name, perm_dict, 'grant')
+            else:
+                _log.debug('Permission %s=%s already granted', key, value)
+        else:
+            self.role.client.modify_role(role_name, perm_dict, 'grant')
+
+        dict.__setitem__(self, key, value)
+
+    def __delitem__(self, key):
+        self.role.client.modify_role(self.role.name, {key: 'RW'}, 'revoke')
+        dict.__delitem__(self, key)
diff --git a/src/etcd/tests/integration/helpers.py b/src/etcd/tests/integration/helpers.py
index 3314be9..1f1d22b 100644
--- a/src/etcd/tests/integration/helpers.py
+++ b/src/etcd/tests/integration/helpers.py
@@ -38,6 +38,12 @@ class EtcdProcessHelper(object):
                 '-initial-cluster', initial_cluster,
                 '-initial-cluster-state', 'new'
             ])
+        else:
+            proc_args.extend([
+                '-initial-cluster', 'test-node-0=http://127.0.0.1:{}'.format(self.internal_port_range_start),
+                '-initial-cluster-state', 'new'
+            ])
+
         for i in range(0, number):
             self.add_one(i, proc_args)
 
diff --git a/src/etcd/tests/integration/test_authentication.py b/src/etcd/tests/integration/test_authentication.py
new file mode 100644
index 0000000..52ba001
--- /dev/null
+++ b/src/etcd/tests/integration/test_authentication.py
@@ -0,0 +1,195 @@
+import unittest
+import shutil
+import tempfile
+
+import time
+
+import etcd
+import etcd.auth
+from etcd.tests.integration.test_simple import EtcdIntegrationTest
+from etcd.tests.integration import helpers
+
+
+class TestAuthentication(unittest.TestCase):
+    def setUp(self):
+        # Restart etcd for each test (since some tests will lock others out)
+        program = EtcdIntegrationTest._get_exe()
+        self.directory = tempfile.mkdtemp(prefix='python-etcd')
+        self.processHelper = helpers.EtcdProcessHelper(
+            self.directory,
+            proc_name=program,
+            port_range_start=6001,
+            internal_port_range_start=8001)
+        self.processHelper.run(number=1)
+        self.client = etcd.auth.AuthClient(port=6001)
+
+        # Wait for sync, to avoid:
+        # "Not capable of accessing auth feature during rolling upgrades."
+        time.sleep(0.5)
+
+    def tearDown(self):
+        self.processHelper.stop()
+        shutil.rmtree(self.directory)
+
+    def test_create_user(self):
+        user = self.client.create_user('username', 'password')
+        assert user.name == 'username'
+        assert len(user.roles) == 0
+
+    def test_create_user_with_role(self):
+        user = self.client.create_user('username', 'password', roles=['root'])
+        assert user.name == 'username'
+        assert user.roles == ('root',)
+
+    def test_create_user_add_role(self):
+        user = self.client.create_user('username', 'password')
+        self.client.create_role('role')
+
+        # Empty to [root]
+        user.roles = ['root']
+        user = self.client.get_user('username')
+        assert user.roles == ('root',)
+
+        # [root] to [root,role]
+        user.roles = ['root', 'role']
+        user = self.client.get_user('username')
+        assert user.roles == ('role', 'root')
+
+        # [root,role] to [role]
+        user.roles = ['role']
+        user = self.client.get_user('username')
+        assert user.roles == ('role',)
+
+    def test_usernames_empty(self):
+        assert len(self.client.usernames) == 0
+
+    def test_usernames(self):
+        self.client.create_user('username', 'password', roles=['root'])
+        assert self.client.usernames == ['username']
+
+    def test_users(self):
+        self.client.create_user('username', 'password', roles=['root'])
+        users = self.client.users
+        assert len(users) == 1
+        assert users[0].name == 'username'
+
+    def test_get_user(self):
+        self.client.create_user('username', 'password', roles=['root'])
+        user = self.client.get_user('username')
+        assert user.roles == ('root',)
+
+    def test_get_user_not_found(self):
+        self.assertRaises(etcd.EtcdException, self.client.get_user, 'username')
+
+    def test_set_user_password(self):
+        self.client.create_user('username', 'password', roles=['root'])
+        user = self.client.get_user('username')
+        assert not user.password
+        user.password = 'new_password'
+        assert not user.password
+
+    def test_create_role(self):
+        role = self.client.create_role('role')
+        assert role.name == 'role'
+        assert len(role.permissions) == 0
+
+    def test_grant_role(self):
+        role = self.client.create_role('role')
+
+        # Read access to keys under /foo
+        role.permissions['/foo/*'] = 'R'
+        assert len(role.permissions) == 1
+        assert role.permissions['/foo/*'] == 'R'
+
+        # Write access to the key at /foo/bar
+        role.permissions['/foo/bar'] = 'W'
+        assert len(role.permissions) == 2
+
+        # Full access to keys under /pub
+        role.permissions['/pub/*'] = 'RW'
+        assert len(role.permissions) == 3
+
+        # Fresh fetch to bust cache:
+        role = self.client.get_role('role')
+        assert len(role.permissions) == 3
+
+    def test_get_role(self):
+        role = self.client.create_role('role')
+        role.permissions['/foo/*'] = 'R'
+
+        role = self.client.get_role('role')
+        assert len(role.permissions) == 1
+
+    def test_revoke_role(self):
+        role = self.client.create_role('role')
+        role.permissions['/foo/*'] = 'R'
+
+        del role.permissions['/foo/*']
+
+        role = self.client.get_role('role')
+        assert len(role.permissions) == 0
+
+    def test_modify_role_invalid(self):
+        role = self.client.create_role('role')
+        self.assertRaises(ValueError, role.permissions.__setitem__, '/foo/*',
+                          '')
+
+    def test_modify_role_permissions(self):
+        role = self.client.create_role('role')
+        role.permissions['/foo/*'] = 'R'
+
+        # Replace R with W
+        role.permissions['/foo/*'] = 'W'
+        assert role.permissions['/foo/*'] == 'W'
+        role = self.client.get_role('role')
+        assert role.permissions['/foo/*'] == 'W'
+
+        # Extend W to RW
+        role.permissions['/foo/*'] = 'WR'
+        role = self.client.get_role('role')
+        assert role.permissions['/foo/*'] == 'RW'
+
+        # NO-OP RW to RW
+        role.permissions['/foo/*'] = 'RW'
+        role = self.client.get_role('role')
+        assert role.permissions['/foo/*'] == 'RW'
+
+        # Reduce RW to W
+        role.permissions['/foo/*'] = 'W'
+        role = self.client.get_role('role')
+        assert role.permissions['/foo/*'] == 'W'
+
+    def test_role_names_empty(self):
+        assert self.client.role_names == ['root']
+
+    def test_role_names(self):
+        self.client.create_role('role')
+        assert self.client.role_names == ['role', 'root']
+
+    def test_roles(self):
+        self.client.create_role('role')
+        assert len(self.client.roles) == 2
+
+    def test_enable_auth(self):
+        # Store a value, lock out guests
+        self.client.write('/foo', 'bar')
+        self.client.create_user('root', 'rootpassword')
+        # Creating role before auth is enabled prevents default permissions
+        self.client.create_role('guest')
+        self.client.toggle_auth(True)
+
+        # Now we can't access key:
+        try:
+            self.client.get('/foo')
+            self.fail('Expected exception')
+        except etcd.EtcdException as e:
+            assert 'Insufficient credentials' in str(e)
+
+        # But an authenticated client can:
+        root_client = etcd.Client(port=6001,
+                                  username='root',
+                                  password='rootpassword')
+        assert root_client.get('/foo').value == 'bar'
+
+    def test_enable_auth_before_root_created(self):
+        self.assertRaises(etcd.EtcdException, self.client.toggle_auth, True)