def check_user_session(user):
print('check_user_session')
existing_sessions = {}
for file_name in glob.glob(
parameters.get_parameter('dcv', 'session_location') + '/*'):
if user in file_name:
ignore = False
session_info = open_yaml(file_name)
try:
existing_sessions[int(
session_info['session_number'])] = session_info
except:
print("No session number detected for " + str(file_name) +
". This may not be a DCV file")
ignore = True
if ignore is False:
session_job_id = session_info['job_id']
print(session_job_id + ' detected')
try:
check_job_cmd = run_command([
parameters.get_parameter('pbs', 'qstat'), '-f',
session_job_id, '-F', 'json'
], 'check_output')
check_job_status = json.loads(
check_job_cmd["output"].decode('utf-8'))
for job, job_data in check_job_status['Jobs'].items():
if 'exec_host' in job_data.keys():
exec_host = job_data['exec_host'].split('/')[0]
update_yaml(file_name, exec_host)
except Exception as e:
clean_session(user, session_info['session_number'])
print(existing_sessions)
return existing_sessions
def clean_session(user, session_number):
print('clean session')
for file_name in glob.glob(
parameters.get_parameter('dcv', 'session_location') + '/*'):
if user in file_name:
if str(session_number) in file_name:
with open(file_name) as f:
session_info = yaml.safe_load(f)
print(
'remove auth_dir/session_name on remote host and qdel job')
commands = [
parameters.get_parameter('dcv', 'bin') +
' close-session ' + session_info['session_id'],
'rm -rf ' + parameters.get_parameter('dcv', 'auth_dir') +
'/' + session_info['session_id'],
parameters.get_parameter('pbs', 'qdel') + ' ' +
session_info['job_id']
]
proc = subprocess.Popen(["ssh " + session_info['host']],
preexec_fn=demote(user),
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
universal_newlines=True,
bufsize=0,
shell=True)
proc.stdin.write('\n'.join(commands))
proc.stdin.close()
print('remove yaml')
os.remove(file_name)
return True
def validate_ldap(username, password):
ldap_host = parameters.get_parameter('ldap', 'host')
base_dn = parameters.get_parameter('ldap', 'base_dn')
user_dn = 'uid={},ou=people,{}'.format(username, base_dn)
con = ldap.initialize('ldap://{}'.format(ldap_host))
try:
con.bind_s(user_dn, password, ldap.AUTH_SIMPLE)
session['username'] = username
# Check if user has sudo permissions
sudoers_search_base = "ou=Sudoers,dc=soca,dc=local"
sudoers_search_scope = ldap.SCOPE_SUBTREE
sudoers_filter = 'cn=' + username
is_sudo = con.search_s(sudoers_search_base, sudoers_search_scope, sudoers_filter)
if is_sudo.__len__() > 0:
session['sudoers'] = True
else:
session['sudoers'] = False
return {'success': True,
'message': ''}
except ldap.INVALID_CREDENTIALS:
return {'success': False,
'message': 'Invalid credentials.'}
except ldap.SERVER_DOWN:
return {'success': False,
'message': 'LDAP server is down.'}
def verify_sudo_permissions(username):
ldap_host = parameters.get_parameter('ldap', 'host')
base_dn = parameters.get_parameter('ldap', 'base_dn')
user_dn = 'uid={},ou=people,{}'.format(username, base_dn)
con = ldap.initialize('ldap://{}'.format(ldap_host))
try:
# Check if user has sudo permissions
sudoers_search_base = "ou=Sudoers,dc=soca,dc=local"
sudoers_search_scope = ldap.SCOPE_SUBTREE
sudoers_filter = 'cn=' + username
is_sudo = con.search_s(sudoers_search_base, sudoers_search_scope, sudoers_filter)
if is_sudo.__len__() > 0:
return {'success': True,
'message': username + ' is a valid SUDO user'}
else:
return {'success': False,
'message': username + ' is not a SUDO user'}
except ldap.INVALID_CREDENTIALS:
return {'success': False,
'message': 'Invalid credentials.'}
except ldap.SERVER_DOWN:
return {'success': False,
'message': 'LDAP server is down.'}
def f():
username = session['username']
user_private_key_path = '/data/home/' + username + '/.ssh/id_rsa'
generate_ppk = [
'unix/puttygen', user_private_key_path, '-o',
parameters.get_parameter('ssh', 'private_key_location') + '/' +
username + '_soca_privatekey.ppk'
]
subprocess.call(generate_ppk)
os.chmod(
parameters.get_parameter('ssh', 'private_key_location') + '/' +
username + '_soca_privatekey.ppk', 0o700)
return send_file(parameters.get_parameter('ssh', 'private_key_location') +
'/' + username + '_soca_privatekey.ppk',
as_attachment=True,
attachment_filename=username + '_soca_privatekey.ppk')
def oauth():
next_url = request.args.get('state')
sso_auth = auth.sso_authorization(request.args.get('code'))
cognito_root_url = parameters.get_parameter("cognito", "cognito_root_url")
if sso_auth['success'] is True:
return redirect(cognito_root_url+next_url)
else:
if sso_auth['message'] == 'user_not_found':
return redirect(cognito_root_url)
else:
return str(sso_auth['message'])
def validate_account():
if 'username' in session:
return f()
else:
enable_sso = (parameters.get_parameter("cognito",
"enable_sso")).lower()
if enable_sso == "true":
data = {
'redirect_uri':
parameters.get_parameter("cognito",
"cognito_callback_url"),
'client_id':
parameters.get_parameter("cognito", "cognito_app_id"),
'response_type':
'code',
'state':
request.path
}
oauth_url = parameters.get_parameter(
"cognito", "cognito_oauth_authorize_endpoint"
) + '?' + urllib.parse.urlencode(data)
return redirect(oauth_url)
else:
return redirect('/login')
def get_all_users():
all_ldap_users = {}
ldap_host = parameters.get_parameter('ldap', 'host')
user_search_base = "ou=People,dc=soca,dc=local"
user_search_scope = ldap.SCOPE_SUBTREE
user_filter = 'uid=*'
con = ldap.initialize('ldap://{}'.format(ldap_host))
users = con.search_s(user_search_base, user_search_scope, user_filter)
for user in users:
user_base = user[0]
username = user[1]['uid'][0].decode('utf-8')
all_ldap_users[username] = user_base
return all_ldap_users
def build_dcv_connect_client(user, session_number):
yaml_file = 'dcv_' + user + '_' + str(session_number) + '.yml'
session_data = open_yaml(
parameters.get_parameter('dcv', 'session_location') + '/' + yaml_file)
session_file = '''
[version]
format=1.0
[connect]
host=''' + parameters.get_aligo_configuration()['LoadBalancerDNSName'] + '''
port=443
weburlpath=/''' + session_data['host'] + '''
sessionid=''' + session_data['session_id'] + '''
user=''' + user + '''
authToken=''' + session_data['session_password'] + '''
'''
return session_file
def sso_authorization(code):
authorization = 'Basic ' + base64.b64encode(
parameters.get_parameter("cognito", "cognito_app_id").encode() +
':'.encode() + parameters.get_parameter(
"cognito", "cognito_app_secret").encode()).decode()
headers = {
'Content-Type': 'application/x-www-form-urlencoded',
'Authorization': authorization
}
data = {
'grant_type': 'authorization_code',
'client_id': parameters.get_parameter("cognito", "cognito_app_id"),
'code': code,
'redirect_uri': parameters.get_parameter("cognito",
"cognito_callback_url")
}
oauth_token = requests.post(parameters.get_parameter(
"cognito", "cognito_oauth_token_endpoint"),
data=data,
headers=headers).json()
id_token = oauth_token['id_token']
access_token = oauth_token['access_token']
headers = jwt.get_unverified_headers(id_token)
keys = requests.get(
parameters.get_parameter(
"cognito", "cognito_jws_keys_endpoint")).json().get('keys')
key = list(filter(lambda k: k['kid'] == headers['kid'], keys)).pop()
claims = jwt.decode(id_token,
key,
access_token=access_token,
algorithms=[key['alg']],
audience=parameters.get_parameter(
"cognito", "cognito_app_id"))
if claims:
try:
username = claims['email'].split('@')[0]
except Exception as err:
return {
'success':
False,
'message':
'Error reading SSO claims. ' + str(claims) + ' Err: ' +
str(err)
}
# Simply check if user exist
# We could do a simply ldap lookup, but making sure user has a private key is more important
# without private key, a user (even with valid ldap account) won't be able to do anything
check_if_file_exist = os.path.isfile('/data/home/' + username +
'/.ssh/id_rsa')
if check_if_file_exist is True:
# Valid user, create session
session['username'] = username
# verify sudo permission
if openldap.verify_sudo_permissions(username)["success"] is True:
session["sudoers"] = True
else:
session["sudoers"] = False
return {'success': True, 'message': ''}
else:
return {'success': False, 'message': 'user_not_found'}
else:
return {'success': False, 'message': 'SSO error. ' + str(claims)}
def build_qsub(session_owner, session_number, walltime, instance_type):
session_id = str(uuid.uuid4())
command_dcv_create = parameters.get_parameter(
'dcv', 'bin'
) + " create-session --user " + session_owner + " --owner " + session_owner + " " + session_id
session_password = ''.join(
random.choice(string.ascii_uppercase + string.digits +
string.ascii_lowercase) for _ in range(80))
params = {
'pbs_job_name':
'Desktop' + str(session_number),
'pbs_queue':
'desktop',
'pbs_project':
'gui',
'instance_type':
instance_type,
'dcv_create_session':
command_dcv_create,
'session_password':
session_password,
'session_password_b64':
(base64.b64encode(session_password.encode('utf-8'))).decode('utf-8'),
'walltime':
walltime
}
qsub_command = '''<<EOF
#PBS -N ''' + params['pbs_job_name'] + '''
#PBS -q ''' + params['pbs_queue'] + '''
#PBS -P ''' + params['pbs_project'] + '''
#PBS -l walltime=''' + params['walltime'] + '''
#PBS -l instance_type=''' + params['instance_type'] + '''
#PBS -e /dev/null
#PBS -o /dev/null
# Create the DCV Session
''' + params['dcv_create_session'] + '''
# Query dcvsimpleauth with add-user
echo ''' + params[
'session_password_b64'] + ''' | base64 --decode | /usr/libexec/dcvsimpleextauth.py add-user --user ''' + session_owner + ''' --session ''' + session_id + ''' --auth-dir ''' + parameters.get_parameter(
'dcv', 'auth_dir') + '''
# Uncomment if you want to disable Gnome Lock Screen (require webui restart)
# GSETTINGS=$(which gsettings)
# $GSETTINGS set org.gnome.desktop.lockdown disable-lock-screen true
# $GSETTINGS set org.gnome.desktop.session idle-delay 0
# Keep job open
while true
do
session_keepalive=$(/usr/bin/dcv list-sessions | grep ''' + session_id + ''' | wc -l)
if [ $session_keepalive -ne 1 ]
then
exit 0
fi
sleep 3600
done
EOF
'''
yaml_config = parameters.get_parameter(
'dcv', 'session_location') + '/dcv_' + session_owner + '_' + str(
session_number) + '.yml'
if path.exists(yaml_config):
print(yaml_config + ' already exist.')
return False
launch_job_session = run_command(
['su', session_owner, '-c', '/opt/pbs/bin/qsub ' + qsub_command],
"check_output")
if launch_job_session["success"] is True:
job_id = ((launch_job_session["output"].decode('utf-8')
).rstrip().lstrip()).split("\n")[-1].split('.')[0]
dcv_session_data = dict(job_id=job_id,
host='tbd',
state='pending',
session_id=session_id,
session_password=session_password,
session_number=int(session_number))
with open(yaml_config, 'w') as outfile:
yaml.dump(dcv_session_data, outfile, default_flow_style=False)
os.chmod(yaml_config, 0o700)
return True
else:
return str(launch_job_session["output"])
