Manage groups from tsctl and add tests

timesketch/models/acl.py

@@ -146,6 +146,8 @@ def _get_ace(self, permission, user=None, group=None):
         ace = self.AccessControlEntry.query.filter_by(
             user=user, group=None, permission=permission, parent=self).all()
 
+        # TODO(jbn) Make this more efficient. For now we don't expect too many
+        # groups per user so this should be OK for now.
         if user and not ace:
             for group in user.groups:
                 ace = self.AccessControlEntry.query.filter_by(
@@ -206,15 +208,13 @@ def grant_permission(self, permission, user=None, group=None):
         """
         # Grant permission to a group.
         if group and not self._get_ace(permission=permission, group=group):
-            print "grant group {0}".format(group.name)
             self.acl.append(
                 self.AccessControlEntry(permission=permission, group=group))
             db_session.commit()
             return
 
         # Grant permission to a user.
         if not self._get_ace(permission=permission, user=user):
-            print "grant user {0}".format(user)
             self.acl.append(
                 self.AccessControlEntry(permission=permission, user=user))
             db_session.commit()
@@ -231,7 +231,6 @@ def revoke_permission(self, permission, user=None, group=None):
         if group:
             group_ace = self._get_ace(permission=permission, group=group)
             if group_ace:
-                print "revoke group {0}".format(group.name)
                 for ace in group_ace:
                     self.acl.remove(ace)
                 db_session.commit()
@@ -240,7 +239,6 @@ def revoke_permission(self, permission, user=None, group=None):
         # Revoke permission for a user.
         user_ace = self._get_ace(permission=permission, user=user)
         if user_ace:
-            print "revoke user {0}".format(user)
             for ace in user_ace:
                 self.acl.remove(ace)
             db_session.commit()

timesketch/models/acl_test.py

@@ -32,6 +32,18 @@ def test_change_permission(self):
                 self.sketch1.has_permission(
                     permission=permission, user=self.user1))
 
+            # Test group permissions
+            self.sketch1.grant_permission(
+                permission=permission, group=self.group)
+            self.assertTrue(
+                self.sketch1.has_permission(
+                    permission=permission, user=self.user1))
+            self.sketch1.revoke_permission(
+                permission=permission, group=self.group)
+            self.assertFalse(
+                self.sketch1.has_permission(
+                    permission=permission, user=self.user1))
+
     def test_change_public(self):
         """Test toggle the public permission on a sketch."""
         self.sketch1.grant_permission(permission=u'read')

tsctl

@@ -29,10 +29,13 @@ from flask_script import Option
 from flask_script import prompt_bool
 from flask_script import prompt_pass
 
+from sqlalchemy.exc import IntegrityError
+
 from timesketch import create_app
 from timesketch.lib.datastores.elastic import ElasticSearchDataStore
 from timesketch.models import db_session
 from timesketch.models import drop_all
+from timesketch.models.user import Group
 from timesketch.models.user import User
 from timesketch.models.sketch import SearchIndex
 
@@ -85,6 +88,74 @@ class AddUser(Command):
         sys.stdout.write(u'User {0:s} created/updated\n'.format(username))
 
 
+class AddGroup(Command):
+    """Create a new Timesketch group."""
+    option_list = (
+        Option(u'--name', u'-n', dest=u'name', required=True),
+    )
+
+    def __init__(self):
+        super(AddGroup, self).__init__()
+
+    # pylint: disable=arguments-differ, method-hidden
+    def run(self, name):
+        """Creates the group."""
+        name = unicode(name.decode(encoding=u'utf-8'))
+        group = Group.get_or_create(name=name)
+        db_session.add(group)
+        db_session.commit()
+        sys.stdout.write(u'Group {0:s} created\n'.format(name))
+
+
+class GroupManager(Command):
+    """Manage group memberships."""
+    option_list = (
+        Option(
+            u'--add', u'-a', dest=u'add', action=u'store_true', required=False,
+            default=False),
+        Option(
+            u'--remove', u'-r', dest=u'remove', action=u'store_true',
+            required=False, default=False),
+        Option(u'--group', u'-g', dest=u'group_name', required=True),
+        Option(u'--user', u'-u', dest=u'user_name', required=True),
+    )
+
+    def __init__(self):
+        super(GroupManager, self).__init__()
+
+    # pylint: disable=arguments-differ, method-hidden
+    def run(self, add, remove, group_name, user_name):
+        """Add the user to the group."""
+        group_name = unicode(group_name.decode(encoding=u'utf-8'))
+        user_name = unicode(user_name.decode(encoding=u'utf-8'))
+        group = Group.query.filter_by(name=group_name).first()
+        user = User.query.filter_by(username=user_name).first()
+
+        # Add or remove user from group
+        if remove:
+            try:
+                user.groups.remove(group)
+                sys.stdout.write(
+                    u'{0:s} removed from group {1:s}\n'.format(
+                        user_name, group_name))
+                db_session.commit()
+            except ValueError:
+                sys.stdout.write(
+                    u'{0:s} is not a member of group {1:s}\n'.format(
+                        user_name, group_name))
+        else:
+            user.groups.append(group)
+            try:
+                db_session.commit()
+                sys.stdout.write(
+                    u'{0:s} added to group {1:s}\n'.format(
+                        user_name, group_name))
+            except IntegrityError:
+                sys.stdout.write(
+                    u'{0:s} is already a member of group {1:s}\n'.format(
+                        user_name, group_name))
+
+
 class AddSearchIndex(Command):
     """Create a new Timesketch searchindex."""
     option_list = (
@@ -343,6 +414,8 @@ if __name__ == '__main__':
     # Setup Flask-script command manager and register commands.
     shell_manager = Manager(create_app)
     shell_manager.add_command(u'add_user', AddUser())
+    shell_manager.add_command(u'add_group', AddGroup())
+    shell_manager.add_command(u'manage_group', GroupManager())
     shell_manager.add_command(u'add_index', AddSearchIndex())
     shell_manager.add_command(u'csv2ts', CreateTimelineFromCsv())
     shell_manager.add_command(u'db', MigrateCommand)