use a single arg for adding and removing of "passwords" and "hashed_passwords"acl

also ensure unicode values are properly encoded

2 files changed, 109 insertions, 92 deletions

diff --git a/redis/client.py b/redis/client.py
index 05c43da..5ace783 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -954,9 +954,9 @@ class Redis(object):
         return self.execute_command('ACL SAVE')
 
     def acl_setuser(self, username, enabled=False, nopass=False,
-                    add_passwords=None, remove_passwords=None, add_hashes=None,
-                    remove_hashes=None, categories=None, commands=None,
-                    keys=None, reset=False):
+                    passwords=None, hashed_passwords=None, categories=None,
+                    commands=None, keys=None, reset=False, reset_keys=False,
+                    reset_passwords=False):
         """
         Create or update an ACL user.
 
@@ -970,116 +970,133 @@ class Redis(object):
         ``nopass`` is a boolean indicating whether the can authenticate without
         a password. This cannot be True if ``passwords`` are also specified.
 
-        ``add_passwords`` if specified is a list of new plain text passwords
-        that this user can authenticate with. For convenience, the value of
-        ``add_passwords`` can also be a simple string when adding a single
-        password. Do not prefix passwords with '>'.
-
-        ``add_hashes`` if specified is a list of new hashed passwords that
-        this user can authenticate with. This is useful when you want to add a
-        password but do not want to specify the password as plain text. For
-        convenience, the value of ``add_hashes`` can also be a simple string
-        when adding a single password. All hashes must be SHA-256 hashes. Do
-        not prefix hashed passwords with '#'.
-
-        ``remove_passwords`` if specified is a list of plain text passwords
-        to remove from this user. For convenience, the value of
-        ``remove_passwords`` can also be a simple string when removing a
-        single password. Do not prefix passwords with '<'.
-
-        ``remove_hashes`` if specified is a list of hashed passwords to remove
-        from this user. This is useful when you want to remove a password but
-        do not want to specify the password as plain text. For convenience, the
-        value of ``remove_hashes`` can also be a simple string when removing
-        a single password. All hashes must be SHA-256 hashes. Do not prefix
-        hashed passwords with '!'.
+        ``passwords`` if specified is a list of plain text passwords
+        to add to or remove from the user. Each password must be prefixed with
+        a '+' to add or a '-' to remove. For convenience, the value of
+        ``add_passwords`` can be a simple prefixed string when adding or
+        removing a single password.
+
+        ``hashed_passwords`` if specified is a list of SHA-256 hashed passwords
+        to add to or remove from the user. Each hashed password must be
+        prefixed with a '+' to add or a '-' to remove. For convenience,
+        the value of ``hashed_passwords`` can be a simple prefixed string when
+        adding or removing a single password.
 
         ``categories`` if specified is a list of strings representing category
-        permissions. Each string must be prefixed with either a "+@" or "-@"
-        to indicate whether access to the category is allowed or denied.
+        permissions. Each string must be prefixed with either a '+' to add the
+        category permission or a '-' to remove the category permission.
 
         ``commands`` if specified is a list of strings representing command
-        permissions. Each string must be prefixed with either a "+" or "-"
-        to indicate whether access to the command is allowed or denied.
+        permissions. Each string must be prefixed with either a '+' to add the
+        command permission or a '-' to remove the command permission.
 
         ``keys`` if specified is a list of key patterns to grant the user
         access to. Keys patterns allow '*' to support wildcard matching. For
         example, '*' grants access to all keys while 'cache:*' grants access
-        to all keys that are prefixed with 'cache:'.
+        to all keys that are prefixed with 'cache:'. ``keys`` should not be
+        prefixed with a '~'.
 
         ``reset`` is a boolean indicating whether the user should be fully
         reset prior to applying the new ACL. Setting this to True will
         remove all existing passwords, flags and privileges from the user and
-        then apply the specified rules. If this is False, the existing user's
-        passwords, flags and privileges will be kept and the specified rules
+        then apply the specified rules. If this is False, the user's existing
+        passwords, flags and privileges will be kept and any new specified
+        rules will be applied on top.
+
+        ``reset_keys`` is a boolean indicating whether the user's key
+        permissions should be reset prior to applying any new key permissions
+        specified in ``keys``. If this is False, the user's existing
+        key permissions will be kept and any new specified key permissions
         will be applied on top.
+
+        ``reset_passwords`` is a boolean indicating whether to remove all
+        existing passwords and the 'nopass' flag from the user prior to
+        applying any new passwords specified in 'passwords' or
+        'hashed_passwords'. If this is False, the user's existing passwords
+        and 'nopass' status will be kept and any new specified passwords
+        or hashed_passwords will be applied on top.
         """
+        encoder = self.connection_pool.get_encoder()
         pieces = [username]
+
+        if reset:
+            pieces.append(b'reset')
+
+        if reset_keys:
+            pieces.append(b'resetkeys')
+
+        if reset_passwords:
+            pieces.append(b'resetpass')
+
         if enabled:
-            pieces.append('on')
+            pieces.append(b'on')
         else:
-            pieces.append('off')
+            pieces.append(b'off')
 
-        if (add_passwords or add_hashes) and nopass:
+        if (passwords or hashed_passwords) and nopass:
             raise DataError('Cannot set \'nopass\' and supply '
-                            '\'add_passwords\' or \'add_hashes\'')
+                            '\'passwords\' or \'hashed_passwords\'')
 
-        if remove_passwords:
+        if passwords:
             # as most users will have only one password, allow remove_passwords
             # to be specified as a simple string or a list
-            remove_passwords = list_or_args(remove_passwords, [])
-            for password in remove_passwords:
-                pieces.append('<%s' % password)
-
-        if remove_hashes:
-            # as most users will have only one password, allow remove_passwords
-            # to be specified as a simple string or a list
-            remove_hashes = list_or_args(remove_hashes, [])
-            for password in remove_hashes:
-                pieces.append('!%s' % password)
-
-        if add_passwords:
-            # as most users will have only one password, allow add_passwords
-            # to be specified as a simple string or a list
-            add_passwords = list_or_args(add_passwords, [])
-            for password in add_passwords:
-                pieces.append('>%s' % password)
+            passwords = list_or_args(passwords, [])
+            for i, password in enumerate(passwords):
+                password = encoder.encode(password)
+                if password.startswith(b'+'):
+                    pieces.append(b'>%s' % password[1:])
+                elif password.startswith(b'-'):
+                    pieces.append(b'<%s' % password[1:])
+                else:
+                    raise DataError('Password %d must be prefixeed with a '
+                                    '"+" to add or a "-" to remove' % i)
 
-        if add_hashes:
+        if hashed_passwords:
             # as most users will have only one password, allow remove_passwords
             # to be specified as a simple string or a list
-            add_hashes = list_or_args(add_hashes, [])
-            for password in add_hashes:
-                pieces.append('#%s' % password)
+            hashed_passwords = list_or_args(hashed_passwords, [])
+            for i, hashed_password in enumerate(hashed_passwords):
+                hashed_password = encoder.encode(hashed_password)
+                if hashed_password.startswith(b'+'):
+                    pieces.append(b'#%s' % hashed_password[1:])
+                elif hashed_password.startswith(b'-'):
+                    pieces.append(b'!%s' % hashed_password[1:])
+                else:
+                    raise DataError('Hashed %d password must be prefixeed '
+                                    'with a "+" to add or a "-" to remove' % i)
 
         if nopass:
             pieces.append(b'nopass')
 
         if categories:
-            for cat in categories:
-                # ensure the category starts with either a -@ or +@
-                if not cat.startswith('-@') and not cat.startswith('+@'):
+            for category in categories:
+                category = encoder.encode(category)
+                if category.startswith(b'+@'):
+                    pieces.append(category)
+                elif category.startswith(b'+'):
+                    pieces.append(b'+@%s' % category[1:])
+                elif category.startswith(b'-@'):
+                    pieces.append(category)
+                elif category.startswith(b'-'):
+                    pieces.append(b'-@%s' % category[1:])
+                else:
                     raise DataError('Category "%s" must be prefixed with '
-                                    '"+@" or "-@"' % cat)
-                pieces.append(cat.encode())
-
+                                    '"+" or "-"'
+                                    % encoder.decode(category, force=True))
         if commands:
             for cmd in commands:
-                if not cmd.startswith('+') and not cmd.startswith('-'):
+                cmd = encoder.encode(cmd)
+                if not cmd.startswith(b'+') and not cmd.startswith(b'-'):
                     raise DataError('Command "%s" must be prefixed with '
-                                    '"+" or "-"' % cmd)
-                pieces.append(cmd.encode())
+                                    '"+" or "-"'
+                                    % encoder.decode(cmd, force=True))
+                pieces.append(cmd)
 
         if keys:
             for key in keys:
-                if not key.startswith('~'):
-                    key = '~%s' % key
-                pieces.append(key)
+                key = encoder.encode(key)
+                pieces.append(b'~%s' % key)
 
-        if reset:
-            if not self.execute_command('ACL SETUSER', username, 'reset'):
-                # we first reset the user so we're working with a clean slate
-                raise RedisError('Failed to reset user: %s' % username)
         return self.execute_command('ACL SETUSER', *pieces)
 
     def acl_users(self):
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 04f1cd1..00752fa 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -128,8 +128,8 @@ class TestRedisCommands(object):
 
         # test all args
         assert r.acl_setuser(username, enabled=True, reset=True,
-                             add_passwords=['pass1', 'pass2'],
-                             categories=['+@set', '+@hash', '-@geo'],
+                             passwords=['+pass1', '+pass2'],
+                             categories=['+set', '+@hash', '-geo'],
                              commands=['+get', '+mget', '-hset'],
                              keys=['cache:*', 'objects:*'])
         acl = r.acl_getuser(username)
@@ -142,46 +142,46 @@ class TestRedisCommands(object):
 
         # test reset=False keeps existing ACL and applies new ACL on top
         assert r.acl_setuser(username, enabled=True, reset=True,
-                             add_passwords=['pass1'],
+                             passwords=['+pass1'],
                              categories=['+@set'],
-                             commands=['+set'],
+                             commands=['+get'],
                              keys=['cache:*'])
         assert r.acl_setuser(username, enabled=True,
-                             add_passwords=['pass2'],
+                             passwords=['+pass2'],
                              categories=['+@hash'],
-                             commands=['+mset'],
+                             commands=['+mget'],
                              keys=['objects:*'])
         acl = r.acl_getuser(username)
         assert set(acl['categories']) == set(['-@all', '+@set', '+@hash'])
-        assert set(acl['commands']) == set(['+set', '+mset'])
+        assert set(acl['commands']) == set(['+get', '+mget'])
         assert acl['enabled'] is True
         assert acl['flags'] == ['on']
         assert set(acl['keys']) == set([b'cache:*', b'objects:*'])
         assert len(acl['passwords']) == 2
 
-        # test remove_passwords
+        # test removal of passwords
         assert r.acl_setuser(username, enabled=True, reset=True,
-                             add_passwords=['pass1', 'pass2'])
+                             passwords=['+pass1', '+pass2'])
         assert len(r.acl_getuser(username)['passwords']) == 2
         assert r.acl_setuser(username, enabled=True,
-                             remove_passwords=['pass2'])
+                             passwords=['-pass2'])
         assert len(r.acl_getuser(username)['passwords']) == 1
 
-        # Resets and tests that hashed values are set properly.
+        # Resets and tests that hashed passwords are set properly.
         hashed_password = ('5e884898da28047151d0e56f8dc629'
                            '2773603d0d6aabbdd62a11ef721d1542d8')
         assert r.acl_setuser(username, enabled=True, reset=True,
-                             add_hashes=[hashed_password])
+                             hashed_passwords=['+' + hashed_password])
         acl = r.acl_getuser(username)
         assert acl['passwords'] == [hashed_password]
 
-        # test remove_passwords for hash removal
+        # test removal of hashed passwords
         assert r.acl_setuser(username, enabled=True, reset=True,
-                             add_hashes=[hashed_password],
-                             add_passwords=['pass1'])
+                             hashed_passwords=['+' + hashed_password],
+                             passwords=['+pass1'])
         assert len(r.acl_getuser(username)['passwords']) == 2
         assert r.acl_setuser(username, enabled=True,
-                             remove_hashes=[hashed_password])
+                             hashed_passwords=['-' + hashed_password])
         assert len(r.acl_getuser(username)['passwords']) == 1
 
     @skip_if_server_version_lt('5.9.101')
@@ -227,7 +227,7 @@ class TestRedisCommands(object):
         request.addfinalizer(teardown)
 
         with pytest.raises(exceptions.DataError):
-            r.acl_setuser(username, add_passwords='mypass', nopass=True)
+            r.acl_setuser(username, passwords='+mypass', nopass=True)
 
     @skip_if_server_version_lt('5.9.101')
     def test_acl_users(self, r):