__init__.py

"""
This module persistently stores informion on tunnels in an in memory structure.
"""
import datetime
import yaml
import threading

class SshAgentException(Exception):
    pass

class SSHSession:
    """Interfaces for working with processes forked from flask
    in particular, we fork processes for ssh-agent and ssh tunnels and execution
    """
    def __init__(self,sshsessid,**kwargs):
        self.last = datetime.datetime.now()
        self.first = datetime.datetime.now()
        self.socket = None
        self.token = None
        self.port = {}
        self.pids = []
        self.authtok = None
        self.__dict__.update(kwargs)
        self.sshagent = 'ssh-agent'
        self.sshadd = '/usr/bin/ssh-add'
        self.sshkeygen = 'ssh-keygen'
        self.lock = threading.Lock()
        self.tunnels = []
        self.sshsessid = sshsessid

    def start_agent(self):
        import subprocess
        from .. import app
        import logging
        import os
        logger = logging.getLogger()
        if app.config['ENABLELAUNCH'] and 'SSH_AUTH_SOCK' in os.environ and os.environ['SSH_AUTH_SOCK']:
            self.socket = os.environ['SSH_AUTH_SOCK']
            return
        p = subprocess.Popen([self.sshagent],start_new_session=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
        (stdout,stderr) = p.communicate()
        for l in stdout.decode().split(';'):
            if 'SSH_AUTH_SOCK=' in l:
                socket = l.split('=')[1]
                self.socket = socket
            if 'SSH_AGENT_PID=' in l:
                pid = l.split('=')[1]
                self.pids.append(pid)

    def get_certs(self):
        import subprocess
        import os
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = self.socket
        cmd = [self.sshadd,'-L']
        p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE, env=env)
        (stdout,stderr) = p.communicate()
        return stdout

    def add_keycert(self,key,cert):
        import tempfile
        import os
        import subprocess
        import logging
        logger = logging.getLogger()
        self.lock.acquire()
        try:
            if self.socket is None:
                self.start_agent()
        except Exception as e:
            self.lock.release()
            raise e
        self.lock.release()
        keyf = tempfile.NamedTemporaryFile(mode='w',delete=False)
        keyname = keyf.name
        keyf.write(key)
        keyf.close()
        certf = open(keyname+'-cert.pub',mode='w')
        certf.write(cert)
        certf.close()
        p = subprocess.Popen([self.sshkeygen,'-L','-f','-'],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
        keygenout,keygenerr = p.communicate(cert.encode())
        logger.debug('inspecting keycert pid {}'.format(p.pid))
        # Examine the cert to determine its expiry. Use the -t flag to automatically remove from the ssh-agent when the cert expires
        certcontents = SSHSession.parse_cert_contents(keygenout.decode().splitlines())
        endtime = datetime.datetime.strptime(certcontents['Valid'][0].split()[3],"%Y-%m-%dT%H:%M:%S")
        delta = endtime - datetime.datetime.now() # I *think* the output of ssh-keygen -L is in the current timezone even though I assume the certs validity is in UTC
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = self.socket
        cmd = [self.sshadd,'-t',"{}".format(int(delta.total_seconds()))]
        cmd.append(keyname)
        if int(delta.total_seconds()) > 0:
            p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE,env=env)
            logger.debug('adding keycert pid {}'.format(p.pid))
            (stdout,stderr) = p.communicate()
            if p.returncode != 0:
                logger.error("Couldn't add key and cert")
                logger.error(stdout)
                logger.error(stderr)
                raise SshAgentException()
        os.unlink(keyname+'-cert.pub')
        os.unlink(keyname)

    def rm_fp(self,fp):
        cmd = [self.sshadd,'-L']
        p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE,env=env)
        (stdout,stderr) = p.communicate()
        for l in stdout.splitlines():
            p = subprocess.Popen([self.sshkeygen,'-l','-f','-'],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
            keygenout,keygenerr = p.communicate(l)
            if fp in keygenout:
                self.rm_pub(self,l)

    def rm_pub(self,pub):
        import tempfile
        pubf = tempfile.NamedTemporaryFile(mode='w',delete=False)
        pubname = pubf.name
        pubf.write(pub)
        pubf.close()
        cmd=[self.sshadd,'-d',pubname]
        p = subprocess.Popen(cmd,stdin=None,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
        (stderr,stdout) = p.communicate()
        os.unlink(pubname)


    def get_cert_contents(self):
        import os
        import subprocess
        import logging
        logger=logging.getLogger()
        res=[]
        if self.socket is None:
            return res
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = self.socket
        cmd = [self.sshadd,'-L']
        p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE,env=env)
        (stdout,stderr) = p.communicate()
        logger.debug('query agent pid {}'.format(p.pid))
        for l in stdout.splitlines():
            if b'cert' in l:
                p = subprocess.Popen([self.sshkeygen,'-L','-f','-'],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
                keygenout,keygenerr = p.communicate(l)
                logger.debug('inspecting agent contents pid {}'.format(p.pid))
                certcontents = SSHSession.parse_cert_contents(keygenout.decode().splitlines())
                res.append(certcontents)
        return res

    @staticmethod
    def parse_cert_contents(lines):
        key = None
        values = []
        res = {}
        for l in lines:
            l = l.rstrip().lstrip()
            if ':' in l:
                if key is not None:
                    res[key] = values
                values = []
                (key,v) = l.split(':',1)
                v = v.lstrip().rstrip()
                if v != '':
                    values = [v]
            else:
                if l != '':
                    values.append(l)
        return res



    def refresh(self):
        import datetime
        import logging
        self.last = datetime.datetime.now()
        logger = logging.getLogger()

    def addkey(self,key,cert):
        pass

    def kill(self):
        import os
        import signal
        import logging
        import time
        logger=logging.getLogger()
        logger.debug("shuting down ssh session for {} last seen at {}".format(self.authtok,self.last))
        for pid in self.pids:
            logger.debug("killing pid {}".format(pid))
            try:
                os.killpg(os.getpgid(int(pid)), signal.SIGTERM) # Sometimes this fails and I don't know why
                try:
                    os.kill(int(pid), 0) # If the first kill worked, this will raise a ProcessLookupError
                    time.sleep(2)
                    os.killpg(os.getpgid(int(pid)),signal.SIGKILL)
                    logger.error('resorting to sigkill for pid {}'.format(pid))
                except ProcessLookupError:
                    logger.debug(f'sigterm succesfull for {pid}')
                    pass
                logger.debug("killed {}".format(pid))
            except ProcessLookupError as e:
                logger.debug("process {} not found".format(pid))
        for tunnel in self.tunnels:
            tunnel.kill()
            (stdout, stderr) = tunnel.communicate()

    @staticmethod
    def test_sshsession(sess):
        import os
        import subprocess
        import logging
        logger=logging.getLogger()
        env = os.environ.copy()
        sess.lock.acquire()
        try:
            if sess.socket is None:
                sess.start_agent()
        except Exception as e:  
            sess.lock.release()
            raise e
        sess.lock.release()
        env['SSH_AUTH_SOCK'] = sess.socket
        cmd = [sess.sshadd,'-l']
        p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE,env=env)
        (stdout,stderr) = p.communicate()
        if p.returncode != 0:
            """
            A non-zero return code can occur if the agent is running
            but there are no keys loaded
            This is actually not an error condition
            """
            if b'The agent has no identities' in stdout:
                return
            logger.error("Couldn't communicate with the ssh agent")
            logger.error(stdout)
            logger.error(stderr)
            raise SshAgentException()


    @staticmethod
    def get_sshsession():
        import random
        import string
        from .. import sshsessions
        from flask import session, request
        import logging
        import uwsgi
        uwsgi.lock()
        
        logger = logging.getLogger()
        sshsessid = session.get('sshsessid', None)
        N = 8
        while sshsessid is None:
            key = ''.join(random.SystemRandom().choice(string.ascii_uppercase + string.digits) for _ in range(N))
            if key not in session:
                sshsessid = key
                session['sshsessid'] = sshsessid
        if sshsessid not in sshsessions:
            logger.debug('creating a new SSH session ID {}'.format(sshsessid))
            sshsessions[sshsessid] = SSHSession(sshsessid)

        sshsession = sshsessions[sshsessid]
        try:
            SSHSession.test_sshsession(sshsession)
        except SshAgentException:
            logger = logging.getLogger()
            logger.error('SSHAgentException, killing the agent and restarted. This is highly unexpected')
            sshsession.kill()
            sshsessions[sshsessid] = SSHSession(sshsessid)
            sshsession = sshsessions[sshsessid]
            SSHSession.test_sshsession(sshsession)

        sshsession.refresh()
        uwsgi.unlock()
        return sshsession

    @staticmethod
    def remove_sshsession():
        import random
        import string
        from .. import sshsessions
        from flask import session
        sshsessid = session.get('sshsessid', None)
        del sshsessions[sshsessid]