<reponame>movermeyer/poseidon
from __future__ import print_function
import os
import getpass
from cStringIO import StringIO

try:
    import paramiko
except ImportError:
    print("Please install paramiko to use SSH connection")
    raise

# make these optional so not everyone has to build C binaries
try:
    import pandas as pd
    if pd.__version__ <= '0.13.1':
        raise ImportError
    has_pandas = True
except ImportError:
    print("Pandas not found or out of date. SSHClient.ps will return str data")
    has_pandas = False


class SSHClient(object):
    """
    Thin wrapper to connect to client over SSH and execute commands
    """

    def __init__(self, host, username='root', password=<PASSWORD>, port=None,
                 interactive=False):
        """
        Parameters
        ----------
        interactive: bool, default False
            If True then prompts for password whenever necessary
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.interactive = interactive
        self.pwd = '~'
        self._con = None

    @property
    def con(self):
        if self._con is None:
            self._connect()
        return self._con

    def _connect(self):
        self._con = paramiko.SSHClient()
        self._con.set_missing_host_key_policy(
            paramiko.AutoAddPolicy())
        kwargs = {}
        for k in ['username', 'password', 'port']:
            if getattr(self, k, None):
                kwargs[k] = getattr(self, k)
        self._con.connect(self.host, **kwargs)

    def chdir(self, new_pwd, relative=True):
        """
        Parameters
        ----------
        new_pwd: str,
            Directory to change to
        relative: bool, default True
            If True then the given directory is treated as relative to the
            current directory
        """
        if new_pwd and self.pwd and relative:
            new_pwd = os.path.join(self.pwd, new_pwd)
        self.pwd = <PASSWORD>

    def add_public_key(self, key_path, validate_password=True):
        # TODO unit test. Not sure this works
        if validate_password:
            self.password = self.validate_password(self.password)
        key_path = os.path.expanduser(key_path)
        with open(key_path, 'r') as fp:
            cmd = 'mkdir -p ~/.ssh && echo "%s" >> ~/.ssh/authorized_keys'
            self.wait(cmd % fp.read())

    def close(self):
        if self._con is not None:
            self._con.close()
            self._con = None

    def exec_command(self, cmd):
        """
        Proceed with caution, if you run a command that causes a prompt and then
        try to read/print the stdout it's going to block forever

        Returns
        -------
        (stdin, stdout, stderr)
        """
        if self.pwd is not None:
            cmd = 'cd %s ; %s' % (self.pwd, cmd)
        if self.interactive:
            print(cmd)
        return self.con.exec_command(cmd)

    def wait(self, cmd, raise_on_error=True):
        """
        Execute command and wait for it to finish. Proceed with caution because
        if you run a command that causes a prompt this will hang
        """
        _, stdout, stderr = self.exec_command(cmd)
        stdout.channel.recv_exit_status()
        output = stdout.read()
        if self.interactive:
            print(output)
        errors = stderr.read()
        if self.interactive:
            print(errors)
        if errors and raise_on_error:
            raise ValueError(errors)
        return output

    def nohup(self, cmd):
        """
        Execute the command using nohup and &
        """
        cmd = "nohup %s &" % cmd
        self.exec_command(cmd)

    def sudo(self, password=None):
        """
        Enter sudo mode
        """
        if self.username == 'root':
            raise ValueError('Already root user')
        password = self.validate_password(password)
        stdin, stdout, stderr = self.exec_command('sudo su')
        stdin.write("%s\n" % password)
        stdin.flush()
        errors = stderr.read()
        if errors:
            raise ValueError(errors)

    def validate_password(self, password):
        if password is None:
            password = self.password
        if password is None and self.interactive:
            password = get<PASSWORD>pass()
        if password is None:
            raise ValueError("Password must not be empty")
        return password

    def unsudo(self):
        """
        Assume already in sudo
        """
        self.wait('exit')

    def apt(self, package_names, raise_on_error=False):
        """
        Install specified packages using apt-get. -y options are
        automatically used. Waits for command to finish.

        Parameters
        ----------
        package_names: list-like of str
        raise_on_error: bool, default False
            If True then raise ValueError if stderr is not empty
            debconf often gives tty error
        """
        if isinstance(package_names, basestring):
            package_names = [package_names]
        cmd = "apt-get install -y %s" % (' '.join(package_names))
        return self.wait(cmd, raise_on_error=raise_on_error)

    def curl(self, url, raise_on_error=True, **kwargs):
        import simplejson as json
        def format_param(name):
            if len(name) == 1:
                prefix = '-'
            else:
                prefix = '--'
            return prefix + name
        def format_value(value):
            if value is None:
                return ''
            return json.dumps(value)
        options = ['%s %s' % (format_param(k), format_value(v))
                   for k, v in kwargs.items()]
        cmd = 'curl %s "%s"' % (' '.join(options), url)
        return self.wait(cmd, raise_on_error=raise_on_error)

    def pip(self, package_names, raise_on_error=True):
        """
        Install specified python packages using pip. -U option added
        Waits for command to finish.

        Parameters
        ----------
        package_names: list-like of str
        raise_on_error: bool, default True
            If True then raise ValueError if stderr is not empty
        """
        if isinstance(package_names, basestring):
            package_names = [package_names]
        cmd = "pip install -U %s" % (' '.join(package_names))
        return self.wait(cmd, raise_on_error=raise_on_error)

    def pip_freeze(self, raise_on_error=True):
        """
        Run `pip freeze` and return output
        Waits for command to finish.
        """
        return self.wait('pip freeze', raise_on_error=raise_on_error)

    def pip_r(self, requirements, raise_on_error=True):
        """
        Install all requirements contained in the given file path
        Waits for command to finish.

        Parameters
        ----------
        requirements: str
            Path to requirements.txt
        raise_on_error: bool, default True
            If True then raise ValueError if stderr is not empty
        """
        cmd = "pip install -r %s" % requirements
        return self.wait(cmd, raise_on_error=raise_on_error)

    def ps(self, args=None, options='', all=True, verbose=True,
           as_frame='auto', raise_on_error=True):
        if args is None:
            args = ''
        if all:
            args += 'A'
        if verbose:
            args += 'f'
        if len(args) > 0 and args[0] != '-':
            args = '-' + args

        results = self.wait(('ps %s %s' % (args, options)).strip(),
                            raise_on_error=raise_on_error)

        if as_frame == 'auto':
            as_frame = has_pandas

        if as_frame:
            if not has_pandas:
                raise ImportError("Unable to import pandas")
            df = pd.read_fwf(StringIO(results))
            cmd_loc = df.columns.get_loc('CMD')
            if cmd_loc < len(df.columns):
                col = cmd_loc.fillna('')
                for i in range(cmd_loc + 1, len(df.columns)):
                    col = col + df.icol(i).fillna('')
                df['CMD'] = col
            return df

        return results

    def top(self):
        return self.ps('o', TOP_OPTIONS)

    def git(self, username, repo, alias=None, token=None):
        """
        Parameters
        ----------
        token: str, default None
            Assumes you have GITHUB_TOKEN in envvar if None

        https://github.com/blog/1270-easier-builds-and-deployments-using-git-
        over-https-and-oauth
        """
        if alias is None:
            alias = repo
        if token is None:
            token = os.environ.get('GITHUB_TOKEN')
        self.wait('mkdir -p %s' % alias)
        old_dir = self.pwd
        try:
            self.chdir(alias, relative=True)
            cmd = 'git init && git pull https://%s@github.com/%s/%s.git'
            # last line to stderr
            return self.wait(cmd % (token, username, repo),
                             raise_on_error=False)
        finally:
            self.chdir(old_dir, relative=False)

TOP_OPTIONS = '%cpu,%mem,user,comm'

from __future__ import print_function
import os
import getpass
from cStringIO import StringIO

try:
    import paramiko
except ImportError:
    print("Please install paramiko to use SSH connection")
    raise

# make these optional so not everyone has to build C binaries
try:
    import pandas as pd
    if pd.__version__ <= '0.13.1':
        raise ImportError
    has_pandas = True
except ImportError:
    print("Pandas not found or out of date. SSHClient.ps will return str data")
    has_pandas = False


class SSHClient(object):
    """
    Thin wrapper to connect to client over SSH and execute commands
    """

    def __init__(self, host, username='root', password=<PASSWORD>, port=None,
                 interactive=False):
        """
        Parameters
        ----------
        interactive: bool, default False
            If True then prompts for password whenever necessary
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.interactive = interactive
        self.pwd = '~'
        self._con = None

    @property
    def con(self):
        if self._con is None:
            self._connect()
        return self._con

    def _connect(self):
        self._con = paramiko.SSHClient()
        self._con.set_missing_host_key_policy(
            paramiko.AutoAddPolicy())
        kwargs = {}
        for k in ['username', 'password', 'port']:
            if getattr(self, k, None):
                kwargs[k] = getattr(self, k)
        self._con.connect(self.host, **kwargs)

    def chdir(self, new_pwd, relative=True):
        """
        Parameters
        ----------
        new_pwd: str,
            Directory to change to
        relative: bool, default True
            If True then the given directory is treated as relative to the
            current directory
        """
        if new_pwd and self.pwd and relative:
            new_pwd = os.path.join(self.pwd, new_pwd)
        self.pwd = <PASSWORD>

    def add_public_key(self, key_path, validate_password=True):
        # TODO unit test. Not sure this works
        if validate_password:
            self.password = self.validate_password(self.password)
        key_path = os.path.expanduser(key_path)
        with open(key_path, 'r') as fp:
            cmd = 'mkdir -p ~/.ssh && echo "%s" >> ~/.ssh/authorized_keys'
            self.wait(cmd % fp.read())

    def close(self):
        if self._con is not None:
            self._con.close()
            self._con = None

    def exec_command(self, cmd):
        """
        Proceed with caution, if you run a command that causes a prompt and then
        try to read/print the stdout it's going to block forever

        Returns
        -------
        (stdin, stdout, stderr)
        """
        if self.pwd is not None:
            cmd = 'cd %s ; %s' % (self.pwd, cmd)
        if self.interactive:
            print(cmd)
        return self.con.exec_command(cmd)

    def wait(self, cmd, raise_on_error=True):
        """
        Execute command and wait for it to finish. Proceed with caution because
        if you run a command that causes a prompt this will hang
        """
        _, stdout, stderr = self.exec_command(cmd)
        stdout.channel.recv_exit_status()
        output = stdout.read()
        if self.interactive:
            print(output)
        errors = stderr.read()
        if self.interactive:
            print(errors)
        if errors and raise_on_error:
            raise ValueError(errors)
        return output

    def nohup(self, cmd):
        """
        Execute the command using nohup and &
        """
        cmd = "nohup %s &" % cmd
        self.exec_command(cmd)

    def sudo(self, password=None):
        """
        Enter sudo mode
        """
        if self.username == 'root':
            raise ValueError('Already root user')
        password = self.validate_password(password)
        stdin, stdout, stderr = self.exec_command('sudo su')
        stdin.write("%s\n" % password)
        stdin.flush()
        errors = stderr.read()
        if errors:
            raise ValueError(errors)

    def validate_password(self, password):
        if password is None:
            password = self.password
        if password is None and self.interactive:
            password = get<PASSWORD>pass()
        if password is None:
            raise ValueError("Password must not be empty")
        return password

    def unsudo(self):
        """
        Assume already in sudo
        """
        self.wait('exit')

    def apt(self, package_names, raise_on_error=False):
        """
        Install specified packages using apt-get. -y options are
        automatically used. Waits for command to finish.

        Parameters
        ----------
        package_names: list-like of str
        raise_on_error: bool, default False
            If True then raise ValueError if stderr is not empty
            debconf often gives tty error
        """
        if isinstance(package_names, basestring):
            package_names = [package_names]
        cmd = "apt-get install -y %s" % (' '.join(package_names))
        return self.wait(cmd, raise_on_error=raise_on_error)

    def curl(self, url, raise_on_error=True, **kwargs):
        import simplejson as json
        def format_param(name):
            if len(name) == 1:
                prefix = '-'
            else:
                prefix = '--'
            return prefix + name
        def format_value(value):
            if value is None:
                return ''
            return json.dumps(value)
        options = ['%s %s' % (format_param(k), format_value(v))
                   for k, v in kwargs.items()]
        cmd = 'curl %s "%s"' % (' '.join(options), url)
        return self.wait(cmd, raise_on_error=raise_on_error)

    def pip(self, package_names, raise_on_error=True):
        """
        Install specified python packages using pip. -U option added
        Waits for command to finish.

        Parameters
        ----------
        package_names: list-like of str
        raise_on_error: bool, default True
            If True then raise ValueError if stderr is not empty
        """
        if isinstance(package_names, basestring):
            package_names = [package_names]
        cmd = "pip install -U %s" % (' '.join(package_names))
        return self.wait(cmd, raise_on_error=raise_on_error)

    def pip_freeze(self, raise_on_error=True):
        """
        Run `pip freeze` and return output
        Waits for command to finish.
        """
        return self.wait('pip freeze', raise_on_error=raise_on_error)

    def pip_r(self, requirements, raise_on_error=True):
        """
        Install all requirements contained in the given file path
        Waits for command to finish.

        Parameters
        ----------
        requirements: str
            Path to requirements.txt
        raise_on_error: bool, default True
            If True then raise ValueError if stderr is not empty
        """
        cmd = "pip install -r %s" % requirements
        return self.wait(cmd, raise_on_error=raise_on_error)

    def ps(self, args=None, options='', all=True, verbose=True,
           as_frame='auto', raise_on_error=True):
        if args is None:
            args = ''
        if all:
            args += 'A'
        if verbose:
            args += 'f'
        if len(args) > 0 and args[0] != '-':
            args = '-' + args

        results = self.wait(('ps %s %s' % (args, options)).strip(),
                            raise_on_error=raise_on_error)

        if as_frame == 'auto':
            as_frame = has_pandas

        if as_frame:
            if not has_pandas:
                raise ImportError("Unable to import pandas")
            df = pd.read_fwf(StringIO(results))
            cmd_loc = df.columns.get_loc('CMD')
            if cmd_loc < len(df.columns):
                col = cmd_loc.fillna('')
                for i in range(cmd_loc + 1, len(df.columns)):
                    col = col + df.icol(i).fillna('')
                df['CMD'] = col
            return df

        return results

    def top(self):
        return self.ps('o', TOP_OPTIONS)

    def git(self, username, repo, alias=None, token=None):
        """
        Parameters
        ----------
        token: str, default None
            Assumes you have GITHUB_TOKEN in envvar if None

        https://github.com/blog/1270-easier-builds-and-deployments-using-git-
        over-https-and-oauth
        """
        if alias is None:
            alias = repo
        if token is None:
            token = os.environ.get('GITHUB_TOKEN')
        self.wait('mkdir -p %s' % alias)
        old_dir = self.pwd
        try:
            self.chdir(alias, relative=True)
            cmd = 'git init && git pull https://%s@github.com/%s/%s.git'
            # last line to stderr
            return self.wait(cmd % (token, username, repo),
                             raise_on_error=False)
        finally:
            self.chdir(old_dir, relative=False)

TOP_OPTIONS = '%cpu,%mem,user,comm'
# make these optional so not everyone has to build C binaries Thin wrapper to connect to client over SSH and execute commands Parameters         ----------         interactive: bool, default False             If True then prompts for password whenever necessary Parameters         ----------         new_pwd: str,             Directory to change to         relative: bool, default True             If True then the given directory is treated as relative to the             current directory # TODO unit test. Not sure this works Proceed with caution, if you run a command that causes a prompt and then         try to read/print the stdout it's going to block forever          Returns         -------         (stdin, stdout, stderr) Execute command and wait for it to finish. Proceed with caution because         if you run a command that causes a prompt this will hang Execute the command using nohup and & Enter sudo mode Assume already in sudo Install specified packages using apt-get. -y options are         automatically used. Waits for command to finish.          Parameters         ----------         package_names: list-like of str         raise_on_error: bool, default False             If True then raise ValueError if stderr is not empty             debconf often gives tty error Install specified python packages using pip. -U option added         Waits for command to finish.          Parameters         ----------         package_names: list-like of str         raise_on_error: bool, default True             If True then raise ValueError if stderr is not empty Run `pip freeze` and return output         Waits for command to finish. Install all requirements contained in the given file path         Waits for command to finish.          Parameters         ----------         requirements: str             Path to requirements.txt         raise_on_error: bool, default True             If True then raise ValueError if stderr is not empty Parameters         ----------         token: str, default None             Assumes you have GITHUB_TOKEN in envvar if None          https://github.com/blog/1270-easier-builds-and-deployments-using-git-         over-https-and-oauth # last line to stderr