Skip to content
Snippets Groups Projects
__init__.py 10.2 KiB
Newer Older
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
        self.last = datetime.datetime.now()
Chris Hines's avatar
Chris Hines committed
        self.first = datetime.datetime.now()
        self.socket = None
        self.token = None
        self.pids = []
        self.authtok = None
        self.__dict__.update(kwargs)
        self.sshadd = '/usr/bin/ssh-add'
        self.lock = threading.Lock()
Chris Hines's avatar
Chris Hines committed
        self.tunnels = []
    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)
Chris Hines's avatar
Chris Hines committed
        import subprocess
Chris Hines's avatar
Chris Hines committed
        import os
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = self.socket
        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)
Chris Hines's avatar
Chris Hines committed
        keygenout,keygenerr = p.communicate(cert.encode())
Chris Hines's avatar
Chris Hines committed
        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
Chris Hines's avatar
Chris Hines committed
        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)


        import os
        import subprocess
        import logging
        logger=logging.getLogger()
        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()
Chris Hines's avatar
Chris Hines committed
        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)
Chris Hines's avatar
Chris Hines committed
                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:
            if ':' in l:
                if key is not None:
                    res[key] = values
                values = []
                (key,v) = l.split(':',1)
Chris Hines's avatar
Chris Hines committed
                if v != '':
Chris Hines's avatar
Chris Hines committed
                if l != '':
                    values.append(l)
    def refresh(self):
        import datetime
Chris Hines's avatar
Chris Hines committed
        import logging
        self.last = datetime.datetime.now()
Chris Hines's avatar
Chris Hines committed
        logger = logging.getLogger()
    def addkey(self,key,cert):
        pass

    def kill(self):
        import os
        import signal
        import logging
Chris Hines's avatar
Chris Hines committed
        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))
                os.killpg(os.getpgid(int(pid)), signal.SIGTERM) # Sometimes this fails and I don't know why
Chris Hines's avatar
Chris Hines committed
                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)
Chris Hines's avatar
Chris Hines committed
                    logger.error('resorting to sigkill for pid {}'.format(pid))
                except ProcessLookupError:
                    logger.debug(f'sigterm succesfull for {pid}')
Chris Hines's avatar
Chris Hines committed
                    pass
                logger.debug("killed {}".format(pid))
            except ProcessLookupError as e:
Ubuntu's avatar
Ubuntu committed
                logger.debug("process {} not found".format(pid))
Chris Hines's avatar
Chris Hines committed
        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()

    def get_sshsession():
        import random
        import string
        from .. import sshsessions
Chris Hines's avatar
Chris Hines committed
        from flask import session, request
Chris Hines's avatar
Chris Hines committed
        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:
Chris Hines's avatar
Chris Hines committed
            logger.debug('creating a new SSH session ID {}'.format(sshsessid))
            sshsessions[sshsessid] = SSHSession(sshsessid)
Chris Hines's avatar
Chris Hines committed
        sshsession = sshsessions[sshsessid]
Chris Hines's avatar
Chris Hines committed
            SSHSession.test_sshsession(sshsession)
            logger = logging.getLogger()
            logger.error('SSHAgentException, killing the agent and restarted. This is highly unexpected')
Chris Hines's avatar
Chris Hines committed
            sshsession.kill()
            sshsessions[sshsessid] = SSHSession(sshsessid)
Chris Hines's avatar
Chris Hines committed
            sshsession = sshsessions[sshsessid]
            SSHSession.test_sshsession(sshsession)
Chris Hines's avatar
Chris Hines committed
        sshsession.refresh()
Chris Hines's avatar
Chris Hines committed
        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]