diff --git a/sshmux/__init__.py b/sshmux/__init__.py index 792d600..e69de29 100644 --- a/sshmux/__init__.py +++ b/sshmux/__init__.py @@ -1 +0,0 @@ -# diff --git a/sshmux/main.py b/sshmux/main.py index 0ec0f48..92ea036 100644 --- a/sshmux/main.py +++ b/sshmux/main.py @@ -5,7 +5,6 @@ import multiprocessing from sshmux import validate from sshmux.ssh import ssh -from getpass import getpass from os import environ @@ -14,16 +13,11 @@ multiple=True, help='IP address or hostname') @click.option('--username', '-u', callback=validate.validate_user, default='', help='ssh username') -@click.option('--password', '-p', default=False, help='ssh password') @click.option('--key', '-k', default=environ['HOME'] + '/.ssh/id_rsa', help='ssh private key') -def main(hostname, username, password, key): +def main(hostname, username, key): """Open ssh session with each ip and execute a command from stdin.""" - if not password: - key = validate.validate_key(key) - if password: - password = getpass() - password = validate.validate_pass(password) + key = validate.validate_key(key) print("Enter your commands below:\n") command = input("sshmux > ") @@ -31,7 +25,7 @@ def main(hostname, username, password, key): procs = [] for server in hostname: procs.append(multiprocessing.Process( - target=ssh, args=(server, command, username, password, key))) + target=ssh, args=(server, command, username, key))) for proc in procs: proc.start() for proc in procs: diff --git a/sshmux/ssh.py b/sshmux/ssh.py index a0f5c4f..9017174 100644 --- a/sshmux/ssh.py +++ b/sshmux/ssh.py @@ -1,56 +1,38 @@ -from __future__ import print_function - -import pexpect -import tempfile -from os import unlink -from sshmux.errors import MuxError - - -def print_output(server, output): - """parse ssh output and print to stdout""" - print(server + ":\n") - for line in output.split('\n')[1:]: - print(line) - - -def ssh(host, cmd, user, password, key, timeout=10, bg_run=False): - """connect to host via ssh""" - output_file = tempfile.NamedTemporaryFile(delete=False) - option = ["-q", "-oStrictHostKeyChecking=no", - "-oUserKnownHostsFile=/dev/null"] - if password: - option.append("-oPubkeyAuthentication=no") - if not password: - option.append("-o PreferredAuthentications=publickey") - if bg_run: - option.append('-f') - - options = " ".join(option) - ssh_cmd = None - if not password: - ssh_cmd = 'ssh -i {0} {1}@{2} {3} "{4}"'.format( - key, user, host, options, cmd) - elif password: - ssh_cmd = 'ssh {0}@{1} {2} "{3}"'.format(user, host, options, cmd) - - child = pexpect.spawn(ssh_cmd, timeout=timeout) - if password: - child.expect(['Password for']) - child.sendline(password) - - child.logfile = output_file - child.expect(pexpect.EOF) - child.close() - output_file.close() - - # BUG(rjrhaverkamp): U mode is deprecated - read_file = open(output_file.name, 'rU') - stdout = read_file.read() - output_file.close() - read_file.close() - unlink(output_file.name) - - if child.exitstatus != 0: - raise MuxError(stdout) - print_output(host, stdout) - return stdout +from __future__ import print_function + +from subprocess import Popen, PIPE, STDOUT +from sshmux.errors import MuxError + + +def print_output(server, output): + """parse ssh output and print to stdout""" + print(server + ":\n") + for line in output.split('\n')[1:]: + print(line) + + +def ssh(host, cmd, user, key, bg_run=False): + """connect to host via ssh""" + option = ["-q", "-oStrictHostKeyChecking=no", + "-oUserKnownHostsFile=/dev/null", + "-o PreferredAuthentications=publickey"] + if bg_run: + option.append('-f') + options = " ".join(option) + ssh_cmd = None + ssh_cmd = 'ssh -i {0} {1}@{2} {3} "{4}"'.format( + key, user, host, options, cmd) + + run = Popen(ssh_cmd, stdout=PIPE, stderr=STDOUT, shell=True) + run.wait() + + if run.returncode == 127: + raise MuxError("command {0} does not exist.".format(cmd)) + if run.returncode != 0: + raise MuxError("failed to run {0} on {1}. Exited with: {2}".format( + cmd, host, run.returncode)) + + output, _ = run.communicate() + stdout = output.decode("utf-8") + print_output(host, stdout) + return stdout diff --git a/sshmux/validate.py b/sshmux/validate.py index e291ede..6bcf308 100644 --- a/sshmux/validate.py +++ b/sshmux/validate.py @@ -19,13 +19,6 @@ def validate_hostname(ctx, param, hostname): return hostname -def validate_pass(password): - """validate password lenght""" - if len(password) == 0 or len(password) > 100: - raise MuxError('password length is not valid') - return password - - def validate_user(ctx, param, username): """validate username length""" if len(username) == 0 or len(username) > 100: diff --git a/tests/test_ssh.py b/tests/test_ssh.py index fdb57ec..597111f 100644 --- a/tests/test_ssh.py +++ b/tests/test_ssh.py @@ -8,7 +8,7 @@ class TestSSH(unittest.TestCase): def test_ssh_output(self): output = ssh.ssh(environ['sshmux_test_host'], 'echo "hello"', environ[ - 'sshmux_test_user'], '', environ['sshmux_test_key']) + 'sshmux_test_user'], environ['sshmux_test_key']) self.assertEqual(output, 'hello\n') def test_wrong_cmd(self): @@ -16,7 +16,12 @@ def test_wrong_cmd(self): user = environ['sshmux_test_user'] key = environ['sshmux_test_key'] self.assertRaises(errors.MuxError, ssh.ssh, host, - 'does_not_exist', user, '', key) + 'does_not_exist', user, key) + + def test_background_run(self): + output = ssh.ssh(environ['sshmux_test_host'], 'echo "hello"&', environ['sshmux_test_user'], environ['sshmux_test_key'], bg_run=True) # NOQA + self.assertEqual(output, 'hello\n') + if __name__ == '__main__': unittest.main() diff --git a/tests/test_validations.py b/tests/test_validations.py index 4a72afa..07329ab 100644 --- a/tests/test_validations.py +++ b/tests/test_validations.py @@ -3,7 +3,6 @@ from sshmux import validate from sshmux.errors import MuxError from os import environ -import sys from click.testing import CliRunner @@ -11,14 +10,14 @@ class TestValidations(unittest.TestCase): @click.command() - @click.option('--hostname', '-h', callback=validate.validate_hostname, multiple=True, - help='IP address or hostname') + @click.option('--hostname', '-h', callback=validate.validate_hostname, + multiple=True, help='IP address or hostname') def check_hostname(hostname): click.echo('sucess') @click.command() - @click.option('--username', '-u', callback=validate.validate_user, default='', - help='ssh username') + @click.option('--username', '-u', callback=validate.validate_user, + default='', help='ssh username') def check_username(username): click.echo('sucess') @@ -61,14 +60,6 @@ def test_key_fail(self): key = environ['HOME'] + '/.ssh/id_rsa_that_does_not_exist' self.assertRaises(MuxError, validate.validate_key, key) - def test_password_check(self): - password = "testpassword" - valid_password = validate.validate_pass(password) - self.assertEqual(password, valid_password) - - def test_password_fail(self): - self.assertRaises(MuxError, - validate.validate_pass, "testpassword" * 12) if __name__ == '__main__': unittest.main()