blob: 1130dc7a92d7fe24aa7c6353172091dfd680dcf6 [file] [log] [blame]
# Copyright 2012 Google Inc. All Rights Reserved.
"""Run the given shell command for each of the given hosts."""
import os
import re
import select
import subprocess
import sys
import config
from portsh import options
__author__ = 'Avery Pennarun ('
optspec = """
run [options...] <hosts> -- <command string...>
r,raw Raw output (don't redirect stdin/stdout/stderr)
q,quiet Quiet: don't print to stderr for nonzero exit codes
def _FixNull(s):
if s is None:
return ''
return str(s)
def _ArgSub(argv, host):
for arg in argv:
argout = []
for part in re.split(r'(\$\w+)', arg):
if part.startswith('$'):
for var in config.Host.__slots__:
if part[1:] == var:
part = _FixNull(getattr(host, var))
yield ''.join(argout)
class Handler(object):
"""When data is received from the given file, print it line by line."""
def __init__(self, prefix, fileobj):
self.prefix = prefix
self.fileobj = fileobj
self.buf = ''
def __del__(self):
if self.buf and not self.buf.endswith('\n'):
self.buf += '\n'
def fileno(self): #gpylint: disable-msg=C6409
return self.fileobj.fileno()
def PrintBuf(self):
while '\n' in self.buf:
line, self.buf = self.buf.split('\n', 1)
print '%s: %s' % (self.prefix, line)
def Run(self):
"""Read from this handler and write its results to stdout."""
got =, 65536)
self.buf += got
return got
def main():
o = options.Options(optspec)
args = []
cmd = []
for i in range(1, len(sys.argv)):
if sys.argv[i] == '--':
args = sys.argv[1:i]
cmd = sys.argv[i+1:]
opt, _, hostnames = o.parse(args)
if not hostnames or not cmd:
o.fatal('you must specify hosts, --, and a command string')
hosts = [config.hosts.FindOne(name=i) for i in hostnames]
except KeyError, e:
o.fatal('host not configured: %s' % e)
procs = []
handlers = []
for host in hosts:
env = dict(os.environ)
env.update(dict((key, _FixNull(getattr(host, key)))
for key in config.Host.__slots__))
if len(cmd) > 1:
argv = list(_ArgSub(cmd, host))
shell = False
argv = cmd[0]
shell = True
if opt.raw:
p = subprocess.Popen(argv, env=env, shell=shell)
p = subprocess.Popen(argv, env=env, shell=shell,
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
handlers.append(Handler(, p.stdout))
handlers.append(Handler(, p.stderr))
procs.append((host, p))
while handlers:
r, _, _ =, [], [])
for handler in r:
if not handler.Run():
final_rv = 0
for host, p in procs:
rv = p.wait()
if rv != 0:
if not opt.quiet:
sys.stderr.write('-- %s: exited with error code %d\n'
% (, rv))
final_rv = 1
return final_rv
if __name__ == '__main__':