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:
(dict): The JSON representation of a BlacklistedToken object.
"""
return {
'id': self.id,
'jti': self.jti,
'user': self.user_identity,
'expires': str(self.expires)
}
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(JSONType)
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 = subscriptions
except json.JSONDecodeError:
self.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": self.subscriptions,
"note": self.note
}
@staticmethod
def from_json(name, subscription_json):
""" Forms a CaseSubscription object from the provided JSON object.
Args:
name (str): The name of the case
subscription_json (dict): A JSON representation of the subscription
Returns:
The CaseSubscription object parsed from the JSON object.
"""
return CaseSubscription(name, subscriptions=subscription_json)
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 ScheduledTask(db.Model, TrackModificationsMixIn):
"""A SqlAlchemy table representing a a task scheduled for periodic execution
Attributes:
id (int): The primary key
name (str): The name of the task
description (str): A description of the task
status (str): The status of the task. either "running" or "stopped"
workflows (list[ScheduledWorkflow]): The workflows attached to this task
trigger_type (str): The type of trigger to use for the scheduler. Either "date", "interval", "cron", or
"unspecified"
trigger_args (str): The arguments for the scheduler trigger
Args:
name (str): The name of the task
description (str, optional): A description of the task. Defaults to empty string
workflows (list[str], optional): The uuids of the workflows attached to this task. Defaults to empty list,
task_trigger (dict): A dict containing two fields: "type", which contains the type of trigger to use for the
scheduler ("date", "interval", "cron", or "unspecified"), and "args", which contains the arguments for the
scheduler trigger
"""
__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(workflow_id=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):
"""Updates this task from a JSON representation of it
Args:
json_in (dict): The JSON representation of the updated task
"""
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):
"""Start executing this task
"""
if self.status != 'running':
self.status = 'running'
if self.trigger_type != 'unspecified':
self._start_workflows()
def stop(self):
"""Stop executing this scheduled task
"""
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 flask import current_app
trigger = trigger if trigger is not None else construct_trigger(self._reconstruct_scheduler_args())
current_app.running_context.scheduler.schedule_workflows(self.id,
current_app.running_context.executor.execute_workflow,
self._get_workflow_ids_as_list(), trigger)
def _stop_workflows(self):
from flask import current_app
current_app.running_context.scheduler.unschedule_workflows(self.id, self._get_workflow_ids_as_list())
def _modify_workflows(self, json_in, trigger):
from flask import current_app
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(workflow_id=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:
current_app.running_context.scheduler.schedule_workflows(self.id,
current_app.running_context.executor.execute_workflow,
new, trigger)
if removed:
current_app.running_context.scheduler.unschedule_workflows(self.id, removed)
def _update_scheduler(self, trigger):
from flask import current_app
current_app.running_context.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_ids_as_list(self):
return [workflow.workflow_id for workflow in self.workflows]
def __get_different_workflows(self, json_in):
original_workflows = set(self._get_workflow_ids_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):
"""Gets a JSON representation of this ScheduledTask
Returns:
(dict): The JSON representation of this ScheduledTask
"""
return {'id': self.id,
'name': self.name,
'description': self.description,
'status': self.status,
'workflows': self._get_workflow_ids_as_list(),
'task_trigger': self._reconstruct_scheduler_args()}
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":
utc_as_rfc_datetime(self.last_login_at),
"current_login_at":
utc_as_rfc_datetime(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
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_id = db.Column(UUIDType(binary=False), nullable=False)
requires_reauth = db.Column(db.Boolean, default=False)
requires_response = db.Column(db.Boolean, default=False)
created_at = db.Column(db.DateTime, default=datetime.utcnow)
history = db.relationship('MessageHistory', backref='message', lazy=True)
def __init__(self,
subject,
body,
workflow_execution_id,
users,
requires_reauth=False,
requires_response=False):
self.subject = subject
self.body = body
self.workflow_execution_id = workflow_execution_id
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': utc_as_rfc_datetime(self.created_at),
'awaiting_response': self.requires_response and not responded
}
if user:
ret['is_read'] = self.user_has_read(user)
last_read_at = self.user_last_read_at(user)
if last_read_at:
ret['last_read_at'] = utc_as_rfc_datetime(last_read_at)
if not summary:
ret.update({
'body': json.loads(self.body),
'workflow_execution_id': str(self.workflow_execution_id),
'requires_reauthorization': self.requires_reauth,
'requires_response': self.requires_response
})
if responded:
ret['responded_at'] = utc_as_rfc_datetime(responded_at)
ret['responded_by'] = responded_by
if with_read_by:
ret['read_by'] = list(self.get_read_by())
return ret
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:
(dict): 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 Message(db.Model):
"""Flask-SqlAlchemy Table which holds messages for users.
It has a many to many relationship with both the users and roles tables.
Attributes:
id (int): The primary key
subject (str): The subject of the message
body (str): The body of the message as a JSON string
users (list[User]): The users to which this message was sent and who haven't deleted the message
roles (list[Role]): The roles to which this message was sent
workflow_execution_id (UUID): The execution id of the workflow which sent the message
requires_reauth (bool): Does the message require reauthentication to address it?
requires_response (bool): Does the message require a response?
created_at (datetime): Timestamp of message creation
history (list[MessageHistory]): The timeline of actions taken on this message
Args:
subject (str): The subject of the message
body (str): The body of the message as a JSON string
users (list[User]): The users to which this message was sent and who haven't deleted the message
roles (list[Role]): The roles to which this message was sent
requires_reauth (bool): Does the message require reauthentication to address it?
requires_response (bool): Does the message require a response?
"""
__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'))
roles = db.relationship('Role', secondary=role_messages_association, backref=db.backref('messages', lazy='dynamic'))
workflow_execution_id = db.Column(UUIDType(binary=False), nullable=False)
requires_reauth = db.Column(db.Boolean, default=False)
requires_response = db.Column(db.Boolean, default=False)
created_at = db.Column(db.DateTime, default=datetime.utcnow)
history = db.relationship('MessageHistory', backref='message', lazy=True)
def __init__(self, subject, body, workflow_execution_id, users=None, roles=None, requires_reauth=False,
requires_response=False):
self.subject = subject
self.body = body
self.workflow_execution_id = workflow_execution_id
if not (users or roles):
message = 'Message must have users and/or roles, but has neither.'
logger.error(message)
raise ValueError(message)
self.users = users if users else []
self.roles = roles if roles else []
self.requires_reauth = requires_reauth
self.requires_response = requires_response
def record_user_action(self, user, action):
"""Records an action taken by a user on this message
Args:
user (User): The user taking the action
action (MessageAction): The action taken
"""
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):
"""Determines if a user has read the message
Args:
user (User): The user of the query
Returns:
(bool): Has the user read the message?
"""
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):
"""Gets the last time the user has read the message
Args:
user (User): The user of the query
Returns:
(datetime|None): The timestamp of the last time the user has read the message or None if the message has not
been read by this 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):
"""Gets all the usernames of the users who have read this message
Returns:
set(str): The usernames of the users who have read this message
"""
return {entry.username for entry in self.history if entry.action == MessageAction.read}
def is_responded(self):
"""Has this message been responded to?
Returns:
tuple(bool, datetime|None, str|None): A tuple of if the message has been responded to, if it has then the
datetime of when it was responded to and username of the user who responded to it.
"""
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 is_authorized(self, user_id=None, role_ids=None):
"""Is a user authorized to respond to this message?
Args:
user_id (int): The ID of the user
role_ids (list[int]): The ids of the roles the user has
Returns:
(bool)
"""
if user_id:
for user in self.users:
if user_id == user.id:
return True
if role_ids:
if isinstance(role_ids, int):
role_ids = [role_ids]
for role_id in role_ids:
for role in self.roles:
if role_id == role.id:
return True
return False
def as_json(self, with_read_by=True, user=None, summary=False):
"""Gets a JSON representation of the message
Args:
with_read_by (bool, optional): Should the JSON include who has read the message? Defaults to True.
user (User, optional): If provided, information specific to the user is included.
summary (bool, optional): If True, only give a brief summary of the messsage. Defaults to False.
Returns:
"""
responded, responded_at, responded_by = self.is_responded()
ret = {'id': self.id,
'subject': self.subject,
'created_at': utc_as_rfc_datetime(self.created_at),
'awaiting_response': self.requires_response and not responded}
if user:
ret['is_read'] = self.user_has_read(user)
last_read_at = self.user_last_read_at(user)
if last_read_at:
ret['last_read_at'] = utc_as_rfc_datetime(last_read_at)
if not summary:
ret.update({'body': json.loads(self.body),
'workflow_execution_id': str(self.workflow_execution_id),
'requires_reauthorization': self.requires_reauth,
'requires_response': self.requires_response})
if responded:
ret['responded_at'] = utc_as_rfc_datetime(responded_at)
ret['responded_by'] = responded_by
if with_read_by:
ret['read_by'] = list(self.get_read_by())
return ret
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(JSONType)
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 = subscriptions
except json.JSONDecodeError:
self.subscriptions = '[]'
finally:
subscriptions = {
subscription['id']: 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": 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 = [{
'id': uid,
'events': events
} for uid, events in case_subs.items()]
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(workflow_id=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.scheduler.schedule_workflows(
self.id, running_context.executor.execute_workflow,
self._get_workflow_ids_as_list(), trigger)
def _stop_workflows(self):
from walkoff.server.flaskserver import running_context
running_context.scheduler.unschedule_workflows(
self.id, self._get_workflow_ids_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(workflow_id=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.scheduler.schedule_workflows(
self.id, running_context.executor.execute_workflow, new,
trigger)
if removed:
running_context.scheduler.unschedule_workflows(
self.id, removed)
def _update_scheduler(self, trigger):
from walkoff.server.flaskserver import running_context
running_context.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_ids_as_list(self):
return [workflow.workflow_id for workflow in self.workflows]
def __get_different_workflows(self, json_in):
original_workflows = set(self._get_workflow_ids_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_ids_as_list(),
'task_trigger': self._reconstruct_scheduler_args()
}