Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow keyboard interrupts #30

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 33 additions & 31 deletions psshlib/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,8 @@ def common_parser():
'spaces, quotes, and backslashes')
connection_group.add_option('-X', '--extra-arg', dest='extra', action='append',
metavar='ARG', help='Extra command-line argument')
connection_group.add_option('-k', '--allow-keyboard-interrupts', dest='allow_keyboard_interrupts',
action='store_true', help='Keyboard interrupts will terminate the process.')

map(parser.add_option_group, [connection_group, output_group, filter_group])
parser.group_map = { # used so subparsers can easily find option groups
Expand Down Expand Up @@ -172,7 +174,7 @@ def run(self, hosts=[], args=None, opts=None):
args = args or self.args
opts = opts or self.opts
hosts = hosts or ServerPool(opts)

if args is None:
raise Exception
elif not hosts:
Expand All @@ -181,7 +183,7 @@ def run(self, hosts=[], args=None, opts=None):
self.setup(opts)

manager = self.setup_manager(hosts, args, opts)

psshutil.run_manager(manager)

exitcode = self.teardown_manager(manager)
Expand Down Expand Up @@ -229,19 +231,19 @@ def pssh_option_parser():
help='when used with the --script option, will do two things differently: '
'1) transfer the script to /root instead of /tmp, 2) run the script '
'as root, not the login user')
pssh_group.add_option('--args', dest='script_args',
pssh_group.add_option('--args', dest='script_args',
help='companion option for --script. Passes SCRIPT_ARGS as arguments'
' to the script run on the remote host.')
pssh_group.add_option('--env', dest='env', action='append', metavar='SCRIPT_ENV',
default=[],
help='specify key=value pairs to inject into the environment of a running '
'--script, e.g. --env="FOO=BAR". Can be specified multiple times')
pssh_group.add_option('--runtime',
pssh_group.add_option('--runtime',
help='specify the runtime to use when running the script from --script')
pssh_group.add_option('--copy-to', default='/tmp',
help='where to remotely copy scripts passed via --script (defaults to '
'/root if --sudo is passed, otherwise /tmp)')

parser.add_option_group(pssh_group)
parser.group_map['pssh_group'] = pssh_group

Expand All @@ -254,10 +256,10 @@ def parse_args(self):
defaults = common_defaults(timeout=_DEFAULT_TIMEOUT)
parser.set_defaults(**defaults)
opts, args = parser.parse_args()

if len(args) == 0 and not opts.send_input and not opts.script:
parser.error('Command not specified.')

if not opts.host_files and not opts.host_strings:
parser.error('Hosts not specified.')

Expand All @@ -266,7 +268,7 @@ def parse_args(self):

if opts.copy_to and not opts.copy_to.startswith('/'):
parser.error('Remote script directory must be a path')

return opts, args

def setup(self, opts):
Expand Down Expand Up @@ -321,7 +323,7 @@ def _generate_script_envelope(self):
else:
envelope = (
"cat > %(script)s; CATRET=$?; chmod 700 %(script)s; %(environ)s %(runner)s %(script_args)s; RET=$((CATRET+$?));"
"rm -f %(script)s; exit $RET"
"rm -f %(script)s; exit $RET"
)
return envelope % {
'script': script,
Expand Down Expand Up @@ -365,7 +367,7 @@ def setup_manager(self, hosts, args, opts):
cmd.append(cmdline)
t = SshTask(host, port, user, cmd, cmdline, opts, stdin)
manager.add_task(t)

return manager

def teardown_manager(self, manager):
Expand Down Expand Up @@ -404,16 +406,16 @@ def parse_args(self):
defaults = common_defaults()
parser.set_defaults(**defaults)
opts, args = parser.parse_args()

if len(args) < 1:
parser.error('Paths not specified.')

if len(args) < 2:
parser.error('Remote path not specified.')

if not opts.host_files and not opts.host_strings:
parser.error('Hosts not specified.')

return opts, args

def setup(self, opts):
Expand Down Expand Up @@ -473,16 +475,16 @@ def parse_args(self):
defaults = common_defaults(timeout=_DEFAULT_TIMEOUT)
parser.set_defaults(**defaults)
opts, args = parser.parse_args()

if len(args) < 1:
parser.error('Pattern not specified.')

if len(args) > 1:
parser.error('Extra arguments given after the pattern.')

if not opts.host_files and not opts.host_strings:
parser.error('Hosts not specified.')

return opts, args

def setup(self, opts):
Expand Down Expand Up @@ -547,19 +549,19 @@ def parse_args(self):
defaults = common_defaults()
parser.set_defaults(**defaults)
opts, args = parser.parse_args()

if len(args) < 1:
parser.error('Paths not specified.')

if len(args) < 2:
parser.error('Remote path not specified.')

if len(args) > 2:
parser.error('Extra arguments given after the remote path.')

if not opts.host_files and not opts.host_strings:
parser.error('Hosts not specified.')

return opts, args

def setup(self, opts):
Expand All @@ -585,7 +587,7 @@ def setup_manager(self, hosts, args, opts):
ssh += ['-p', port]
if opts.ssh_args:
ssh += [opts.ssh_args]

cmd = ['rsync', '-e', ' '.join(ssh)]
if opts.verbose:
cmd.append('-v')
Expand Down Expand Up @@ -629,7 +631,7 @@ def pslurp_option_parser():
action='store_true', help='recusively copy directories (OPTIONAL)')
pslurp_group.add_option('-L', '--localdir', dest='localdir',
help='output directory for remote file copies')

parser.add_option_group(pslurp_group)
parser.group_map['pslurp_group'] = pslurp_group

Expand All @@ -641,19 +643,19 @@ def parse_args(self):
defaults = common_defaults()
parser.set_defaults(**defaults)
opts, args = parser.parse_args()

if len(args) < 1:
parser.error('Paths not specified.')

if len(args) < 2:
parser.error('Local path not specified.')

if len(args) > 2:
parser.error('Extra arguments given after the local path.')

if not opts.host_files and not opts.host_strings:
parser.error('Hosts not specified.')

return opts, args

def setup(self, opts):
Expand Down Expand Up @@ -712,7 +714,7 @@ def teardown_manager(self, manager):
return 3
for status in statuses:
if status == 255:
return 4
return 4
for status in statuses:
if status != 0:
return 5
Expand Down
20 changes: 12 additions & 8 deletions psshlib/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ def __init__(self, opts):

self.progress_bar = opts.progress_bar
self.test_cases = opts.test_cases

def _setup_progress_bar(self):
""" This should be called after ``self.tasks`` is populated
"""
Expand All @@ -61,7 +61,7 @@ def _setup_progress_bar(self):

def _split_manager(self):
# set up the test manager and add first n tasks
new_opts = deepcopy(self.opts)
new_opts = deepcopy(self.opts)
new_opts.__dict__['test_cases'] = None # remove test_cases option, or there'll be a recursion error
new_opts.__dict__['summary'] = None # don't print summary now, do it later
test_man = self.__class__(new_opts)
Expand All @@ -82,7 +82,7 @@ def _split_manager(self):
# add remaining tasks
map(finish_man.add_task, self.tasks[slice(self.test_cases, None)])
psshutil.run_manager(finish_man)

return test_man, finish_man

def run(self):
Expand All @@ -92,7 +92,7 @@ def run(self):
man1, man2 = self._split_manager()
self.done = man1.done + man2.done
else:
self._run()
self._run()

self.tally_results()

Expand All @@ -102,15 +102,15 @@ def run(self):
if self.opts.fork_hosts:
failed_file = open(self.opts.fork_hosts + '.failed.lst', 'w')
passed_file = open(self.opts.fork_hosts + '.passed.lst', 'w')

for i in self.ssh_failed + self.killed + self.cmd_failed:
failed_file.write(i.host + '\n')

for i in self.succeeded:
passed_file.write(i.host + '\n')

return [task.exitstatus for task in self.done]

def _run(self):
try:
if self.outdir or self.errdir:
Expand All @@ -136,11 +136,15 @@ def _run(self):
# This exception handler tries to clean things up and prints
# out a nice status message for each interrupted host.
self.interrupted()
if self.opts.allow_keyboard_interrupts:
raise


except KeyboardInterrupt:
# This exception handler doesn't print out any fancy status
# information--it just stops.
pass
if self.opts.allow_keyboard_interrupts:
raise

if writer:
writer.signal_quit()
Expand Down