pslurp

#!/usr/bin/env python
# -*- Mode: python -*-

# Copyright (c) 2009-2012, Andrew McNabb
# Copyright (c) 2003-2008, Brent N. Chun
# Modified 2012 by knktc

"""Parallel scp from the set of nodes in hosts.txt.

For each node, we essentially do a scp [-r] user@host:remote
outdir/<node>/local.  This program also uses the -q (quiet) and -C
(compression) options.  Note that remote must be an absolute path.

Modify note:
2012-5-5
I've delete some security parts in this tool, so that you can write a hosts.txt
with password in it. It is much convenient when you manage servers with different
password.
Remember! This modified tool is INSECURITY! You must take the RISK by yourself!
"""

import os
import re
import sys

parent, bindir = os.path.split(os.path.dirname(os.path.abspath(sys.argv[0])))
if os.path.exists(os.path.join(parent, 'psshlib')):
    sys.path.insert(0, parent)

from psshlib import psshutil
from psshlib.task import Task
from psshlib.manager import Manager, FatalError
from psshlib.cli import common_parser, common_defaults

def option_parser():
    parser = common_parser()
    parser.usage = "%prog [OPTIONS] remote local"
    parser.epilog = ("Example: pslurp -h hosts.txt -L /tmp/outdir -l irb2 " +
            "         /home/irb2/foo.txt foo.txt")

    parser.add_option('-r', '--recursive', dest='recursive',
            action='store_true', help='recusively copy directories (OPTIONAL)')
    parser.add_option('-L', '--localdir', dest='localdir',
            help='output directory for remote file copies')

    return parser

def parse_args():
    parser = option_parser()
    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 do_pslurp(hosts, remote, local, opts):
    if opts.localdir and not os.path.exists(opts.localdir):
        os.makedirs(opts.localdir)
    if opts.outdir and not os.path.exists(opts.outdir):
        os.makedirs(opts.outdir)
    if opts.errdir and not os.path.exists(opts.errdir):
        os.makedirs(opts.errdir)
    for host, port, user, password in hosts:
        if opts.localdir:
            dirname = "%s/%s" % (opts.localdir, host)
        else:
            dirname = host
        if not os.path.exists(dirname):
            os.mkdir(dirname)
    manager = Manager(opts)
    for host, port, user, password in hosts:
        if opts.localdir:
            localpath = "%s/%s/%s" % (opts.localdir, host, local)
        else:
            localpath = "%s/%s" % (host, local)
        cmd = ['scp', '-qC', '-o',
                'StrictHostKeyChecking=no']
        if opts.options:
            for opt in opts.options:
                cmd += ['-o', opt]
        if port:
            cmd += ['-P', port]
        if opts.recursive:
            cmd.append('-r')
        if opts.extra:
            cmd.extend(opts.extra)
        if user:
            cmd.append('%s@%s:%s' % (user, host, remote))
        else:
            cmd.append('%s:%s' % (host, remote))
        cmd.append(localpath)
        t = Task(host, port, user, password, cmd, opts)
        manager.add_task(t)
    try:
        statuses = manager.run()
    except FatalError:
        sys.exit(1)

    if min(statuses) < 0:
        # At least one process was killed.
        sys.exit(3)
    for status in statuses:
        if status == 255:
            sys.exit(4)
    for status in statuses:
        if status != 0:
            sys.exit(5)

if __name__ == "__main__":
    opts, args = parse_args()
    remote = args[0]
    local = args[1]
    if not re.match("^/", remote):
        print("Remote path %s must be an absolute path" % remote)
        sys.exit(3)
    try:
        hosts = psshutil.read_host_files(opts.host_files,
            default_user=opts.user)
    except IOError:
        _, e, _ = sys.exc_info()
        sys.stderr.write('Could not open hosts file: %s\n' % e.strerror)
        sys.exit(1)
    if opts.host_strings:
        for s in opts.host_strings:
            hosts.extend(psshutil.parse_host_string(s, default_user=opts.user))
    
    #check if the hosts if empty
    if hosts == []:
        print "There are no hosts specified!"
        sys.exit(1)    
    
    do_pslurp(hosts, remote, local, opts)