def create(reservation: Dict[str, Any]) -> Tuple[Content, HttpStatusCode]:
try:
new_reservation = Reservation(title=reservation['title'],
description=reservation['description'],
resource_id=reservation['resourceId'],
user_id=reservation['userId'],
start=reservation['start'],
end=reservation['end'])
user = User.get(get_jwt_identity())
if (is_admin() or __is_reservation_owner(new_reservation)) \
and ReservationVerifier.is_reservation_allowed(user, new_reservation):
new_reservation.save()
content = {
'msg': RESERVATION['create']['success'],
'reservation': new_reservation.as_dict()
}
status = 201
else:
content = {'msg': RESERVATION['create']['failure']['forbidden']}
status = 403
except AssertionError as e:
content = {
'msg': RESERVATION['create']['failure']['invalid'].format(reason=e)
}
status = 422
except Exception as e:
print(e)
content = {'msg': GENERAL['internal_error'] + str(e)}
status = 500
finally:
return content, status
def add_claims_to_access_token(current_user_id):
try:
current_user = User.get(current_user_id)
roles = current_user.role_names
except Exception:
roles = []
finally:
return {'roles': roles}
def get_by_id(id):
# Dead code, web app is currently not using it
try:
user = User.get(id)
except NoResultFound as e:
log.warning(e)
content, status = {'msg': U['not_found']}, 404
except Exception as e:
log.critical(e)
content, status = {'msg': G['internal_error']}, 500
else:
content, status = {'msg': U['get']['success'], 'user': user.as_dict}, 200
finally:
return content, status
def update(id: ReservationId,
newValues: Dict[str, Any]) -> Tuple[Content, HttpStatusCode]:
new_values = newValues
allowed_fields = {'title', 'description', 'resourceId', 'end'}
try:
reservation = Reservation.get(id)
if reservation.end < datetime.utcnow() and not is_admin():
raise ForbiddenException('reservation already finished')
if reservation.start > datetime.utcnow() or is_admin():
allowed_fields.add('start')
if not set(new_values.keys()).issubset(allowed_fields):
raise ForbiddenException('invalid field is present')
for field_name, new_value in new_values.items():
field_name = snakecase(field_name)
assert (field_name is not None) and hasattr(reservation, field_name), \
'reservation has no {} field'.format(field_name)
setattr(reservation, field_name, new_value)
user = User.get(get_jwt_identity())
if not (is_admin() or __is_reservation_owner(reservation)) or not \
ReservationVerifier.is_reservation_allowed(user, reservation):
raise ForbiddenException("reservation not allowed")
reservation.is_cancelled = False
reservation.save()
content, status = {
'msg': RESERVATION['update']['success'],
'reservation': reservation.as_dict()
}, 201
except ForbiddenException as fe:
content, status = {
'msg':
RESERVATION['update']['failure']['forbidden'].format(reason=fe)
}, 403
except NoResultFound:
content, status = {'msg': RESERVATION['not_found']}, 404
except AssertionError as e:
content, status = {
'msg':
RESERVATION['update']['failure']['assertions'].format(reason=e)
}, 422
except Exception as e:
log.critical(e)
content, status = {'msg': GENERAL['internal_error'] + str(e)}, 500
finally:
return content, status
def test_on_signup_user_gets_added_to_all_default_groups_if_there_are_more_than_one(
tables, client, new_group):
new_group.is_default = True
new_group.save()
another_default_group = Group(name='AnotherDefaultGroup', is_default=True)
another_default_group.save()
data = {
'email': '*****@*****.**',
'username': '******',
'password': '******'
}
resp = client.post(ENDPOINT + '/create',
data=json.dumps(data),
headers=HEADERS)
resp_json = json.loads(resp.data.decode('utf-8'))
assert resp.status_code == HTTPStatus.CREATED
assert len(resp_json['user']['groups']) == 2
assert resp_json['user']['groups'][0]['id'] == new_group.id
assert resp_json['user']['groups'][1]['id'] == another_default_group.id
assert new_group in User.get(resp_json['user']['id']).groups
assert another_default_group in User.get(resp_json['user']['id']).groups
def create(reservation: Dict[str, Any]) -> Tuple[Content, HttpStatusCode]:
try:
new_reservation = Reservation(title=reservation['title'],
description=reservation['description'],
resource_id=reservation['resourceId'],
user_id=reservation['userId'],
start=reservation['start'],
end=reservation['end'])
if not is_admin() and not __is_reservation_owner(new_reservation):
raise ForbiddenException(
"Cannot reserve resources in another user's name")
reservation_start = DateUtils.try_parse_string(new_reservation.start)
request_time_limit = timedelta(minutes=1)
starts_in_the_future = (reservation_start +
request_time_limit) >= datetime.utcnow()
if not is_admin() and not starts_in_the_future:
raise ForbiddenException("Cannot reserve resources in the past")
user = User.get(get_jwt_identity())
if not ReservationVerifier.is_reservation_allowed(
user, new_reservation):
raise ForbiddenException("Reservation not allowed")
new_reservation.save()
content = {
'msg': RESERVATION['create']['success'],
'reservation': new_reservation.as_dict()
}
status = 201
except ForbiddenException as e:
content = {
'msg':
RESERVATION['create']['failure']['forbidden'].format(reason=e)
}
status = 403
except AssertionError as e:
content = {
'msg': RESERVATION['create']['failure']['invalid'].format(reason=e)
}
status = 422
except Exception as e:
print(e)
content = {'msg': GENERAL['internal_error'] + str(e)}
status = 500
finally:
return content, status
def get_selected(user_id: Optional[UserId], group_id: Optional[GroupId], resource_id: Optional[ResourceId],
schedule_id: Optional[ScheduleId], include_user_groups: Optional[bool] = False) \
-> Tuple[Union[List[Any], Content], HttpStatusCode]:
try:
# If a specific group is selected then groups are not included in the restriction information in response
# The same applies to users and resources
include_groups = group_id is None
include_users = user_id is None
include_resources = schedule_id is None
restrictions = [] # type: List[Restriction]
if user_id is not None:
user = User.get(user_id)
restrictions.extend(
user.get_restrictions(include_group=include_user_groups))
if group_id is not None:
group = Group.get(group_id)
restrictions.extend(group.get_restrictions())
if resource_id is not None:
resource = Resource.get(resource_id)
restrictions.extend(resource.get_restrictions())
if schedule_id is not None:
schedule = RestrictionSchedule.get(schedule_id)
restrictions.extend(schedule.restrictions)
# Take unique restrictions
result = set(restrictions)
except NoResultFound as e:
log.warning(e)
content, status = {
'msg': GENERAL['bad_request']
}, HTTPStatus.BAD_REQUEST.value
except Exception as e:
log.critical(e)
content, status = {
'msg': GENERAL['internal_error']
}, HTTPStatus.INTERNAL_SERVER_ERROR.value
else:
content = [
restriction.as_dict(
include_groups=include_groups,
include_users=include_users, # type: ignore
include_resources=include_resources) for restriction in result
]
status = HTTPStatus.OK.value
finally:
return content, status
def test_on_signup_user_does_not_belong_to_any_group_if_no_default_group_exists(
tables, client, new_group):
new_group.save()
data = {
'email': '*****@*****.**',
'username': '******',
'password': '******'
}
resp = client.post(ENDPOINT + '/create',
data=json.dumps(data),
headers=HEADERS)
resp_json = json.loads(resp.data.decode('utf-8'))
assert resp.status_code == HTTPStatus.CREATED
assert len(resp_json['user']['groups']) == 0
assert len(User.get(resp_json['user']['id']).groups) == 0
def get_by_id(id: UserId) -> Tuple[Content, HttpStatusCode]:
try:
user = User.get(id)
except NoResultFound as e:
log.warning(e)
content, status = {'msg': USER['not_found']}, 404
except Exception as e:
log.critical(e)
content, status = {'msg': GENERAL['internal_error']}, 500
else:
current_user_id = get_jwt_identity()
claims = get_jwt_claims()
include_private = 'admin' in claims['roles'] or id == current_user_id
content, status = {
'msg': USER['get']['success'],
'user': user.as_dict(include_private=include_private)
}, 200
finally:
return content, status
def delete(id: UserId) -> Tuple[Content, HttpStatusCode]:
try:
current_user_id = get_jwt_identity()
# User is not allowed to delete his own account
assert id != current_user_id, USER['delete']['self']
# Fetch the user and destroy
user_to_destroy = User.get(id)
user_to_destroy.destroy()
except AssertionError as error_message:
content, status = {'msg': str(error_message)}, 403
except NoResultFound:
content, status = {'msg': USER['not_found']}, 404
except Exception as e:
content, status = {'msg': GENERAL['internal_error'] + str(e)}, 500
else:
content, status = {'msg': USER['delete']['success']}, 200
finally:
return content, status
def remove_from_user(restriction_id: RestrictionId,
user_id: UserId) -> Tuple[Content, HttpStatusCode]:
restriction = None
try:
restriction = Restriction.get(restriction_id)
user = User.get(user_id)
restriction.remove_from_user(user)
ReservationVerifier.update_user_reservations_statuses(
user, have_users_permissions_increased=False)
except NoResultFound:
if restriction is None:
content, status = {
'msg': RESTRICTION['not_found']
}, HTTPStatus.NOT_FOUND.value
else:
content, status = {
'msg': USER['not_found']
}, HTTPStatus.NOT_FOUND.value
except InvalidRequestException:
content, status = {
'msg': RESTRICTION['users']['remove']['failure']['not_found']
}, HTTPStatus.NOT_FOUND.value
except AssertionError as e:
content, status = {'msg': RESTRICTION['users']['remove']['failure']['assertions'].format(reason=e)}, \
HTTPStatus.UNPROCESSABLE_ENTITY.value
except Exception as e:
log.critical(e)
content, status = {
'msg': GENERAL['internal_error']
}, HTTPStatus.INTERNAL_SERVER_ERROR.value
else:
content, status = {
'msg':
RESTRICTION['users']['remove']['success'],
'restriction':
restriction.as_dict(include_groups=True,
include_users=True,
include_resources=True)
}, HTTPStatus.OK.value
finally:
return content, status
def update(newValues: Dict[str, Any]) -> Tuple[Content, HttpStatusCode]:
user = newValues
if user.get('id') is not None:
try:
found_user = User.get(user['id'])
updatable_field_names = ['username', 'password', 'email', 'roles']
for field_name in updatable_field_names:
if user.get(field_name) is not None:
if field_name == 'roles':
new_value = [
Role(name=role_name) for role_name in user['roles']
]
else:
new_value = user[field_name]
setattr(found_user, field_name, new_value)
found_user.save()
except AssertionError as e:
content = {
'msg': USER['update']['failure']['invalid'].format(reason=e)
}
status = 422
except Exception:
content = {'msg': GENERAL['internal_error']}
status = 500
else:
content = {
'msg': USER['update']['success'],
'reservation': found_user.as_dict(include_private=True)
}
status = 201
else:
content = {'msg': GENERAL['bad_request']}
status = 400
return content, status
def update(user):
print('REQ', user)
if user.get('id') is not None:
try:
found_user = User.get(user['id'])
updateable_field_names = ['username', 'password', 'email']
for field_name in updateable_field_names:
if user.get(field_name) is not None:
if field_name == 'roles':
new_value = [
Role(name=role_name) for role_name in user['roles']
]
else:
new_value = user[field_name]
setattr(found_user, field_name, new_value)
found_user.save()
except AssertionError as e:
content = {
'msg': R['update']['failure']['invalid'].format(reason=e)
}
status = 422
except Exception:
content = {'msg': G['internal_error']}
status = 500
else:
content = {
'msg': R['update']['success'],
'reservation': found_user.as_dict
}
status = 201
else:
content = {'msg': G['bad_request']}
status = 400
return content, status
def remove_user(group_id: GroupId,
user_id: UserId) -> Tuple[Content, HttpStatusCode]:
group = None
try:
group = Group.get(group_id)
user = User.get(user_id)
group.remove_user(user)
ReservationVerifier.update_user_reservations_statuses(
user, have_users_permissions_increased=False)
except NoResultFound:
if group is None:
content, status = {
'msg': GROUP['not_found']
}, HTTPStatus.NOT_FOUND.value
else:
content, status = {
'msg': USER['not_found']
}, HTTPStatus.NOT_FOUND.value
except InvalidRequestException:
content, status = {
'msg': GROUP['users']['remove']['failure']['not_found']
}, HTTPStatus.NOT_FOUND.value
except AssertionError as e:
content, status = {'msg': GROUP['users']['remove']['failure']['assertions'].format(reason=e)}, \
HTTPStatus.UNPROCESSABLE_ENTITY.value
except Exception as e:
log.critical(e)
content, status = {
'msg': GENERAL['internal_error']
}, HTTPStatus.INTERNAL_SERVER_ERROR.value
else:
content, status = {
'msg': GROUP['users']['remove']['success'],
'group': group.as_dict()
}, HTTPStatus.OK.value
finally:
return content, status
random.choice(string.ascii_uppercase) for x in range(8))
random_username, random_email = rand_str(
), rand_str() + '@test.com'
user = User(password='******',
email=random_email,
username=random_username)
user.save()
for _ in range(3):
# TODO add spawnAt, terminateAt
content, status = business_create(
dict(userId=user.id, hostname=host, command=cmd))
print(content, status)
print(user)
elif action == '9':
user_id = input('User ID > ')
user = User.get(user_id)
print('[BEFORE] User has {} tasks.'.format(len(user.tasks)))
user.destroy()
tasks_after = Task.query.filter(Task.user_id == user_id).all()
print('[AFTER] User has now {} tasks.'.format(len(tasks_after)))
elif action == '10':
task_id = input('ID > ')
if input('Request full content / Only last lines (y/Enter) > '
) == 'y':
content, status = business_get_log(int(task_id), tail=False)
else:
content, status = business_get_log(int(task_id), tail=True)
print(content, status)
print()
print('\n'.join(content.get('output_lines') or []))
else:
def do_run(self):
time_func = time.perf_counter
start_time = time_func()
# 1. Get list of current reservations
current_reservations = Reservation.current_events()
# FIXME DEBUG ONLY
log.debug(
json.dumps([r.as_dict() for r in current_reservations], indent=4))
for reservation in current_reservations:
# 1. Extract reservation info
uuid = reservation.resource_id
hostname = self.find_hostname(uuid)
user = User.get(reservation.user_id)
username = user.username
if hostname is None or username is None:
log.warning(
'Unable to process the reservation ({}@{}), skipping...'.
format(username, hostname))
continue
# 2. Establish connection to node and find all tty sessions
node_connection = self.connection_manager.single_connection(
hostname)
node_sessions = self.node_tty_sessions(node_connection)
node_processes = self.node_gpu_processes(hostname)
reserved_gpu_process_owners = self.gpu_users(node_processes, uuid)
is_unprivileged = lambda sess: sess[
'USER'] in reserved_gpu_process_owners
intruder_ttys = [
sess for sess in node_sessions if is_unprivileged(sess)
]
try:
# Priviliged user can be ignored on this list
reserved_gpu_process_owners.remove(username)
except ValueError:
pass
finally:
unprivileged_gpu_process_owners = reserved_gpu_process_owners
# 3. Execute protection handlers
for intruder in unprivileged_gpu_process_owners:
violation_data = {
'INTRUDER_USERNAME': intruder,
'RESERVATION_OWNER_USERNAME': username,
'RESERVATION_OWNER_EMAIL': user.email,
'RESERVATION_END': utc2local(reservation.end),
'UUID': uuid,
'GPU_NAME': self.gpu_attr(hostname, uuid,
attribute='name'),
'GPU_ID': self.gpu_attr(hostname, uuid, attribute='index'),
'HOSTNAME': hostname,
'TTY_SESSIONS': intruder_ttys,
'SSH_CONNECTION': node_connection
}
for handler in self.violation_handlers:
handler.trigger_action(violation_data)
end_time = time_func()
execution_time = end_time - start_time
# Hold on until next interval
if execution_time < self.interval:
gevent.sleep(self.interval - execution_time)
waiting_time = time_func() - end_time
total_time = execution_time + waiting_time
log.debug(
'ProtectionService loop took: {:.2f}s (waiting {:.2f}) = {:.2f}'.
format(execution_time, waiting_time, total_time))