__init__.py

"""
This module handles SSH Connections
"""
import subprocess

class SshAgentException(Exception):
    pass
class SshValidationException(Exception):
    pass

class Ssh(object):
    """
    Ssh class can execute or create tunnelstat
    """

    @staticmethod
    def validate_port(port):
        try:
            port=int(port)
        except:
            raise SshValidationException("port number {} was not an integer".format(port))

    @staticmethod
    def validate_username(user):
        import re
        username_re = re.compile(r'^[a-z_]([a-z0-9_-]{0,31}|[a-z0-9_\.-]{0,30}\$)$')
        if not username_re.match(user):
            raise SshValidationException("username {} is not valid according to posix standards".format(user))

    @staticmethod
    def validate_hostname(host):
        import re
        hostname_re= re.compile(r'^[a-z0-9_\.-]+$')
        if not hostname_re.match(host):
            raise SshValidationException("hostname {} is not valid".format(host))

    @staticmethod
    def validate_command(cmd):
        """So in theory, due the the format of the Popen command
        we can take any input as a commnand to execute on the login host, including
        ; and other forms of injection,
        and the worst that can happen is the user screws up their own account
        """
        pass


    @staticmethod
    def execute(sess, host, user, cmd, stdin=None):
        """
        execute the command cmd on the host via ssh
        # assume the environment is already setup with an
        # SSH_AUTH_SOCK that allows login
        """
        import os
        if cmd is None and stdin is None:
            return {'stdout': b'', 'stderr': b'No command given to execute'}
        env = os.environ.copy()
        if sess.socket is None:
            raise SshAgentException("No ssh-agent yet")
        env['SSH_AUTH_SOCK'] = sess.socket
        Ssh.validate_username(user)
        Ssh.validate_hostname(host)
        Ssh.validate_command(cmd)
        exec_p = subprocess.Popen(['ssh', '-A', '-o', 'Stricthostkeychecking=no', '-l',
                                   user, host, cmd],
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  stdin=subprocess.PIPE, env=env)
        if stdin is not None:
            (stdout, stderr) = exec_p.communicate(stdin.encode())
        else:
            (stdout, stderr) = exec_p.communicate()
        return {'stdout':stdout, 'stderr':stderr}

    @staticmethod
    def tunnel(sess, port, batchhost, user, host, internalfirewall = True, localbind = True, authtok = None):
        """
        the double tunnel is used if the server we run on the batch host is only
        addressiable on localhost
        e.g. jupyter in its default config runs on localhost:8888 so a web
        browser on the login node can not connect to it
        """
        if port == 22:
            return Ssh.singletunnel(sess, port, batchhost, user, host, authtok = authtok)
        if not internalfirewall and not localbind:
            return Ssh.singletunnel(sess, port, batchhost, user, host, authtok = authtok)
        return Ssh.doubletunnel(sess, port, batchhost, user, host, authtok = authtok)

    @staticmethod
    def addkey(sess,key,cert):
        pass

    @staticmethod
    def singletunnel(sess, port, batchhost, user, host, authtok):
        """
        # fork a daemon process to hold a tunnel open on a given port like
        # ssh -l user -L tunnel host -N
        """
        import os
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = sess.socket
        Ssh.validate_port(port)
        Ssh.validate_hostname(batchhost)
        Ssh.validate_username(user)
        Ssh.validate_hostname(host)
        localport = Ssh.get_free_port()
        tunnel_p = subprocess.Popen(['ssh', '-N', '-o', 'Stricthostkeychecking=no', '-L',
                                     '{localport}:{batchhost}:{port}'.format(port=port,
                                                                             localport=localport,
                                                                             batchhost=batchhost),
                                     '-l', user, host], stdin=subprocess.PIPE, env=env)
        Ssh.wait_for_tunnel(localport)
        if sess.port is None:
            sess.port = {}

        sess.port.update({authtok,localport})
        sess.pids.append(tunnel_p.pid)
        return localport, [tunnel_p.pid]

    @staticmethod
    def doubletunnel(sess, port, batchhost, user, host, authtok):
        """
        # fork a daemon process to hold a tunnel open on a given port like
        # ssh -l user -L tunnel host -N
        """
        import os
        env = os.environ.copy()
        env['SSH_AUTH_SOCK'] = sess.socket
        pids = []
        localport1 = Ssh.get_free_port()
        Ssh.validate_port(port)
        Ssh.validate_hostname(batchhost)
        Ssh.validate_username(user)
        Ssh.validate_hostname(host)
        tunnel_p = subprocess.Popen(['ssh', '-N', '-o', 'Stricthostkeychecking=no', '-L',
                                     '{localport1}:{batchhost}:22'.format(localport1=localport1,
                                                                          batchhost=batchhost),
                                     '-l', user, host], stdin=subprocess.PIPE, env=env)
        pids.append(tunnel_p.pid)
        Ssh.wait_for_tunnel(localport1)
        localport2 = Ssh.get_free_port()
        tunnel_p = subprocess.Popen(['ssh', '-N', '-o', 'Stricthostkeychecking=no', '-L',
                                     '{localport2}:localhost:{port}'.format(port=port,
                                                                            localport2=localport2),
                                     '-l', user, 'localhost', '-p',
                                     '{}'.format(localport1)], stdin=subprocess.PIPE, stdout = subprocess.PIPE, stderr = subprocess.PIPE, env=env)
        pids.append(tunnel_p.pid)
        Ssh.wait_for_tunnel(localport2)
        if sess.port is None:
            sess.port = {}
        sess.port.update({authtok:localport2})
        pv = tunnel_p.poll()
        sess.pids.extend(pids)
        return localport2, pids

    @staticmethod
    def wait_for_tunnel(localport):
        """
        # In order to avoid a race condition, we wait for the tunnel to be established
        """
        import socket
        notopen = True
        while notopen:
            ssock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            ssock.setblocking(True)
            try:
                ssock.connect(('127.0.0.1', localport))
                notopen = False
                ssock.close()
            except socket.error:
                ssock.close()
        return

    @staticmethod
    def get_free_port():
        """
        # Finds a port which the local server can listen on.
        """
        import socket
        serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        for testport in range(1025, 65500):
            try:
                serversocket.bind(('127.0.0.1', testport))
                port = testport
                serversocket.close()
                return port
            except OSError:
                pass