#!/usr/bin/env python

# Should work with python 2 or 3
# Iain Murray, 2019

"""
Use like scp(1) but get root permissions (obtained by sudo) at the remote end.

You can override which scp command to use at the local and remote ends by
setting the LOCAL_SCP_CMD and REMOTE_SCP_CMD environment variables.

This wrapper needs to be able to use sudo without a password at the remote end,
as is enabled by default on (for example) most AWS instances.

You can configure sudo not to need a password by running "sudo visudo" on the
remote host and adding:
yourusername    ALL=(ALL) NOPASSWD:/usr/bin/scp
to the last line. This change is dangerous! scp can overwrite any file, so
a malicious process running as yourusername can now easily escalate its
permissions.

You could instead get a password to sudo, by setting REMOTE_SCP_CMD to point to
a script on the remote host with something like the following:
    #!/bin/sh
    export SUDO_ASKPASS=some_command_to_get_password
    sudo -A scp "$@"
But you'll need to write a some_command_to_get_password script to get the
password. You won't be able to interact directly with this command!
Maybe you'll communicate with it over a forwarded port, or by ssh-ing
in separately?
"""


import os, sys

args = sys.argv[:]

if (len(args) < 2) or (args[1] == '-h') or (args[1] == '--help'):
    sys.stderr.write("Help for %s\n" % sys.argv[0] + __doc__ + '\n')
    sys.exit(1)

if os.getenv('SCP_SUDO_DO_SSH', False):
    # Wrap ssh, but rewrite scp command to use sudo
    SSH = os.getenv('SCP_SUDO_SSH', 'ssh')
    try:
        idx = args.index('--') + 2
        assert((args[idx] == 'scp') or (args[idx].startswith('scp ')))
        args[idx] = os.getenv('REMOTE_SCP_CMD', 'sudo scp') + args[idx][3:]
    except:
        sys.stderr.write("\n%s error\n" % sys.argv[0])
        sys.stderr.write("ssh command-line arguments don't seem to be from scp?:\n\n")
        sys.stderr.write('\n'.join(args) + '\n\n')
        sys.exit(255)
    
    args[0] = SSH
    os.execvp(SSH, args)
else:
    # Call scp, telling it to use this script as the ssh command
    # (wrapping any existing custom ssh command)
    SCP = os.getenv('LOCAL_SCP_CMD', 'scp')
    args[0] = SCP
    set_ssh = False
    for idx, opt in enumerate(args):
        if opt == '--':
            break
        if opt == '-S':
            set_ssh = True
            os.environ['SCP_SUDO_SSH'] = args[idx + 1]
            args[idx + 1] = sys.argv[0]
        elif opt.startswith('-S'):
            set_ssh = True
            os.environ['SCP_SUDO_SSH'] = opt[2:]
            args[idx] = '-S' + sys.argv[0]
    if not set_ssh:
        args = [args[0], '-S', sys.argv[0]] + args[1:]
    os.environ['SCP_SUDO_DO_SSH'] = '1'
    os.execvp(SCP, args)