class Permission(db.Model): __tablename__ = 'permission' id = db.Column(db.Integer, primary_key=True, autoincrement=True) name = db.Column(db.String(255), nullable=False) resource_id = db.Column(db.Integer, db.ForeignKey('resource.id')) def __init__(self, name): self.name = name
class BlacklistedToken(db.Model): id = db.Column(db.Integer, primary_key=True) jti = db.Column(db.String(36), nullable=False) user_identity = db.Column(db.String(50), nullable=False) expires = db.Column(db.DateTime, nullable=False) def as_json(self): """Get the JSON representation of a BlacklistedToken object. Returns: The JSON representation of a BlacklistedToken object. """ return { 'id': self.id, 'jti': self.jti, 'user': self.user_identity, 'expires': self.expires }
class Resource(db.Model, TrackModificationsMixIn): __tablename__ = 'resource' id = db.Column(db.Integer, primary_key=True, autoincrement=True) name = db.Column(db.String(255), nullable=False) role_id = db.Column(db.Integer, db.ForeignKey('role.id')) permissions = db.relationship('Permission', backref=db.backref('resource'), cascade='all, delete-orphan') def __init__(self, name, permissions): """Initializes a new Resource object, which is a type of Resource that a Role may have access to. Args: name (str): Name of the Resource object. permissions (list[str]): List of permissions ("create", "read", "update", "delete", "execute") for the Resource """ self.name = name self.set_permissions(permissions) def set_permissions(self, new_permissions): """Adds the given list of permissions to the Resource object. Args: new_permissions (list|set[str]): A list of permission names with which the Resource will be associated. These permissions must be in the set ["create", "read", "update", "delete", "execute"]. """ self.permissions = [] new_permission_names = set(new_permissions) self.permissions.extend( [Permission(permission) for permission in new_permission_names]) def as_json(self, with_roles=False): """Returns the dictionary representation of the Resource object. Args: with_roles (bool, optional): Boolean to determine whether or not to include Role objects associated with the Resource in the JSON representation. Defaults to False. """ out = { 'id': self.id, 'name': self.name, 'permissions': [permission.name for permission in self.permissions] } if with_roles: out["role"] = self.role.name return out
class CaseSubscription(db.Model, TrackModificationsMixIn): """ The ORM for the case subscriptions configuration """ __tablename__ = 'case_subscription' id = db.Column(db.Integer, primary_key=True, autoincrement=True) name = db.Column(db.String(255), nullable=False) subscriptions = db.Column(db.Text()) note = db.Column(db.String) def __init__(self, name, subscriptions=None, note=''): """ Constructs an instance of a CaseSubscription. Args: name (str): Name of the case subscription. subscriptions (list(dict)): A subscription JSON object. note (str, optional): Annotation of the event. """ self.name = name self.note = note if subscriptions is None: subscriptions = [] try: self.subscriptions = json.dumps(subscriptions) except json.JSONDecodeError: self.subscriptions = '[]' finally: subscriptions = { subscription['uid']: subscription['events'] for subscription in subscriptions } walkoff.case.subscription.add_cases({name: subscriptions}) def as_json(self): """ Gets the JSON representation of the CaseSubscription object. Returns: The JSON representation of the CaseSubscription object. """ return { "id": self.id, "name": self.name, "subscriptions": json.loads(self.subscriptions), "note": self.note } @staticmethod def update(case_name): """ Synchronizes the subscription from the subscription in memory in the core. Args: case_name (str): The name of case to synchronize. """ case = CaseSubscription.query.filter_by(name=case_name).first() if case and case_name in walkoff.case.subscription.subscriptions: case_subs = walkoff.case.subscription.subscriptions[case_name] case.subscriptions = json.dumps([{ 'uid': uid, 'events': events } for uid, events in case_subs.items()])
class Role(db.Model, TrackModificationsMixIn): __tablename__ = 'role' id = db.Column(db.Integer, primary_key=True, autoincrement=True) name = db.Column(db.String(80), unique=True, nullable=False) description = db.Column(db.String(255)) resources = db.relationship('Resource', backref=db.backref('role'), cascade='all, delete-orphan') def __init__(self, name, description='', resources=None): """Initializes a Role object. Each user has one or more Roles associated with it, which determines the user's permissions. Args: name (str): The name of the Role. description (str, optional): A description of the role. resources (list(dict[name:resource, permissions:list[permission])): A list of dictionaries containing the name of the resource, and a list of permission names associated with the resource. Defaults to None. """ self.name = name self.description = description self.resources = [] if resources: self.set_resources(resources) def set_resources(self, new_resources): """Adds the given list of resources to the Role object. Args: new_resources (list(dict[name:resource, permissions:list[permission])): A list of dictionaries containing the name of the resource, and a list of permission names associated with the resource. """ new_resource_names = set([resource['name'] for resource in new_resources]) current_resource_names = set([resource.name for resource in self.resources] if self.resources else []) resource_names_to_add = new_resource_names - current_resource_names resource_names_to_delete = current_resource_names - new_resource_names resource_names_intersect = current_resource_names.intersection(new_resource_names) self.resources[:] = [resource for resource in self.resources if resource.name not in resource_names_to_delete] for resource_perms in new_resources: if resource_perms['name'] in resource_names_to_add: self.resources.append(Resource(resource_perms['name'], resource_perms['permissions'])) elif resource_perms['name'] in resource_names_intersect: resource = Resource.query.filter_by(role_id=self.id, name=resource_perms['name']).first() if resource: resource.set_permissions(resource_perms['permissions']) db.session.commit() def as_json(self, with_users=False): """Returns the dictionary representation of the Role object. Args: with_users (bool, optional): Boolean to determine whether or not to include User objects associated with the Role in the JSON representation. Defaults to False. Returns: The dictionary representation of the Role object. """ out = {"id": self.id, "name": self.name, "description": self.description, "resources": [resource.as_json() for resource in self.resources]} if with_users: out['users'] = [user.username for user in self.users] return out
class ScheduledTask(db.Model, TrackModificationsMixIn): __tablename__ = 'scheduled_task' id = db.Column(db.Integer, primary_key=True, autoincrement=True) name = db.Column(db.String(255), nullable=False) description = db.Column(db.String(255), nullable=False) status = db.Column(db.Enum('running', 'stopped')) workflows = db.relationship('ScheduledWorkflow', cascade="all, delete-orphan", backref='post', lazy='dynamic') trigger_type = db.Column(db.Enum('date', 'interval', 'cron', 'unspecified')) trigger_args = db.Column(db.String(255)) def __init__(self, name, description='', status='running', workflows=None, task_trigger=None): self.name = name self.description = description if workflows is not None: for workflow in set(workflows): self.workflows.append(ScheduledWorkflow(uid=workflow)) if task_trigger is not None: construct_trigger( task_trigger) # Throws an error if the args are invalid self.trigger_type = task_trigger['type'] self.trigger_args = json.dumps(task_trigger['args']) else: self.trigger_type = 'unspecified' self.trigger_args = '{}' self.status = status if status in ('running', 'stopped') else 'running' if self.status == 'running' and self.trigger_type != 'unspecified': self._start_workflows() def update(self, json_in): trigger = None if 'task_trigger' in json_in and json_in['task_trigger']: trigger = construct_trigger( json_in['task_trigger'] ) # Throws an error if the args are invalid self._update_scheduler(trigger) self.trigger_type = json_in['task_trigger']['type'] self.trigger_args = json.dumps(json_in['task_trigger']['args']) if 'name' in json_in: self.name = json_in['name'] if 'description' in json_in: self.description = json_in['description'] if 'workflows' in json_in and json_in['workflows']: self._modify_workflows(json_in, trigger=trigger) if 'status' in json_in and json_in['status'] != self.status: self._update_status(json_in) def start(self): if self.status != 'running': self.status = 'running' if self.trigger_type != 'unspecified': self._start_workflows() def stop(self): if self.status != 'stopped': self.status = 'stopped' self._stop_workflows() def _update_status(self, json_in): self.status = json_in['status'] if self.status == 'running': self._start_workflows() elif self.status == 'stopped': self._stop_workflows() def _start_workflows(self, trigger=None): from walkoff.server.flaskserver import running_context trigger = trigger if trigger is not None else construct_trigger( self._reconstruct_scheduler_args()) running_context.controller.schedule_workflows( self.id, self._get_workflow_uids_as_list(), trigger) def _stop_workflows(self): from walkoff.server.flaskserver import running_context running_context.controller.scheduler.unschedule_workflows( self.id, self._get_workflow_uids_as_list()) def _modify_workflows(self, json_in, trigger): from walkoff.server.flaskserver import running_context new, removed = self.__get_different_workflows(json_in) for workflow in self.workflows: self.workflows.remove(workflow) for workflow in json_in['workflows']: self.workflows.append(ScheduledWorkflow(uid=workflow)) if self.trigger_type != 'unspecified' and self.status == 'running': trigger = trigger if trigger is not None else construct_trigger( self._reconstruct_scheduler_args()) if new: running_context.controller.schedule_workflows( self.id, new, trigger) if removed: running_context.controller.scheduler.unschedule_workflows( self.id, removed) def _update_scheduler(self, trigger): from walkoff.server.flaskserver import running_context running_context.controller.scheduler.update_workflows(self.id, trigger) def _reconstruct_scheduler_args(self): return { 'type': self.trigger_type, 'args': json.loads(self.trigger_args) } def _get_workflow_uids_as_list(self): return [workflow.uid for workflow in self.workflows] def __get_different_workflows(self, json_in): original_workflows = set(self._get_workflow_uids_as_list()) incoming_workflows = set(json_in['workflows']) new = incoming_workflows - original_workflows removed = original_workflows - incoming_workflows return new, removed def as_json(self): return { 'id': self.id, 'name': self.name, 'description': self.description, 'status': self.status, 'workflows': self._get_workflow_uids_as_list(), 'task_trigger': self._reconstruct_scheduler_args() }
class ScheduledWorkflow(db.Model): __tablename__ = 'scheduled_workflow' id = db.Column(db.Integer, primary_key=True, autoincrement=True) uid = db.Column(db.String(50), nullable=False) task_id = db.Column(db.Integer, db.ForeignKey('scheduled_task.id'))
class Message(db.Model): __tablename__ = 'message' id = db.Column(db.Integer, primary_key=True, autoincrement=True) subject = db.Column(db.String()) body = db.Column(db.String(), nullable=False) users = db.relationship('User', secondary=user_messages_association, backref=db.backref('messages', lazy='dynamic')) workflow_execution_uid = db.Column(db.String(25)) requires_reauth = db.Column(db.Boolean, default=False) requires_response = db.Column(db.Boolean, default=False) created_at = db.Column(db.DateTime, default=func.current_timestamp()) history = db.relationship('MessageHistory', backref='message', lazy=True) def __init__(self, subject, body, workflow_execution_uid, users, requires_reauth=False, requires_response=False): self.subject = subject self.body = body self.workflow_execution_uid = workflow_execution_uid self.users = users self.requires_reauth = requires_reauth self.requires_response = requires_response def record_user_action(self, user, action): if user in self.users: if ((action == MessageAction.unread and not self.user_has_read(user)) or (action == MessageAction.respond and (not self.requires_response or self.is_responded()[0]))): return elif action == MessageAction.delete: self.users.remove(user) self.history.append(MessageHistory(user, action)) def user_has_read(self, user): user_history = [ history_entry for history_entry in self.history if history_entry.user_id == user.id ] for history_entry in user_history[::-1]: if history_entry.action in (MessageAction.read, MessageAction.unread): if history_entry.action == MessageAction.unread: return False if history_entry.action == MessageAction.read: return True else: return False def user_last_read_at(self, user): user_history = [ history_entry for history_entry in self.history if history_entry.user_id == user.id ] for history_entry in user_history[::-1]: if history_entry.action == MessageAction.read: return history_entry.timestamp else: return None def get_read_by(self): return { entry.username for entry in self.history if entry.action == MessageAction.read } def is_responded(self): if not self.requires_response: return False, None, None for history_entry in self.history[::-1]: if history_entry.action == MessageAction.respond: return True, history_entry.timestamp, history_entry.username else: return False, None, None def as_json(self, with_read_by=True, user=None, summary=False): responded, responded_at, responded_by = self.is_responded() ret = { 'id': self.id, 'subject': self.subject, 'created_at': str(self.created_at), 'awaiting_response': self.requires_response and not responded } if user: ret['is_read'] = self.user_has_read(user) ret['last_read_at'] = str(self.user_last_read_at(user)) if not summary: ret.update({ 'body': json.loads(self.body), 'workflow_execution_uid': self.workflow_execution_uid, 'requires_reauthorization': self.requires_reauth, 'requires_response': self.requires_response }) if responded: ret['responded_at'] = str(responded_at) ret['responded_by'] = responded_by if with_read_by: ret['read_by'] = list(self.get_read_by()) return ret
class User(db.Model, TrackModificationsMixIn): __tablename__ = 'user' id = db.Column(db.Integer, primary_key=True, autoincrement=True) roles = db.relationship('Role', secondary=user_roles_association, backref=db.backref('users', lazy='dynamic')) username = db.Column(db.String(80), unique=True, nullable=False) _password = db.Column('password', db.String(255), nullable=False) active = db.Column(db.Boolean, default=True) last_login_at = db.Column(db.DateTime) current_login_at = db.Column(db.DateTime) last_login_ip = db.Column(db.String(45)) current_login_ip = db.Column(db.String(45)) login_count = db.Column(db.Integer, default=0) def __init__(self, name, password, roles=None): """Initializes a new User object Args: name (str): The username for the User. password (str): The password for the User. roles (list[int]): List of Role ids for the User. Defaults to None. """ self.username = name self._password = pbkdf2_sha512.hash(password) self.roles = [] if roles: self.set_roles(roles) @hybrid_property def password(self): """Returns the password for the user. """ return self._password @password.setter def password(self, new_password): """Sets the password for a user, and encrypts it. Args: new_password (str): The new password for the User. """ self._password = pbkdf2_sha512.hash(new_password) def verify_password(self, password_attempt): """Verifies that the input password matches with the stored password. Args: password_attempt(str): The input password. Returns: True if the passwords match, False if not. """ return pbkdf2_sha512.verify(password_attempt, self._password) def set_roles(self, new_roles): """Sets the roles for a User. Args: new_roles (list[int]|set(int)): A list of Role IDs for the User. """ new_role_ids = set(new_roles) new_roles = Role.query.filter( Role.id.in_(new_role_ids)).all() if new_role_ids else [] self.roles[:] = new_roles roles_not_added = new_role_ids - {role.id for role in new_roles} if roles_not_added: logger.warning( 'Cannot add roles {0} to user {1}. Roles do not exist'.format( roles_not_added, self.id)) def login(self, ip_address): """Tracks login information for the User upon logging in. Args: ip_address (str): The IP address from which the User logged in. """ self.last_login_at = self.current_login_at self.current_login_at = datetime.utcnow() self.last_login_ip = self.current_login_ip self.current_login_ip = ip_address self.login_count += 1 def logout(self): """Tracks login/logout information for the User upon logging out. """ if self.login_count > 0: self.login_count -= 1 else: logger.warning( 'User {} logged out, but login count was already at 0'.format( self.id)) db.session.commit() def has_role(self, role): """Checks if a User has a Role associated with it. Args: role (int): The ID of the Role. Returns: True if the User has the Role, False otherwise. """ return role in [role.id for role in self.roles] def as_json(self, with_user_history=False): """Returns the dictionary representation of a User object. Args: with_user_history (bool, optional): Boolean to determine whether or not to include user history in the JSON representation of the User. Defaults to False. Returns: The dictionary representation of a User object. """ out = { "id": self.id, "username": self.username, "roles": [role.as_json() for role in self.roles], "active": self.active } if with_user_history: out.update({ "last_login_at": self.last_login_at, "current_login_at": self.current_login_at, "last_login_ip": self.last_login_ip, "current_login_ip": self.current_login_ip, "login_count": self.login_count }) return out