| OLD | NEW |
| 1 #!/usr/bin/env python | 1 #!/usr/bin/env python |
| 2 # coding: utf-8 | 2 # coding: utf-8 |
| 3 | 3 |
| 4 import argparse | 4 import argparse |
| 5 import sys | 5 import sys |
| 6 import os | 6 import os |
| 7 import posixpath |
| 7 import re | 8 import re |
| 8 import subprocess | 9 import subprocess |
| 9 import yaml | 10 import yaml |
| 10 | 11 |
| 11 def createArgumentParser(**kwargs): | 12 def createArgumentParser(**kwargs): |
| 12 parser = argparse.ArgumentParser(**kwargs) | 13 parser = argparse.ArgumentParser(**kwargs) |
| 13 parser.add_argument( | 14 parser.add_argument( |
| 14 '-u', '--user', metavar='user', dest='user', type=str, default='vagrant', | 15 '-u', '--user', metavar='user', dest='user', type=str, |
| 15 help='user name for use with SSH, must exist on all hosts' | 16 help='user name for use with SSH, must exist on all hosts' |
| 16 ) | 17 ) |
| 17 | 18 |
| 19 parser.add_argument( |
| 20 '-l', '--local', action='store_false', dest='remote', default=None, |
| 21 help='use the local version of hosts.yaml' |
| 22 ) |
| 23 |
| 24 parser.add_argument( |
| 25 '-r', '--remote', metavar='master', dest='remote', type=str, |
| 26 help='use a remote (puppet-master) version of hosts.yaml' |
| 27 ) |
| 28 |
| 18 return parser | 29 return parser |
| 19 | 30 |
| 20 def parseOptions(args): | 31 def parseOptions(args): |
| 21 description = 'Run a command on the given hosts or groups of hosts' | 32 description = 'Run a command on the given hosts or groups of hosts' |
| 22 parser = createArgumentParser(description=description) | 33 parser = createArgumentParser(description=description) |
| 23 parser.add_argument( | 34 parser.add_argument( |
| 24 '-i', '--ignore-errors', action='store_true', dest='ignore_errors', | 35 '-i', '--ignore-errors', action='store_true', dest='ignore_errors', |
| 25 help='continue execution on next host in case of an error' | 36 help='continue execution on next host in case of an error' |
| 26 ) | 37 ) |
| 27 | 38 |
| 28 hosts = set() | 39 hosts = set() |
| 29 parser.add_argument( | 40 parser.add_argument( |
| 30 '-t', '--target', metavar='host|group', | 41 '-t', '--target', metavar='host|group', |
| 31 help='target host or group, can be specified multiple times', | 42 help='target host or group, can be specified multiple times', |
| 32 type=lambda value: hosts.update([value]) | 43 type=lambda value: hosts.update([value]) |
| 33 ) | 44 ) |
| 34 | 45 |
| 35 parser.add_argument( | 46 parser.add_argument( |
| 36 'args', metavar='command', type=str, nargs='+', | 47 'args', metavar='command', type=str, nargs='+', |
| 37 help='the command to run on the specified hosts' | 48 help='the command to run on the specified hosts' |
| 38 ) | 49 ) |
| 39 | 50 |
| 40 options = parser.parse_args(args) | 51 options = parser.parse_args(args) |
| 41 options.hosts = hosts | 52 options.hosts = hosts |
| 42 return options | 53 return options |
| 43 | 54 |
| 44 def getValidHosts(): | 55 def getValidHosts(options): |
| 45 dirname = os.path.dirname(sys.argv[0]) | 56 path_canonical = ('modules', 'private', 'hiera', 'hosts.yaml') |
| 46 path_name = os.path.join(dirname, 'modules', 'private', 'hiera', 'hosts.yaml') | 57 |
| 47 with open(path_name, 'rb') as handle: | 58 if options.remote: |
| 48 config = yaml.load(handle) | 59 login = ['-l', options.user] if options.user else [] |
| 60 path_name = posixpath.join('/etc/puppet/infrastructure', *path_canonical) |
| 61 command = ['ssh'] + login + [options.remote, '--', 'sudo', 'cat', path_name] |
| 62 child = subprocess.Popen(command, stderr=sys.stderr, stdout=subprocess.PIPE) |
| 63 try: |
| 64 config = yaml.load(child.stdout) |
| 65 finally: |
| 66 child.stdout.close() |
| 67 child.wait() |
| 68 elif options.remote is False: |
| 69 dirname = os.path.dirname(sys.argv[0]) |
| 70 path_name = os.path.join(dirname, *path_canonical) |
| 71 with open(path_name, 'rb') as handle: |
| 72 config = yaml.load(handle) |
| 73 else: |
| 74 sys.exit('Please either specify a --remote host or use --local') |
| 75 |
| 49 servers = config.get('servers', {}) | 76 servers = config.get('servers', {}) |
| 50 return servers | 77 return servers |
| 51 | 78 |
| 52 def resolveHostList(hosts): | 79 def resolveHostList(options): |
| 53 | 80 |
| 54 result = set() | 81 result = set() |
| 55 | 82 |
| 56 try: | 83 try: |
| 57 valid_hosts = getValidHosts() | 84 valid_hosts = getValidHosts(options) |
| 58 except Warning as error: | 85 except Warning as error: |
| 59 print >>sys.stderr, 'Warning: failed to determine valid hosts:', error | 86 print >>sys.stderr, 'Warning: failed to determine valid hosts:', error |
| 60 result.update(hosts) | 87 result.update(options.hosts) |
| 61 else: | 88 else: |
| 62 for name in hosts: | 89 for name in options.hosts: |
| 63 chunk = [ | 90 chunk = [ |
| 64 value.get('dns', key) for (key, value) in valid_hosts.items() | 91 value.get('dns', key) for (key, value) in valid_hosts.items() |
| 65 | 92 |
| 66 if name == key | 93 if name == key |
| 67 or name == '*' | 94 or name == '*' |
| 68 or name == value.get('dns', None) | 95 or name == value.get('dns', None) |
| 69 or name in value.get('groups', ()) | 96 or name in value.get('groups', ()) |
| 70 ] | 97 ] |
| 71 | 98 |
| 72 if len(chunk) == 0: | 99 if len(chunk) == 0: |
| 73 print >>sys.stderr, 'Warning: failed to recognize host or group', name | 100 print >>sys.stderr, 'Warning: failed to recognize host or group', name |
| 74 else: | 101 else: |
| 75 result.update(chunk) | 102 result.update(chunk) |
| 76 | 103 |
| 77 return result | 104 return result |
| 78 | 105 |
| 79 def runCommand(user, host, command, ignore_errors=False): | 106 def runCommand(user, host, command, ignore_errors=False): |
| 80 if not isinstance(command, list): | 107 if not isinstance(command, list): |
| 81 command = [command] | 108 command = [command] |
| 82 command = ['ssh'] + (['-l', user] if user else []) + [host] + command | 109 command = ['ssh'] + (['-l', user] if user else []) + [host] + command |
| 83 if ignore_errors: | 110 if ignore_errors: |
| 84 subprocess.call(command) | 111 subprocess.call(command) |
| 85 else: | 112 else: |
| 86 subprocess.check_call(command) | 113 subprocess.check_call(command) |
| 87 | 114 |
| 88 if __name__ == '__main__': | 115 if __name__ == '__main__': |
| 89 options = parseOptions(sys.argv[1:]) | 116 options = parseOptions(sys.argv[1:]) |
| 90 selectedHosts = resolveHostList(options.hosts) | 117 selectedHosts = resolveHostList(options) |
| 91 if len(selectedHosts) == 0: | 118 if len(selectedHosts) == 0: |
| 92 print >>sys.stderr, 'No valid hosts or groups specified, nothing to do' | 119 print >>sys.stderr, 'No valid hosts or groups specified, nothing to do' |
| 93 sys.exit(0) | 120 sys.exit(0) |
| 94 for host in selectedHosts: | 121 for host in selectedHosts: |
| 95 print >>sys.stderr, 'Running on %s...' % host | 122 print >>sys.stderr, 'Running on %s...' % host |
| 96 runCommand(options.user, host, options.args, options.ignore_errors) | 123 runCommand(options.user, host, options.args, options.ignore_errors) |
| OLD | NEW |