class Host(Model('host')):
def api(self, options_callback = None, message_callback = None):
return zimagi.Client(
user = self.user,
token = self.token,
encryption_key = self.encryption_key if settings.ENCRYPT_COMMAND_API or settings.ENCRYPT_DATA_API else None,
host = self.host,
command_port = self.command_port,
data_port = self.data_port,
options_callback = options_callback,
message_callback = message_callback
)
def command_api(self, options_callback = None, message_callback = None):
return zimagi.command.Client(
user = self.user,
token = self.token,
encryption_key = self.encryption_key if settings.ENCRYPT_COMMAND_API else None,
host = self.host,
port = self.command_port,
options_callback = options_callback,
message_callback = message_callback
)
def data_api(self, options_callback = None):
return zimagi.data.Client(
user = self.user,
token = self.token,
encryption_key = self.encryption_key if settings.ENCRYPT_DATA_API else None,
host = self.host,
port = self.data_port,
options_callback = options_callback
)
class APITokenAuthentication(authentication.TokenAuthentication):
user_class = Model('user')
api_type = None
def get_auth_header(self, request):
header = authentication.get_authorization_header(request)
if header and isinstance(header, (bytes, bytearray)):
header = header.decode('utf-8')
return header
def parse_token(self, token_text):
auth_components = re.split(r'\s+', token_text)
if len(auth_components) != 3:
raise exceptions.AuthenticationFailed(
"Invalid token. Required format: {} <user_name> <token>".
format(self.keyword))
if auth_components[0].lower() != self.keyword.lower():
raise exceptions.AuthenticationFailed(
"Authentication header required: 'Authorization: {} <user_name> <token>'"
.format(self.keyword))
try:
user = self.user_class.facade.retrieve(auth_components[1])
token = auth_components[2]
except self.user_class.DoesNotExist:
raise exceptions.AuthenticationFailed(
'Invalid token: User not found')
return (user, token)
def authenticate(self, request):
token_text = self.get_auth_header(request)
if not token_text:
return None
user, token = self.parse_token(token_text)
try:
if self.api_type:
token = Cipher.get(self.api_type,
user=user.name).decrypt(token)
except Exception as error:
raise exceptions.AuthenticationFailed(
'Invalid token header. Credentials can not be decrypted')
if not user.is_active:
raise exceptions.AuthenticationFailed(
'User account is inactive. Contact administrator')
if not user.check_password(token):
raise exceptions.AuthenticationFailed(
'Invalid token: User credentials are invalid')
user.last_login = now()
user.save()
self.user_class.facade.set_active_user(user)
return (user, token)
def process_response(self, request, response):
if request.path != '/status':
cache_entry = Model('cache').facade.get_or_create(
request.build_absolute_uri())
cache_entry.requests += 1
cache_entry.save()
if not (hasattr(request, '_cache_update_cache')
and request._cache_update_cache):
return response
response['Object-Cache'] = 'MISS'
if response.streaming or response.status_code not in (200, 304):
return response
if not request.COOKIES and response.cookies and has_vary_header(
response, 'Cookie'):
return response
if 'private' in response.get('Cache-Control', ()):
return response
timeout = get_max_age(response)
if timeout is None:
timeout = self.cache_timeout
elif timeout == 0:
return response
patch_response_headers(response, timeout)
if timeout and response.status_code == 200:
cache_key = learn_cache_key(request,
response,
timeout,
self.key_prefix,
cache=self.cache)
if hasattr(response, 'render') and callable(response.render):
response.add_post_render_callback(
lambda r: self.cache.set(cache_key, r, timeout))
else:
self.cache.set(cache_key, response, timeout)
return response
class ScheduledTask(ScheduleModelMixin,
DerivedAbstractModel(celery_beat_models,
'PeriodicTask',
id=None,
name=None,
args=None,
kwargs=None,
interval=None,
crontab=None,
clocked=None,
solar=None), Model('scheduled_task')):
objects = celery_beat_managers.PeriodicTaskManager()
class ScheduledTask(
ScheduleModelMixin,
DerivedAbstractModel(celery_beat_models, 'PeriodicTask',
id = None,
name = None,
args = None,
kwargs = None,
interval = None,
crontab = None,
clocked = None,
solar = None
),
Model('scheduled_task')
):
objects = celery_beat_managers.PeriodicTaskManager()
def validate_unique(self, *args, **kwargs):
super(celery_beat_models.PeriodicTask, self).validate_unique(*args, **kwargs)
schedule_types = ['interval', 'crontab', 'clocked']
selected_schedule_types = [s for s in schedule_types
if getattr(self, s)]
if len(selected_schedule_types) == 0:
raise ValidationError(
'One of clocked, interval, crontab, or solar '
'must be set.'
)
err_msg = 'Only one of clocked, interval, crontab, '\
'or solar must be set'
if len(selected_schedule_types) > 1:
error_info = {}
for selected_schedule_type in selected_schedule_types:
error_info[selected_schedule_type] = [err_msg]
raise ValidationError(error_info)
# clocked must be one off task
if self.clocked and not self.one_off:
err_msg = 'clocked must be one off, one_off must set True'
raise ValidationError(err_msg)
def save(self, *args, **kwargs):
self.exchange = self.exchange or None
self.routing_key = self.routing_key or None
self.queue = self.queue or None
self.headers = self.headers or None
if not self.enabled:
self.last_run_at = None
self._clean_expires()
self.validate_unique()
super(celery_beat_models.PeriodicTask, self).save(*args, **kwargs)
class User(AbstractBaseUser, Model('user')):
USERNAME_FIELD = 'name'
objects = UserManager()
def save(self, *args, **kwargs):
if not self.password and self.name == settings.ADMIN_USER:
self.set_password(settings.DEFAULT_ADMIN_TOKEN)
super().save(*args, **kwargs)
@property
def env_groups(self, **filters):
filters['environment_id'] = Model('environment').facade.get_env()
return self.groups.filter(**filters)
class APITokenAuthentication(authentication.TokenAuthentication):
user_class = Model('user')
def get_auth_header(self, request):
return authentication.get_authorization_header(request)
def validate_token_header(self, auth):
if len(auth) == 1:
msg = 'Invalid token header. No credentials provided.'
logger.warning(msg)
raise exceptions.AuthenticationFailed(msg)
elif len(auth) > 2:
msg = 'Invalid token header. Token string should not contain spaces.'
logger.warning(msg)
raise exceptions.AuthenticationFailed(msg)
def authenticate_credentials(self, auth):
self.validate_token_header(auth)
components = auth[1].split('++')
if len(components) != 2:
raise exceptions.AuthenticationFailed(
'Invalid token. Required format: Token user++token')
try:
user = self.user_class.facade.retrieve(components[0])
token = components[1]
except self.user_class.DoesNotExist:
raise exceptions.AuthenticationFailed(
'Invalid token: User not found')
if not user.is_active:
raise exceptions.AuthenticationFailed(
'User account is inactive. Contact administrator')
if not user.check_password(token):
raise exceptions.AuthenticationFailed(
'Invalid token: User credentials are invalid')
user.last_login = now()
user.save()
self.user_class.facade.set_active_user(user)
return (user, token)
class User(AbstractBaseUser, Model('user')):
USERNAME_FIELD = 'name'
objects = UserManager()
def save(self, *args, **kwargs):
if not self.password and self.name == settings.ADMIN_USER:
self.set_password(settings.DEFAULT_ADMIN_TOKEN)
if not self.encryption_key or self.encryption_key == '<generate>':
self.encryption_key = Cipher.get_provider_class(
'user_api_key').generate_key()
super().save(*args, **kwargs)
@property
def env_groups(self, **filters):
return self.groups.filter(**filters)
class Module(Model('module')):
STATUS_VALID = 'valid'
STATUS_INVALID = 'invalid'
@property
def status(self):
path = self.provider.module_path(self.name, ensure = False)
zimagi_path = os.path.join(path, 'zimagi.yml')
if self.provider.check_system() or os.path.isfile(zimagi_path):
return self.STATUS_VALID
return self.STATUS_INVALID
def save(self, *args, **kwargs):
try:
caches[settings.CACHE_MIDDLEWARE_ALIAS].clear()
caches[settings.CACHE_MIDDLEWARE_ALIAS].close()
except Exception:
pass
super().save(*args, **kwargs)
self.save_deploy_modules()
@classmethod
def save_deploy_modules(cls):
config_facade = model_index().get_facade_index()['config']
deploy_modules = []
for module in cls.facade.all():
if module.remote:
deploy_modules.append({
'remote': module.remote,
'reference': module.reference,
'config': module.config
})
config_facade.store('deploy_modules',
value = serialized_token() + serialize(deploy_modules),
value_type = 'str'
)
class Environment(Model('environment')):
@property
def is_active(self):
return True if self.name == self.facade.get_env() else False
def save(self, *args, **kwargs):
env_name = Runtime.get_env()
if self.name == env_name:
image = self.base_image
if self.runtime_image:
image = self.runtime_image
Runtime.set_env(
self.name,
self.repo,
image
)
super().save(*args, **kwargs)
class Log(Model('log')):
STATUS_QUEUED = 'queued'
STATUS_RUNNING = 'running'
STATUS_SUCCESS = 'success'
STATUS_FAILED = 'failed'
STATUS_ABORTED = 'aborted'
STATUS_UNTRACKED = 'untracked'
def save(self, *args, **kwargs):
if not self.name:
self.name = "{}{}x".format(now().strftime("%Y%m%d%H%M%S"),
self.facade.generate_token(5))
super().save(*args, **kwargs)
def success(self):
return self.status == self.STATUS_SUCCESS
def failed(self):
return self.status == self.STATUS_FAILED
def aborted(self):
return self.status == self.STATUS_ABORTED
def running(self):
return self.status == self.STATUS_RUNNING
def queued(self):
return self.status == self.STATUS_QUEUED
def untracked(self):
return self.status == self.STATUS_UNTRACKED
def set_status(self, status):
if isinstance(status, bool):
self.status = self.STATUS_SUCCESS if status else self.STATUS_FAILED
else:
self.status = status
def env_groups(self, **filters):
filters['environment_id'] = Model('environment').facade.get_env()
return self.groups.filter(**filters)
class Config(Model('config')):
def save(self, *args, **kwargs):
self.value = data.format_value(self.value_type, self.value)
super().save(*args, **kwargs)
class LogMessage(Model('log_message')):
def __str__(self):
return "{} ({})".format(self.log.command, self.data)
class TaskDatetime(ScheduleModelMixin,
DerivedAbstractModel(celery_beat_models,
'ClockedSchedule',
id=None), Model('task_datetime')):
pass
class TaskCrontab(ScheduleModelMixin,
DerivedAbstractModel(celery_beat_models,
'CrontabSchedule',
id=None), Model('task_crontab')):
pass
def get_base_scope(self):
return {'environment_id': Model('environment').facade.get_env()}
class Group(Model('group')):
def __str__(self):
if self.parent:
return "{} ({})".format(self.name, self.parent)
return self.name
class TaskInterval(ScheduleModelMixin,
DerivedAbstractModel(celery_beat_models,
'IntervalSchedule',
id=None), Model('task_interval')):
pass