Skip to content

Commit

Permalink
Merge pull request #53 from rolobio/multi-target-groups
Browse files Browse the repository at this point in the history 
Multi target groups
  • Loading branch information
@rolobio
rolobio committed Feb 18, 2015
2 parents 9815a58 + 623c12a commit 114bb38
Show file tree
Hide file tree
Showing 4 changed files with 55 additions and 22 deletions.
31 changes: 20 additions & 11 deletions sshm/lib.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,19 +16,24 @@
def expand_ranges(to_expand):
"""
Convert a comma-seperated range of integers into a list. Keep any zero
padding the numbers may have.
padding the numbers may have. If the provided string is just a -, then I
will return an entire range.
Example: "1,4,07-10" to ['1', '4', '07', '08', '09', '10']
@param to_expand: Expand this string into a list of integers.
@type to_expand: str
"""
if to_expand == '-':
return [str(i) for i in range(0, 256)]

nums = []
for single, range_str in _match_ranges.findall(to_expand):
if single:
nums.append(single)
if range_str:
i, j = range_str.split('-')
print(i, j)
# Create a string that will pad the integer with its current amount
# of zeroes.
# Example: if i is '03' the string will be '%0.2d'
Expand Down Expand Up @@ -245,12 +250,13 @@ def sshm(servers, command, extra_arguments=None, stdin=None, disable_formatting_
This is a generator to facilitate using the results of each ssh command as
they become available.
@param servers: A string containing the servers to execute "command" on via
@param servers: A list of strings or a string containing the servers to
execute "command" on via
SSH.
Examples:
'example.com'
'example[1-3].com'
'mail[1,3,8].example.com'
['example.com']
['example[1-3].com']
['mail[1,3,8].example.com', 'example.com']
@type servers: str
@param command: A string containing the command to execute.
Expand All @@ -267,6 +273,8 @@ def sshm(servers, command, extra_arguments=None, stdin=None, disable_formatting_
@returns: A list containing (success, handle, message) from each method
call.
"""
if type(servers) == str:
servers = [servers,]
# Disable formatting when requested
global disable_formatting
disable_formatting = disable_formatting_var
Expand All @@ -288,12 +296,13 @@ def sshm(servers, command, extra_arguments=None, stdin=None, disable_formatting_
thread_num = 0
# Only tell the thread to get stdin if there is some.
if_stdin = True if stdin else False
for uri in uri_expansion(servers):
thread = threading.Thread(target=ssh, args=(thread_num, context, uri,
command, extra_arguments, if_stdin))
thread.start()
threads.append(thread)
thread_num += 1
for server_group in servers:
for uri in uri_expansion(server_group):
thread = threading.Thread(target=ssh, args=(thread_num, context, uri,
command, extra_arguments, if_stdin))
thread.start()
threads.append(thread)
thread_num += 1

poller = zmq.Poller()
poller.register(sink, zmq.POLLIN)
Expand Down
18 changes: 17 additions & 1 deletion sshm/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def get_argparse_args(args=None):
parser = argparse.ArgumentParser(
formatter_class=argparse.RawDescriptionHelpFormatter,
description=__long_description__)
parser.add_argument('servers')
parser.add_argument('servers', nargs='+')
parser.add_argument('command')
parser.add_argument('-s', '--sorted-output', action='store_true', default=False,
help='Sort the output by the URI of each instance. This will wait for all instances to finish before showing any output!')
Expand All @@ -46,6 +46,22 @@ def get_argparse_args(args=None):
help='Hide server information on output. This implies sorted.')
parser.add_argument('--version', action='version', version='%(prog)s '+__version__)
args, extra_args = parser.parse_known_args(args=args)

# Move any servers that start with a - to extra_args
new_servers = []
for i in args.servers:
if i.startswith('-'):
extra_args.append(i)
else:
new_servers.append(i)
args.servers = new_servers

# If the comand starts with a -, replace it with the last server and
# move the command to extra_args.
if args.command.startswith('-'):
extra_args.append(args.command)
args.command = args.servers.pop(-1)

if args.quiet:
args.sorted_output = True
return (args, args.command, extra_args)
Expand Down
1 change: 1 addition & 0 deletions sshm/test/test_lib.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ def test_uri_expansion(self):
('10-11.1.2.3-5',['10.1.2.3', '10.1.2.4', '10.1.2.5', '11.1.2.3', '11.1.2.4', '11.1.2.5']),
('192.168.3-5,7.1', ['192.168.3.1', '192.168.4.1', '192.168.5.1', '192.168.7.1']),
('192.168.3-5,7.1:567', ['192.168.3.1:567', '192.168.4.1:567', '192.168.5.1:567', '192.168.7.1:567']),
('192.168.0.-', ['192.168.0.0', '192.168.0.1', '192.168.0.2', '192.168.0.3', '192.168.0.4', '192.168.0.5', '192.168.0.6', '192.168.0.7', '192.168.0.8', '192.168.0.9', '192.168.0.10', '192.168.0.11', '192.168.0.12', '192.168.0.13', '192.168.0.14', '192.168.0.15', '192.168.0.16', '192.168.0.17', '192.168.0.18', '192.168.0.19', '192.168.0.20', '192.168.0.21', '192.168.0.22', '192.168.0.23', '192.168.0.24', '192.168.0.25', '192.168.0.26', '192.168.0.27', '192.168.0.28', '192.168.0.29', '192.168.0.30', '192.168.0.31', '192.168.0.32', '192.168.0.33', '192.168.0.34', '192.168.0.35', '192.168.0.36', '192.168.0.37', '192.168.0.38', '192.168.0.39', '192.168.0.40', '192.168.0.41', '192.168.0.42', '192.168.0.43', '192.168.0.44', '192.168.0.45', '192.168.0.46', '192.168.0.47', '192.168.0.48', '192.168.0.49', '192.168.0.50', '192.168.0.51', '192.168.0.52', '192.168.0.53', '192.168.0.54', '192.168.0.55', '192.168.0.56', '192.168.0.57', '192.168.0.58', '192.168.0.59', '192.168.0.60', '192.168.0.61', '192.168.0.62', '192.168.0.63', '192.168.0.64', '192.168.0.65', '192.168.0.66', '192.168.0.67', '192.168.0.68', '192.168.0.69', '192.168.0.70', '192.168.0.71', '192.168.0.72', '192.168.0.73', '192.168.0.74', '192.168.0.75', '192.168.0.76', '192.168.0.77', '192.168.0.78', '192.168.0.79', '192.168.0.80', '192.168.0.81', '192.168.0.82', '192.168.0.83', '192.168.0.84', '192.168.0.85', '192.168.0.86', '192.168.0.87', '192.168.0.88', '192.168.0.89', '192.168.0.90', '192.168.0.91', '192.168.0.92', '192.168.0.93', '192.168.0.94', '192.168.0.95', '192.168.0.96', '192.168.0.97', '192.168.0.98', '192.168.0.99', '192.168.0.100', '192.168.0.101', '192.168.0.102', '192.168.0.103', '192.168.0.104', '192.168.0.105', '192.168.0.106', '192.168.0.107', '192.168.0.108', '192.168.0.109', '192.168.0.110', '192.168.0.111', '192.168.0.112', '192.168.0.113', '192.168.0.114', '192.168.0.115', '192.168.0.116', '192.168.0.117', '192.168.0.118', '192.168.0.119', '192.168.0.120', '192.168.0.121', '192.168.0.122', '192.168.0.123', '192.168.0.124', '192.168.0.125', '192.168.0.126', '192.168.0.127', '192.168.0.128', '192.168.0.129', '192.168.0.130', '192.168.0.131', '192.168.0.132', '192.168.0.133', '192.168.0.134', '192.168.0.135', '192.168.0.136', '192.168.0.137', '192.168.0.138', '192.168.0.139', '192.168.0.140', '192.168.0.141', '192.168.0.142', '192.168.0.143', '192.168.0.144', '192.168.0.145', '192.168.0.146', '192.168.0.147', '192.168.0.148', '192.168.0.149', '192.168.0.150', '192.168.0.151', '192.168.0.152', '192.168.0.153', '192.168.0.154', '192.168.0.155', '192.168.0.156', '192.168.0.157', '192.168.0.158', '192.168.0.159', '192.168.0.160', '192.168.0.161', '192.168.0.162', '192.168.0.163', '192.168.0.164', '192.168.0.165', '192.168.0.166', '192.168.0.167', '192.168.0.168', '192.168.0.169', '192.168.0.170', '192.168.0.171', '192.168.0.172', '192.168.0.173', '192.168.0.174', '192.168.0.175', '192.168.0.176', '192.168.0.177', '192.168.0.178', '192.168.0.179', '192.168.0.180', '192.168.0.181', '192.168.0.182', '192.168.0.183', '192.168.0.184', '192.168.0.185', '192.168.0.186', '192.168.0.187', '192.168.0.188', '192.168.0.189', '192.168.0.190', '192.168.0.191', '192.168.0.192', '192.168.0.193', '192.168.0.194', '192.168.0.195', '192.168.0.196', '192.168.0.197', '192.168.0.198', '192.168.0.199', '192.168.0.200', '192.168.0.201', '192.168.0.202', '192.168.0.203', '192.168.0.204', '192.168.0.205', '192.168.0.206', '192.168.0.207', '192.168.0.208', '192.168.0.209', '192.168.0.210', '192.168.0.211', '192.168.0.212', '192.168.0.213', '192.168.0.214', '192.168.0.215', '192.168.0.216', '192.168.0.217', '192.168.0.218', '192.168.0.219', '192.168.0.220', '192.168.0.221', '192.168.0.222', '192.168.0.223', '192.168.0.224', '192.168.0.225', '192.168.0.226', '192.168.0.227', '192.168.0.228', '192.168.0.229', '192.168.0.230', '192.168.0.231', '192.168.0.232', '192.168.0.233', '192.168.0.234', '192.168.0.235', '192.168.0.236', '192.168.0.237', '192.168.0.238', '192.168.0.239', '192.168.0.240', '192.168.0.241', '192.168.0.242', '192.168.0.243', '192.168.0.244', '192.168.0.245', '192.168.0.246', '192.168.0.247', '192.168.0.248', '192.168.0.249', '192.168.0.250', '192.168.0.251', '192.168.0.252', '192.168.0.253', '192.168.0.254', '192.168.0.255']),
# URLs
('example.com', ['example.com']),
('example.com:789', ['example.com:789']),
Expand Down
27 changes: 17 additions & 10 deletions sshm/test/test_main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ def test_get_argparse_args(self):
# Valid
provided = ['example.com', 'ls']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, 'example.com')
self.assertEqual(args.servers, ['example.com'])
self.assertFalse(args.strip_whitespace)
self.assertFalse(args.disable_formatting)
self.assertEqual(command, 'ls')
Expand All @@ -31,7 +31,7 @@ def test_get_argparse_args(self):
# Valid
provided = ['example[1-3].com', 'exit']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, 'example[1-3].com')
self.assertEqual(args.servers, ['example[1-3].com'])
self.assertFalse(args.strip_whitespace)
self.assertFalse(args.disable_formatting)
self.assertEqual(command, 'exit')
Expand All @@ -44,27 +44,27 @@ def test_get_argparse_args(self):
self.assertRaises(SystemExit, get_argparse_args, provided)

# Extra arguments
provided = ['example[1-3].com', 'exit', '-o UserKnownHostsFile=/dev/null']
provided = ['example[1-3].com', '"exit"', '-o UserKnownHostsFile=/dev/null']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, 'example[1-3].com')
self.assertEqual(args.servers, ['example[1-3].com'])
self.assertFalse(args.strip_whitespace)
self.assertFalse(args.disable_formatting)
self.assertEqual(command, 'exit')
self.assertEqual(command, '"exit"')
self.assertEqual(extra_args, [provided[2],])

# You can strip whitespace
provided = ['-p', 'example[1-3].com', 'exit', '-o UserKnownHostsFile=/dev/null']
provided = ['-p', '-o UserKnownHostsFile=/dev/null', '-o StrictHostKeyChecking=no', 'example[1-3].com', 'exit']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, 'example[1-3].com')
self.assertEqual(args.servers, ['example[1-3].com'])
self.assertTrue(args.strip_whitespace)
self.assertFalse(args.disable_formatting)
self.assertEqual(command, 'exit')
self.assertEqual(extra_args, [provided[3],])
self.assertEqual(extra_args, provided[1:3])

# Disable formatting
provided = ['-d', 'example.com', 'ls']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, 'example.com')
self.assertEqual(args.servers, ['example.com'])
self.assertFalse(args.strip_whitespace)
self.assertTrue(args.disable_formatting)
self.assertEqual(command, 'ls')
Expand All @@ -73,12 +73,19 @@ def test_get_argparse_args(self):
# You can hide server information output
provided = ['-q', 'example.com', 'ls']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, 'example.com')
self.assertEqual(args.servers, ['example.com'])
self.assertTrue(args.sorted_output)
self.assertTrue(args.quiet)
self.assertEqual(command, 'ls')
self.assertEqual(extra_args, [])

# You can specify multiple groups of targets
provided = ['example.com', 'user@mail.example[1-2].com', 'ls']
args, command, extra_args = get_argparse_args(provided)
self.assertEqual(args.servers, ['example.com', 'user@mail.example[1-2].com'])
self.assertEqual(command, 'ls')
self.assertEqual(extra_args, [])


def test__print_handling_newlines(self):
"""
Expand Down

0 comments on commit 114bb38

Please sign in to comment.