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

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

"""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.
"""

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 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 in hosts:
        if opts.localdir:
            localpath = "%s/%s/%s" % (opts.localdir, host, local)
        else:
            localpath = "%s/%s" % (host, local)
        cmd = ['scp', '-qC']
        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, 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))
    do_pslurp(hosts, remote, local, opts)