__init__.py

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

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,**kwargs):
        self.last = datetime.datetime.now()
        self.socket = None
        self.token = None
        self.port = {}
        self.key = ''
        self.cert = ''
        self.pids = []
        self.authtok = None
        self.__dict__.update(kwargs)
        self.sshagent = 'ssh-agent'
        self.sshadd = '/usr/bin/ssh-add'
        self.sshkeygen = 'ssh-keygen'
        self.ctrl_processes = {}

    def start_agent(self):
        import subprocess
        p = subprocess.Popen([self.sshagent],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 add_keycert(self,key,cert):
        import tempfile
        import os
        import subprocess
        import logging
        logger = logging.getLogger()
        if self.socket is None:
            logger.debug('adding a cert, start the agent first')
            self.start_agent()
            logger.debug('agent socket is now {}'.format(self.socket))
        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()
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = self.socket
        cmd = [self.sshadd]
        cmd.append(keyname)
        p = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE,env=env)
        (stdout,stderr) = p.communicate()
        if p.returncode != 0:
            logger.error("Couldn't add key and cert")
            logger.debug(stdout)
            logger.debug(stderr)
        os.unlink(keyname+'-cert.pub')
        os.unlink(keyname)

    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()
        if stderr is not None:
            logger.debug('called sshadd, got stderr {}'.format(stderr))
        if stdout is not None:
            logger.debug('called sshadd, got stdout {}'.format(stdout))
        for l in stdout.splitlines():
            logger.debug('is {} a cert?'.format(l))
            if b'cert' in l:
                logger.debug('decoding {}'.format(l))
                p = subprocess.Popen([self.sshkeygen,'-L','-f','-'],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
                keygenout,keygenerr = p.communicate(l)
                certcontents = SSHSession.parse_cert_contents(keygenout.decode().splitlines())
                res.append(certcontents)
            else: 
                res.append({'pubkey':l.decode()})
        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 is not '':
                    values = [v]
            else:
                if l is not '':
                    values.append(l)
        return res



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

    # def set_authtok(self,authtok):
    #     self.authtok = authtok

    def addkey(self,key,cert):
        pass

    def kill(self):
        import os
        import signal
        import logging
        logger=logging.getLogger()
        logger.debug("shuting down ssh session for {} last seen at {}".format(self.authtok,self.last))
        for pid in self.pids:
            try:
                os.kill(int(pid), signal.SIGTERM)
            except ProcessLookupError as e:
                logger.debug("process {} not found".format(pid))
        for ctrl in self.ctrl_processes.items():
            ctrl[1].kill()
            ctrl[1].wait(5)
            os.unlink(ctrl[0])


    @staticmethod
    def get_sshsession():
        import random
        import string
        from .. import sshsessions
        from flask import session
        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:
            sshsessions[sshsessid] = SSHSession()

        return sshsessions[sshsessid]

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