def ban_user(auth):
admin = get_admin()
params = get_request_params()
with admin:
user = admin.get_user_by_email(params['email'])
if user is not None:
# Only superadmins can ban admins
if auth['user_type'] != UserType.SUPERADMIN and \
user['user_type'] in [UserType.ADMIN, UserType.SUPERADMIN]:
raise UnauthorizedError()
# Cannot ban yourself
if auth['user_id'] == user['id']:
raise UnauthorizedError()
return jsonify(admin.ban_user(**params))
def create_user(auth):
admin = get_admin()
params = get_request_params()
# Only superadmins can create admins
if auth['user_type'] != UserType.SUPERADMIN and \
params['user_type'] in [UserType.ADMIN, UserType.SUPERADMIN]:
raise UnauthorizedError()
with admin:
return jsonify(admin.create_user(**params))
def delete_model(auth, model_id):
admin = get_admin()
params = get_request_params()
with admin:
# Non-admins cannot delete others' models
if auth['user_type'] in [UserType.MODEL_DEVELOPER]:
model = admin.get_model(model_id)
if auth['user_id'] != model['user_id']:
raise UnauthorizedError()
return jsonify(admin.delete_model(model_id, **params))
def get_inference_jobs_by_user(auth):
admin = get_admin()
params = get_request_params()
assert 'user_id' in params
# Non-admins can only get their own jobs
if auth['user_type'] in [UserType.APP_DEVELOPER, UserType.MODEL_DEVELOPER] \
and auth['user_id'] != params['user_id']:
raise UnauthorizedError()
with admin:
return jsonify(admin.get_inference_jobs_by_user(**params))
def download_model_file(auth, model_id):
admin = get_admin()
params = get_request_params()
with admin:
# Non-admins cannot access others' models
if auth['user_type'] in [UserType.MODEL_DEVELOPER]:
model = admin.get_model(model_id)
if auth['user_id'] != model['user_id']:
raise UnauthorizedError()
model_file = admin.get_model_file(model_id, **params)
res = make_response(model_file)
res.headers.set('Content-Type', 'application/octet-stream')
return res
def create_train_job(auth):
admin = get_admin()
params = get_request_params()
with admin:
# Ensure that datasets are owned by current user
dataset_attrs = ['train_dataset_id', 'val_dataset_id']
for attr in dataset_attrs:
if attr in params:
dataset_id = params[attr]
dataset = admin.get_dataset(dataset_id)
if auth['user_id'] != dataset['owner_id']:
raise UnauthorizedError(
'You have no access to dataset of ID "{}"'.format(
dataset_id))
return jsonify(admin.create_train_job(auth['user_id'], **params))
def generate_user_token():
params = get_request_params()
# Only superadmin can authenticate (other users must use Rafiki Admin)
if not (params['email'] == SUPERADMIN_EMAIL and \
params['password'] == SUPERADMIN_PASSWORD):
raise UnauthorizedError()
auth = {
'user_type': UserType.SUPERADMIN
}
token = generate_token(auth)
return jsonify({
'user_type': auth['user_type'],
'token': token
})
def generate_user_token():
admin = get_admin()
params = get_request_params()
# Error will be thrown here if credentials are invalid
with admin:
user = admin.authenticate_user(**params)
# User cannot be banned
if user.get('banned_date'
) is not None and datetime.now() > user.get('banned_date'):
raise UnauthorizedError('User is banned')
token = generate_token(user)
return jsonify({
'user_id': user['id'],
'user_type': user['user_type'],
'token': token
})
