""" All the API endpoints for controling processes and tunnels via SSH """ import json from flask import session, redirect, request, Response, make_response from flask_restful import Resource import flask_restful from . import api, islocal if islocal: from .localssh import Ssh from .localtunnelstat import Tunnelstat class GetCert(Resource): """ This class is necessary because I'm not reconfiguring SSHAuthZ to support CORS, but the TES does support CORS """ def post(self): """ takes a public key, returns to the SPA a certificate """ print("in GetCert.post") data = request.get_json() print(data) return {'cert':GetCert.get_cert(data['token'], data['pubkey'], data['signing_url'])} @staticmethod def get_cert(access_token, pub_key, url): """ Sign a pub key into a cert """ import requests print("accss_token {}".format(access_token)) print("pub_key {}".format(pub_key)) print("url {}".format(url)) sess = requests.Session() headers = {"Authorization":"Bearer %s"%access_token} data = {"public_key":pub_key} resp = sess.post(url, json=data, headers=headers, verify=False) print("get_cert returned from its external call") data = resp.json() return data['certificate'] class TestAuth(Resource): """ Tests whether the backend can login to the selected compute Resource """ def get(self): """ tell the SPA if the TES is logged in """ return 'token' in session def get_conn_params(): """ Return parameters relating to the backend compute service Retrieve them from the session (ideally) """ return get_m3_params() def get_m3_params(): """ Hard code the parameters for M3. This will be removed latter factored into a site config file for each compute backend parsed and set by the frontend. """ params = {} params['host'] = 'm3.massive.org.au' params['user'] = 'chines' params['cancelcmd'] = 'scancel {jobid}' params['statcmd'] = '/home/chines/jsonstat.py' params['submitcmd'] = 'sbatch --partition=m3f' params['internalfirewall'] = False return params def get_app_params(): """ Return the parameters for the application retrieved from the Session """ keys = ['startscript', 'paramscmd', 'client','localbind'] returnvalue = {} for k in keys: try: returnvalue[k] = session.get(k) except: # This should be a key exception, i.e. if the key doesn't exist on the session pass print("got app params",returnvalue) return returnvalue class TunnelstatEP(Resource): """ Endpoints used by the WS proxy """ def put(self, authtok): """ update the last used time on a tunnel """ Tunnelstat.update(authtok) def get(self, authtok): """ given an authtoken, return the port of the tunnels """ port = Tunnelstat.getport(authtok) return port class Stat(Resource): """ endpoints to return info on jobs on the backend compute Resource """ def get(self): """ get info on the job from the backend """ params = get_conn_params() res = Ssh.execute(host=params['host'], user=params['user'], cmd=params['statcmd']) if not (res['stderr'] == '' or res['stderr'] is None or res['stderr'] == b''): print(res['stderr']) flask_restful.abort(400, message=res['stderr'].decode()) return json.loads(res['stdout'].decode()) class JobCancel(Resource): """ Terminate a job on the compute backend """ def delete(self, jobid): """ Terminate a job on the backend """ print("in jobcancle jobid is {}".format(jobid)) params = get_conn_params() res = Ssh.execute(host=params['host'], user=params['user'], cmd=params['cancelcmd'].format(jobid=jobid)) if not (res['stderr'] == '' or res['stderr'] is None or res['stderr'] == b''): flask_restful.abort(400, message=res['stderr'].decode()) return res['stdout'].decode() class AppSetup(Resource): """ configure the session for the app the user wants """ def post(self): """ post details of the app to be run or connected to """ data = request.get_json() keys = ['startscript', 'paramscmd', 'client','localbind'] for key in keys: session[key] = data[key] print(data[key]) class JobSubmit(Resource): """ Class dealing the starting a new job on the compute backend """ def post(self): """starting a job is a post, since it changes the state of the backend""" params = get_conn_params() appparams = get_app_params() res = Ssh.execute(host=params['host'], user=params['user'], cmd=params['submitcmd'], stdin=appparams['startscript']) if not (res['stderr'] == '' or res['stderr'] is None or res['stderr'] == b''): print(res['stderr']) flask_restful.abort(400, message=res['stderr'].decode()) return res['stdout'].decode() def gen_authtok(): """ generate a random string suitable for an auth token stored in a cookie """ import random import string return ''.join(random.SystemRandom().choice(string.ascii_uppercase + string.digits) for _ in range(16)) class JobConnect(Resource): """ endpoints for connecting to an existing JobCancel """ def create_tunnel(self, params, appparams, batchhost, data): connectparams = {} if 'paramscmd' in appparams and appparams['paramscmd'] is not None: connectparams['batchhost'] = batchhost paramcmd = 'ssh {batchhost} '.format(batchhost=batchhost) + appparams['paramscmd'] res = Ssh.execute(host=params['host'], user=params['user'], cmd=paramcmd.format(data)) try: connectparams.update(json.loads(res['stdout'])) except json.decoder.JSONDecodeError as e: print(res['stdout']) print(res['stderr']) if not (res['stderr'] == '' or res['stderr'] is None or res['stderr'] == b''): flask_restful.abort(400, message=res['stderr'].decode()) if 'port' in connectparams and 'batchhost' in connectparams: tunnelport, pids = Ssh.tunnel(port=connectparams['port'], batchhost=connectparams['batchhost'], user=params['user'], host=params['host'], internalfirewall=params['internalfirewall'], localbind=appparams['localbind']) authtok = gen_authtok() Tunnelstat.newtunnel(authtok, tunnelport, pids) connectparams['localtunnelport'] = tunnelport connectparams['authtok'] = authtok return connectparams def get(self, jobid, batchhost): """ Connecting to a job is a get operation (i.e. it does not make modifications) """ params = get_conn_params() appparams = get_app_params() data = request.get_json() connectparams = self.create_tunnel(params, appparams, batchhost, data) return self.connect(appparams, connectparams) def connect(self, appparams, connectparams): """ perform the connection either by forking a local client or returning a redirect """ import subprocess if 'cmd' in appparams['client'] and appparams['client']['cmd'] is not None: # We need for fork a local process such as vncviewer or a terminal # We may need a wrapper for local processes to find the correct # process on all OS cmdlist = [] for cmdarg in appparams['client']['cmd']: cmdlist.append(cmdarg.format(**connectparams)) app_process = subprocess.Popen(cmdlist, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # stdout, stderr = app_process.communicate() # if stderr is not "": # return "connected with cmd {} but got error {}".format(cmdlist,stderr) elif 'redir' in appparams['client'] and appparams['client']['redir'] is not None: template_response = Response() template_response.set_cookie(key='twsproxyauth', value=connectparams['authtok']) twsproxy = 'http://localhost:4000/' response = make_response(redirect(twsproxy+appparams['client']['redir']. format(**connectparams))) response.set_cookie('twsproxyauth', connectparams['authtok']) return response return "Connecting with cmd {}".format(cmdlist) api.add_resource(TunnelstatEP, '/tunnelstat/<string:authtok>') api.add_resource(GetCert, '/getcert') api.add_resource(Stat, '/stat') api.add_resource(JobCancel, '/cancel/<int:jobid>') api.add_resource(JobSubmit, '/submit') api.add_resource(JobConnect, '/connect/<int:jobid>/<string:batchhost>') api.add_resource(AppSetup, '/appsetup')