Skip to content
Snippets Groups Projects
apiendpoints.py 26.88 KiB
"""
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
from flask import render_template
import flask_restful

from . import api, app
from .sshwrapper import Ssh, SshAgentException, SftpPermissionException, SftpException, SshCtrlException, SshExecException
from .tunnelstat import SSHSession
import subprocess

class AppParamsException(Exception):
    pass

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
        """
        import logging
        logger = logging.getLogger()
        try:
            data = request.get_json()
            if type(data) is not dict:
                return apiabort(400, message="necessary data not provided in post")
            response = {'cert':GetCert.get_cert(data['token'], data['pubkey'], data['signing_url'])}
            return response
        except:
            import traceback
            logger.error('Failed to get certificate')
            logger.error(traceback.format_exc())
            flask_restful.abort(400,message="an error occured generating your certificate")

    @staticmethod
    def get_cert(access_token, pub_key, url):
        """
        Sign a pub key into a cert
        """
        import requests
        import logging
        logger = logging.getLogger()
        sess = requests.Session()
        headers = {"Authorization":"Bearer %s"%access_token}
        data = {"public_key":pub_key}
        resp = sess.post(url, json=data, headers=headers, verify=False)
        data = resp.json()
        return data['certificate']


class SSHAgent(Resource):
    def post(self):
        import logging
        logger = logging.getLogger()
        try:
            session.permanent = True
            from .tunnelstat import SSHSession
            sshsess = SSHSession.get_sshsession()
            data = request.get_json()
            if type(data) is not dict:
                return apiabort(400, message="necessary data not provided in post")
            sshsess.add_keycert(key=data['key'],cert=data['cert'])
            sshsess.refresh()
            return "OK"
        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            logger.error('failed to add ssh key to the agent')
            flask_restful.abort(500,message="failed to add the ssh key to the agent")

    def get(self):
        from .tunnelstat import SSHSession
        import logging
        logger = logging.getLogger()
        try:
            sshsess = SSHSession.get_sshsession()
            if sshsess.socket == None:
                return [] # The agent hasn't even been started yet
            return sshsess.get_cert_contents()
        except Exception as e:
            import traceback
            logger.error('SSHAgent.get: Exception {}'.format(e))
            logger.error(traceback.format_exc())
            flask_restful.abort(500,message="failed to query the ssh-agent")


    def delete(self):
        import logging
        logger = logging.getLogger()
        from .tunnelstat import SSHSession
        try:
            sshsess = SSHSession.get_sshsession()
            sshsess.kill()
            SSHSession.remove_sshsession()
            session.pop('sshsessid', None)
            return []
        except Exception as e:
            import traceback
            logger.error(traceback.format_exc())
            flask_restful.abort(500,message="failed to shut down ssh-agent")

def get_conn_params():
    """
    Return parameters relating to the backend compute service
    Retrieve them from the query string
    """
    import logging
    logger = logging.getLogger()
    identitystr = request.args.get('identity')
    if type(identitystr) is str:
        identityparams = json.loads(identitystr)
    else:
        identityparams = {}
    try:
        interfacestr = request.args.get('interface')
        if type(interfacestr) is str:
            interfaceparams = json.loads(interfacestr)
        else:
            interfaceparams = {}
    except:
        interfaceparams = {}
        pass
    try:
        pathstr = request.args.get('path')
        if type(pathstr) is str:
            pathparams = json.loads(pathstr)
        else:
            pathparams = None
        cdstr = request.args.get('cd')
        if type(cdstr) is str:
            cdparams = json.loads(cdstr)
        else:
            cdparams = None
    except:
        pathparams = None
        cdparams = None
    try:
        appstr = request.args.get('app')
        if type(appstr) is str:
            appparams = json.loads(appstr)
        else:
            appparams = {}
    except:
        appparams = {}
    try:
        appinstancestr = request.args.get('appinstance')
        if type(appinstancestr) is str:
            appinstanceparams = json.loads(appinstancestr)
        else:
            appinstanceparams = {}
    except:
        appinstanceparams = {}
    params = {}

    params['identity'] = identityparams
    params['interface'] = interfaceparams
    params['app'] = appparams
    params['appinstance'] = appinstanceparams


    params.update(interfaceparams)
    params['user'] = identityparams['username']
    params['host'] = identityparams['site']['host']
    params['path'] = pathparams
    params['cd'] = cdparams

    return params



class TunnelstatEP(Resource):
    """
    Endpoints used by the WS proxy
    The Frontend should never call this
    """

    def get(self, authtok):
        """
        given an authtoken, return the port of the tunnels
        """
        import logging
        logger = logging.getLogger()
        from . import sshsessions
        port = None
        try:
            for (sessid,sshsess) in sshsessions.items():
                for (tok,port) in sshsess.port.items():
                    if tok == authtok:
                        return port
        except:
            logger.error("exception in TunnelstatEP.get")
            import traceback
            logger.error(traceback.format_exc())
        return None

class ContactUs(Resource):
    def post(self):
        import tempfile
        data = request.get_json()
        f = tempfile.NamedTemporaryFile(mode='w+b',dir=app.config['MESSAGES'],delete=False)
        f.write(json.dumps(data).encode())
        f.close()
        return

def wrap_execute(sess, host, user, cmd, bastion=None,stdin=None, sshport="22", bastionsshport="22"):
    """
    This function is supposed to interpret all the possible exceptions from Ssh.execute and generate
    appropriate HTTP errors (both status codes and messages)

    There are a range of situations:
    it might raise an SshExecException (non-zero return code)
    it might return valid data or not
    it might return stderr or not
    the data on stderr might be valid json or not.
    """
    import logging
    logger = logging.getLogger()
    try:
        res = Ssh.execute(sess,  host, user, cmd, bastion, stdin, sshport, bastionsshport)
        if not (res['stderr'] == '' or res['stderr'] is None or res['stderr'] == b''):
            logger.error(res['stderr'])
            #flask_restful.abort(400, message=res['stderr'].decode())
            return apiabort(400, message=res['stderr'].decode(), sess=sess)

        try:
            if res['stdout'].decode() != '':
                data = json.loads(res['stdout'].decode())
                return data
            else:
                return None
        except json.decoder.JSONDecodeError:
            return None
            #return apiabort(500, data = {'stdout':res['stdout'].decode(),'stderr':res['stderr'].decode()})
        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            return apiabort(400, message=e, sess=sess)

    except SshCtrlException as e:
        return apiabort(401,message="{} {}".format(host,e), sess=sess, logdata = json.dumps({'user':user,'cmd':cmd,'host':host}))
    # This exception is raised if the remove command had a non-zero return code
    except SshExecException as e:
        return apiabort(400, message="{}".format(e), sess=sess)
    # Any other exceptions. This should never be reached.
    except subprocess.TimeoutExpired as e:
        return apiabort(504, message="{}".format(e), sess=sess)
    except Exception as e:
        import traceback
        logger.error('JobStat.get: Exception {}'.format(e)
        )
        logger.error(traceback.format_exc())
        return apiabort(500,message="{}".format(e), sess=sess)

class JobStat(Resource):

    """
    endpoints to return info on jobs on the backend compute Resource
    """
    def get(self):
        """
        get info on the job from the backend
        """
        import logging
        logger = logging.getLogger()
        logger.debug('entered jobstat')
        try:
            sshsess = SSHSession.get_sshsession()
        except Exception as e:
            return apiabort(500, message="{}".format(e))
        try:
            s = request.args.get('statcmd')
            if type(s) is str:
                cmd = json.loads(s)
            else:
                cmd = None
            s = request.args.get('host')
            if type(s) is str:
                host = json.loads(s)
            else:
                host = None
            s = request.args.get('username')
            if type(s) is str:
                user = json.loads(s)
            else:
                user = None
        except (TypeError, KeyError) as e:
            import traceback
            return apiabort(400, message="Missing required parameter {}\n{}".format(e,traceback.format_exc()), sess=sshsess)

        rv = wrap_execute(sshsess, host=host, user=user, cmd=cmd)
        return rv

class MkDir(Resource):
    def post(self):
        import logging
        logger = logging.getLogger()
        try:
            data = request.get_json()
            params = get_conn_params()
            sshsess = SSHSession.get_sshsession()
            site = params['identity']['site']
            if type(data) is not dict:
                return apiabort(400, message="necessary parameters not provided in the post information")
            if 'dtnport' in site:
                sshport = site['dtnport']
            else:
                sshport = "22"
            try:
                Ssh.sftpmkdir(sshsess, host=params['identity']['site']['host'],
                               user=params['identity']['username'], path=params['path'],name=data['name'], sshport=sshport)
            except SftpPermissionException as e:
                #return json.dumps({'message':"You don't have permission to make a directory there"}), 403
                #flask_restful.abort(403,message="You don't have permission to make a directory there")
                return apiabort(403,message="You don't have permission to make a directory there")
            except SftpException as e:
                #return json.dumps({'message':"Something went wrong making that directory"}), 500
                #flask_restful.abort(500,message="Something went wrong making that directory")
                return apiabort(500,message="Something went wrong making that directory")
            except Exception as e:
                import traceback
                logger.error(traceback.format_exc())
                #flask_restful.abort(500,message="Something went wrong creating that directory, probably a bug")
                return apiabort(500,message="Something went wrong creating that directory, probably a bug")

            return
        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            #flask_restful.abort(500,message="mkdir failed in some unexpected way")
            return apiabort(500,message="mkdir failed in some unexpected way")

class DirList(Resource):
    def get(self):
        import logging
        logger = logging.getLogger()
        try:
            params = get_conn_params()
            sshsess = SSHSession.get_sshsession()
            site = params['identity']['site']
            if 'dtnport' in site:
                sshport = site['dtnport']
            else:
                sshport = "22"
            path = params['path']
            cd = params['cd']
            if path == "":
                path = "."
            if cd == "":
                cd = "."
            try:
                if 'lscmd' in site and site['lscmd'] is not None and site['lscmd'] != "":
                    res = Ssh.execute(sshsess, host=params['identity']['site']['host'], user=params['identity']['username'],
                                        sshport=sshport,
                                        cmd="{} {} {}".format(site['lscmd'],path,cd))
                    try:
                        dirls = json.loads(res['stdout'].decode())
                    except:
                        #return json.dumps({'message':"You don't have permission to view that directory"}), 401
                        #flask_restful.abort(401,message="You don't have permission to view that directory")
                        return apiabort(403,message="You don't have permission to view that directory")
                else:
                    dirls = Ssh.sftpls(sshsess, host=params['identity']['site']['host'],
                                       user=params['identity']['username'], path=params['path'],changepath=params['cd'], sshport=sshport)
                return dirls
            except SshCtrlException as e:
                if ("{}".format(e) != ''):
                    return apiabort(400,message="{}".format(e))
                else:
                    # This is the most common error, if the user disconnected from the network and the api server cleaned up all the agents and control sockets
                    return apiabort(401)
        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            #flask_restful.abort(500,message="dirlist failed in some unexpected way")
            return apiabort(500,message="dirlist failed in some unexpected way")

class JobCancel(Resource):
    """
    Terminate a job on the compute backend
    """
    def delete(self, jobid):
        import logging
        logger = logging.getLogger()
        """
        Terminate a job on the backend
        """
        sshsess = None
        try:
            params = get_conn_params()
            sshsess = SSHSession.get_sshsession()
            res = wrap_execute(sshsess, host=params['identity']['site']['host'], user=params['identity']['username'],
                                  cmd=params['interface']['cancelcmd'].format(jobid=jobid))
            return res

        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            #return json.dumps({'message':"jobcancel failed in some unexpected way"}), 500
            #flask_restful.abort(500,message="jobcancel failed in some unexpected way")
            return apiabort(500,message="jobcancel failed in some unexpected way", sess=sshsess)

class RemoteCommand(Resource):
    # jobsubmit res = wrap_execute(sshsess, host=params['identity']['site']['host'], user=params['identity']['username'], cmd=params['interface']['submitcmd'], stdin=script)
    # appinstance rv = wrap_execute(sshsess, host=batchhost, bastion=loginhost, user=username, cmd=paramscmd) ... part of path
    # jobscancel res = wrap_execute(sshsess, host=params['identity']['site']['host'], user=params['identity']['username'], cmd=params['interface']['cancelcmd'].format(jobid=jobid))
    def post(self):
        import logging
        logger=logging.getLogger()
        try:
            sshsess = SSHSession.get_sshsession()

            data = request.get_json()
            cmd = None
            bastion = None
            if type(data) is dict:
                user = data['user']
                host = data['host']
                if 'bastion' in data:
                    bastion = data['bastion']
                else:
                    bastion = None
                cmd = data['cmd']
            else:
                return apiabort(400, message="necessary data not provided")
            try:
                if 'input' in data:
                    input = data['input'].format(**data)
                else:
                    input = None
                res = wrap_execute(sshsess, host, user, cmd, bastion, stdin=input)
                return res
            except Exception as e:
                import traceback
                logger.error(e)
                logger.error(traceback.format_exc())
                return apiabort(500,message="Remote Command failed in some unexpected way", sess=sshsess)
        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            return apiabort(500,message="Remote command failed in some unexpected way")

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"""
        import logging
        logger=logging.getLogger()
        sshsess = None
        try:
            params = get_conn_params()
            sshsess = SSHSession.get_sshsession()
            data = request.get_json()
            if type(data) is not dict:
                return apiabort(400, message="necessary data not provided in the post information")
            try:
                script = data['app']['startscript'].format(**data)
            except Exception as e:
                import traceback
                logger.error(e)
                logger.error(traceback.format_exc())
                logger.error('formating data')
                logger.error(data)
                logger.error('end formating data')
                logger.error('body')
                logger.error(request.data)
                logger.error('end body')
                #return json.dumps({'message':'Incomplete job information was passed to the backend.'}), 400
                #flask_restful.abort(400, message='Incomplete job information was passed to the backend.')
                return apiabort(400, message='Incomplete job information was passed to the backend.', sess=sshsess)

            res = wrap_execute(sshsess, host=params['identity']['site']['host'], user=params['identity']['username'],
                          cmd=params['interface']['submitcmd'], stdin=script)
            return res
        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            #return json.dumps({'message':"JobSubmission failed in some unexpected way"}), 500
            #flask_restful.abort(500,message="JobSubmission failed in some unexpected way")
            return apiabort(500,message="JobSubmission failed in some unexpected way", sess=sshsess)

def gen_authtok():
    """
    generate a random string suitable for an auth token stored in a cookie
    """
    import random
    import string
    import logging
    logger=logging.getLogger()
    return ''.join(random.SystemRandom().choice(string.ascii_uppercase +
                                                string.digits) for _ in range(16))

def apiabort(code, message=None, data = None, sess = None, logdata = None):
    import logging
    import traceback
    from flask import g
    import time
    logger = logging.getLogger('apiabort')
    logger.error("request time {}".format(time.time() - g.start))
    if sess is not None:
        logger.error("certs {}".format(sess.get_certs()))
    logger.error(logdata)
    logger.error("{} {} {}".format(code,data,message))
    logger.error("{}".format(traceback.format_exc()))
    if data is not None:
        return data, code
    else:
        return {'message':message}, code

class AppUrl(Resource):
    def get(self):
        import logging
        logger = logging.getLogger()
        try:
            s = request.args.get('app')
            if type(s) is str:
                appdef = json.loads(s)
            else:
                return apiabort(400, message="Failed to provide information on the application")
            s = request.args.get('appinst')
            if type(s) is str:
                inst = json.loads(s)
            else:
                return apiabort(400, message="failed to provide information in the application instance")
            inst['twsproxy']='{twsproxy}'
            inst['twsproxyauth'] = request.cookies.get('twsproxyauth')
            logger.debug('twsproxyauth {} {}'.format(inst['twsproxyauth'],request.cookies.get('twsproxyauth')))
            url = "{}/{}".format("{twsproxy}",appdef['client']['redir'].format(**inst))
            return url
        except Exception as e:
            import traceback
            logger.error('AppUrl failed')
            logger.error(e)
            logger.error(traceback.format_exc())
            #flask_restful.abort(500,message="AppUrl failed in some unexpected way")
            return apiabort(500,message="AppUrl failed in some unexpected way")

class SetTWSProxyAuth(Resource):
    def get(self, authtok):
        import urllib.parse
        import logging
        import datetime
        logger = logging.getLogger()
        s = request.args.get('redirect')
        if type(s) is str:
            url = urllib.parse.unquote(s)
        else:
            return apiabort(400, message="no redirect provided")
        logger.debug('SetTWSProxyAuth will redirect to {}'.format(url))
        response = make_response(redirect(url))
        response.set_cookie('twsproxyauth', authtok, secure=True, httponly=True, expires = datetime.datetime.now()+datetime.timedelta(days=30))
        return response


class Ping(Resource):
    def get(self):
        return None, 201

class AppInstance(Resource):
    def get(self, username, loginhost, batchhost, jobid):
        """Run a command to get things like password and port number
        command is passed as a query string"""
        import logging
        import hashlib
        logger=logging.getLogger()
        events = logging.getLogger('connection_events')
        sshsess = SSHSession.get_sshsession()
        sessionId = hashlib.sha1(sshsess.sshsessid.encode()).hexdigest()
        s = request.args.get('cmd')
        if type(s) is str:
            paramscmd = json.loads(s).format(jobid=jobid)
        else:
            return apiabort(400, message="no command provided to find information about the application instance")
        try:
            rv = wrap_execute(sshsess, host=batchhost, bastion=loginhost, user=username, cmd=paramscmd)
            events.info('connect {}'.format(json.dumps({'user':username, 'jobid':jobid, 'loginhost': loginhost, 'sessid':sessionId})))
            return rv
        except Exception as e:
            import traceback
            logger.error('AppInstance failed')
            logger.error(e)
            logger.error(traceback.format_exc())
            #flask_restful.abort(500,message="AppUrl failed in some unexpected way")
            return apiabort(500,message="AppInstance failed in some unexpected way")

class CreateTunnel(Resource):
    @staticmethod
    def validate_connect_params(port, username, host, batchhost):
        try:
            intport = int(port)
        except Exception as e:
            return False
        if ' ' in username or '\n' in username: # This really needs more validation
            return False
        if ' ' in host or '\n' in host: # This really needs more validation
            return False
        return True

    def post(self,username,loginhost,batchhost):
        """
        Create a tunnel using established keys
        parameters for the tunnel (host username port etc)
        will be passed in the body
        """
        import logging
        import json
        logger = logging.getLogger()
        try:
            data = request.get_json()
            if type(data) is not dict:
                return apiabort(400, message="necessary data not provided in the post information")
            try:
                port = data['port']
            except KeyError as missingdata:
                raise AppParamsException("missing value for {} data contains {}".format(missingdata,data))
            if not CreateTunnel.validate_connect_params(port, username, loginhost, batchhost):
                raise AppParamsException("Invalid value: {} {} {} {}".
                                         format(username, loginhost, batchhost, port))
            if 'internalfirewall' in data:
                firewall = data['internalfirewall']
            else:
                firewall = False
            if 'localbind' in data:
                localbind = data['localbind']
            else:
                localbind = True
            sshsess = SSHSession.get_sshsession()
            if 'token' in data:
                authtok = data['token']
                logger.info('using token for authenticating tunnel')
            else:
                authtok = gen_authtok()
            port,pids = Ssh.tunnel(sshsess, port=port, batchhost=batchhost,
                       user=username, host=loginhost,
                       internalfirewall=firewall,
                       localbind=localbind, authtok=authtok)
            response = make_response(json.dumps({'localport':port}),200)
            response.mimetype = 'application/json'
            response.set_cookie('twsproxyauth', authtok, secure=True, samesite='None')
            return response
        except AppParamsException as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            
            return apiabort(400,message="Unable to create the secure tunnel. Please try again.")

        except Exception as e:
            import traceback
            logger.error(e)
            logger.error(traceback.format_exc())
            #flask_restful.abort(500,message="CreateTunnel failed in some unexpected way")
            return apiabort(500,message="CreateTunnel failed in some unexpected way")



api.add_resource(TunnelstatEP, '/tunnelstat/<string:authtok>')
api.add_resource(GetCert, '/getcert')

api.add_resource(JobStat, '/stat') # trying to deprecate
api.add_resource(JobSubmit, '/submit') # trying to deprecate
api.add_resource(AppInstance, '/appinstance/<string:username>/<string:loginhost>/<string:batchhost>/<string:jobid>') # trying to deprecate
api.add_resource(JobCancel, '/cancel/<string:jobid>') # trying to deprecate

api.add_resource(RemoteCommand,'/remotecommand')
api.add_resource(Ping, '/ping')
api.add_resource(SetTWSProxyAuth, '/settwsproxyauth/<string:authtok>')
api.add_resource(CreateTunnel, '/createtunnel/<string:username>/<string:loginhost>/<string:batchhost>')
api.add_resource(AppUrl, '/appurl')
api.add_resource(SSHAgent,'/sshagent')
api.add_resource(DirList,'/ls')
api.add_resource(MkDir,'/mkdir')
api.add_resource(ContactUs,'/contactus')