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

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

"""Parallel rsync to the set of nodes in hosts.txt.

For each node, we essentially do a rsync -rv -e ssh local user@host:remote.
Note that remote must be an absolute path.
"""

import os
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] local remote"
    parser.epilog = ("Example: prsync -r -h hosts.txt -l irb2 foo " +
          "/home/irb2/foo")

    parser.add_option('-r', '--recursive', dest='recursive',
            action='store_true', help='recusively copy directories (OPTIONAL)')
    parser.add_option('-a', '--archive', dest='archive', action='store_true',
            help='use rsync -a (archive mode) (OPTIONAL)')
    parser.add_option('-z', '--compress', dest='compress', action='store_true',
            help='use rsync compression (OPTIONAL)')
    parser.add_option('-S', '--ssh-args', metavar="ARGS", dest='ssh_args',
            action='store', help='extra arguments for ssh')
    parser.add_option('-d', '--download', dest='download', action='store_true',
            help='download files instead, reversing local and remote (OPTIONAL)')
    parser.add_option('-i', '--inline', dest='inline', action='store_true',
            help='inline aggregated output and error for each server')

    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('Remote path not specified.')

    if not opts.host_files and not opts.host_strings:
        parser.error('Hosts not specified.')

    return opts, args

def do_prsync(hosts, local, remote, opts):
    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)
    manager = Manager(opts)
    for host, port, user in hosts:
        ssh = ['ssh']
        if opts.options:
            for opt in opts.options:
                ssh += ['-o', opt]
        if port:
            ssh += ['-p', port]
        if opts.ssh_args:
            ssh += [opts.ssh_args]

        cmd = ['rsync', '-e', ' '.join(ssh)]
        if opts.verbose:
            cmd.append('-v')
        if opts.recursive:
            cmd.append('-r')
        if opts.archive:
            cmd.append('-a')
        if opts.compress:
            cmd.append('-z')
        if opts.extra:
            cmd.extend(opts.extra)
        if opts.download:
            if len(local) != 1:
                sys.exit('can only download one target at a time')
            if user:
                cmd.append('%s@%s:%s' % (user, host, local[0]))
            else:
                cmd.append('%s:%s' % (host, local[0]))
            localpath = "%s/%s" % (remote, host)
            os.makedirs(localpath)
            cmd.append(localpath)
        else:
            cmd.extend(local)
            if user:
                cmd.append('%s@%s:%s' % (user, host, remote))
            else:
                cmd.append('%s:%s' % (host, remote))
        t = Task(host, port, user, cmd, opts)
        manager.add_task(t)
    try:
        statuses = manager.run()
    except FatalError:
        sys.exit(1)

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

if __name__ == "__main__":
    opts, args = parse_args()
    local = args[0:-1]
    remote = args[-1]
    try:
        hosts = psshutil.read_host_files(opts.host_files, opts.host_glob,
            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_prsync(hosts, local, remote, opts)
