def do_create(user: Dict[str, Any]) -> Tuple[Content, HttpStatusCode]:
try:
new_user = User(username=user['username'],
email=user['email'],
password=user['password'],
roles=[Role(name='user')])
new_user.save()
try:
default_groups = Group.get_default_groups()
for group in default_groups:
group.add_user(new_user)
except Exception:
log.warning(
"User has been created, but not added to default group.")
except AssertionError as e:
content = {
'msg': USER['create']['failure']['invalid'].format(reason=e)
}
status = 422
except IntegrityError:
content = {'msg': USER['create']['failure']['duplicate']}
status = 409
except Exception as e:
content = {'msg': GENERAL['internal_error'] + str(e)}
status = 500
else:
content = {
'msg': USER['create']['success'],
'user': new_user.as_dict(include_private=True)
}
status = 201
finally:
return content, status
def login(user: Dict[str, Any]) -> Tuple[Content, HttpStatusCode]:
try:
current_user = User.find_by_username(user['username'])
assert User.verify_hash(user['password'], current_user.password), \
USER['login']['failure']['credentials']
except NoResultFound:
content = {'msg': USER['not_found']}
status = 404
except AssertionError as error_message:
content = {'msg': str(error_message)}
status = 401
except Exception:
content = {'msg': GENERAL['internal_error']}
status = 500
else:
content = {
'msg':
USER['login']['success'].format(username=current_user.username),
'access_token':
create_access_token(identity=current_user.id, fresh=True),
'refresh_token':
create_refresh_token(identity=current_user.id)
}
status = 200
finally:
return content, status
def get() -> Tuple[List[Any], HttpStatusCode]:
claims = get_jwt_claims()
include_private = 'admin' in claims['roles']
return [
user.as_dict(include_private=include_private) for user in User.all()
], 200
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 test_exception_on_creating_user_with_not_unique_username(tables):
with pytest.raises(IntegrityError):
duplicated_username = '******'
password = '******'
existing_user = User(username=duplicated_username, password=password)
duplicated_user = User(username=duplicated_username, password=password)
existing_user.save()
duplicated_user.save()
def do_create(user):
try:
new_user = User(username=user['username'],
email=user['email'],
password=user['password'],
roles=[Role(name='user')])
new_user.save()
except AssertionError as e:
content = {'msg': R['create']['failure']['invalid'].format(reason=e)}
status = 422
except IntegrityError:
content = {'msg': R['create']['failure']['duplicate']}
status = 409
except Exception as e:
content = {'msg': G['internal_error'] + str(e)}
status = 500
else:
content = {'msg': R['create']['success'], 'user': new_user.as_dict}
status = 201
finally:
return content, status
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 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 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 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 trigger_action(self, violation_data: Dict[str, Any]) -> None:
'''Contains business logic for intruder and admin email notifications.
It relies on early returns if any error occures.
:param violation_data: data received from ProtectionService
'''
# Expect certain keys beforehand
assert {
'INTRUDER_USERNAME', 'RESERVATION_OWNER_USERNAME',
'RESERVATION_OWNER_EMAIL', 'RESERVATION_END', 'UUID', 'HOSTNAME'
}.issubset(violation_data), 'Invalid keys in violation_data'
if not self._test_smtp_configuration():
return
try:
# Fetch intruder email address and extend violation data
intruder_email = User.find_by_username(
violation_data['INTRUDER_USERNAME']).email
except NoResultFound as e:
intruder_email = None
log.warning(e)
finally:
violation_data['INTRUDER_EMAIL'] = intruder_email
if not intruder_email:
# Intruder has no account or email assigned, try notify admin then
timer = self._get_timer(violation_data['INTRUDER_USERNAME'])
if MAILBOT.NOTIFY_ADMIN and self._time_to_resend(timer,
to_admin=True):
self._email_admin(violation_data, timer)
return
# Intruder has account and email address, try email him and admin then
timer = self._get_timer(intruder_email)
if MAILBOT.NOTIFY_INTRUDER and self._time_to_resend(timer):
self._email_intruder(intruder_email, violation_data, timer)
if MAILBOT.NOTIFY_ADMIN and self._time_to_resend(timer, to_admin=True):
self._email_admin(violation_data, timer)
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 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
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
spawnAt=timenow,
terminateAt=timenow))
print(content, status)
elif action == '7':
task_id = input('ID > ')
content, status = business_destroy(int(task_id))
print(content, status)
elif action == '8':
import string
import random
rand_str = lambda: ''.join(
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)))
def new_admin():
return User(username='******',
password='******',
roles=[Role(name='user'),
Role(name='admin')])
def new_user_2():
return User(username='******',
password='******',
roles=[Role(name='user')])
def new_user():
return User(username='******',
password='******',
roles=[Role(name='user')])
def get():
return [
user.as_dict for user in User.all()
], 200
def __init__(self):
init_db()
# Prepare empty ORM object
self.new_user = User()
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))
class AccountCreator:
'''
It asks all the necessary questions in order to set up a new account.
Takes user's input from CLI.
Exits on Ctrl+C
'''
def __init__(self):
init_db()
# Prepare empty ORM object
self.new_user = User()
def run_prompt(self):
self._ask_for_username()
self._ask_for_email()
self._ask_for_password()
self._ask_for_role()
self._create_user()
def _create_user(self):
try:
self.new_user.save()
except Exception as e:
click.echo(
'Account creation failed due to an error: {}.'.format(e))
else:
click.echo('Account created successfully.')
def _ask_for_username(self):
'''
Process is repeated until username becomes valid.
If so, it assignes the vaule to the ORM object.
'''
valid_username_provided = False
while not valid_username_provided:
try:
username = click.prompt('[1/4] UNIX username', type=str)
self.new_user.username = username
except click.Abort:
raise
except Exception as e:
click.echo('Invalid username: {reason}.'.format(reason=e))
else:
valid_username_provided = True
def _ask_for_email(self):
valid_email_provided = False
while not valid_email_provided:
try:
email = click.prompt(
'[2/4] email (for TensorHive warnings only)', type=str)
self.new_user.email = email
except click.Abort:
raise
except Exception as e:
click.echo('Invalid email: {reason}.'.format(reason=e))
else:
valid_email_provided = True
def _ask_for_password(self):
# Useful aliases
prompt_for_password = lambda message: click.prompt(
message, type=str, hide_input=True)
password_length_requirement = 'at least {} characters'.format(
self.new_user.min_password_length)
first_password_message = '[3/4] password ({})'.format(
password_length_requirement)
repeated_password_message = '[3/4] repeat password'
valid_password_provided = False
while not valid_password_provided:
try:
password1 = prompt_for_password(message=first_password_message)
self.new_user.password = password1
password2 = prompt_for_password(
message=repeated_password_message)
assert password1 == password2, 'Passwords don\'t match, please try again.'
except click.Abort:
raise
except Exception as error_msg:
click.echo(str(error_msg))
else:
valid_password_provided = True
def _ask_for_role(self):
try:
make_admin = click.confirm('[4/4] admin account?', default=False)
# TODO Refactor roles: admin or not instead of two mutually exclusive 'admin' and 'user
self.new_user.roles.append(Role(name='user'))
if make_admin:
self.new_user.roles.append(Role(name='admin'))
except click.Abort:
raise
except Exception:
click.echo('Unknown error - could not assign role.')
def test_exception_on_creating_user_with_invalid_username(tables, test_name, test_username):
with pytest.raises(AssertionError):
User(username=test_username, password='******').save()
def test_exception_on_creating_user_with_no_password(tables):
with pytest.raises(IntegrityError):
User(username='******').save()