michael@0: #!/usr/bin/python
michael@0: #
michael@0: # This Source Code Form is subject to the terms of the Mozilla Public
michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/.
michael@0: #
michael@0: # When run directly, this script expects the following environment variables
michael@0: # to be set:
michael@0: # UPLOAD_HOST    : host to upload files to
michael@0: # UPLOAD_USER    : username on that host
michael@0: # UPLOAD_PATH    : path on that host to put the files in
michael@0: #
michael@0: # And will use the following optional environment variables if set:
michael@0: # UPLOAD_SSH_KEY : path to a ssh private key to use
michael@0: # UPLOAD_PORT    : port to use for ssh
michael@0: # POST_UPLOAD_CMD: a commandline to run on the remote host after uploading.
michael@0: #                  UPLOAD_PATH and the full paths of all files uploaded will
michael@0: #                  be appended to the commandline.
michael@0: #
michael@0: # All files to be uploaded should be passed as commandline arguments to this
michael@0: # script. The script takes one other parameter, --base-path, which you can use
michael@0: # to indicate that files should be uploaded including their paths relative
michael@0: # to the base path.
michael@0: 
michael@0: import sys, os
michael@0: from optparse import OptionParser
michael@0: from subprocess import PIPE, Popen, check_call
michael@0: 
michael@0: def RequireEnvironmentVariable(v):
michael@0:     """Return the value of the environment variable named v, or print
michael@0:     an error and exit if it's unset (or empty)."""
michael@0:     if not v in os.environ or os.environ[v] == "":
michael@0:         print "Error: required environment variable %s not set" % v
michael@0:         sys.exit(1)
michael@0:     return os.environ[v]
michael@0: 
michael@0: def OptionalEnvironmentVariable(v):
michael@0:     """Return the value of the environment variable named v, or None
michael@0:     if it's unset (or empty)."""
michael@0:     if v in os.environ and os.environ[v] != "":
michael@0:         return os.environ[v]
michael@0:     return None
michael@0: 
michael@0: def FixupMsysPath(path):
michael@0:     """MSYS helpfully translates absolute pathnames in environment variables
michael@0:     and commandline arguments into Windows native paths. This sucks if you're
michael@0:     trying to pass an absolute path on a remote server. This function attempts
michael@0:     to un-mangle such paths."""
michael@0:     if 'OSTYPE' in os.environ and os.environ['OSTYPE'] == 'msys':
michael@0:         # sort of awful, find out where our shell is (should be in msys/bin)
michael@0:         # and strip the first part of that path out of the other path
michael@0:         if 'SHELL' in os.environ:
michael@0:             sh = os.environ['SHELL']
michael@0:             msys = sh[:sh.find('/bin')]
michael@0:             if path.startswith(msys):
michael@0:                 path = path[len(msys):]
michael@0:     return path
michael@0: 
michael@0: def WindowsPathToMsysPath(path):
michael@0:     """Translate a Windows pathname to an MSYS pathname.
michael@0:     Necessary because we call out to ssh/scp, which are MSYS binaries
michael@0:     and expect MSYS paths."""
michael@0:     if sys.platform != 'win32':
michael@0:         return path
michael@0:     (drive, path) = os.path.splitdrive(os.path.abspath(path))
michael@0:     return "/" + drive[0] + path.replace('\\','/')
michael@0: 
michael@0: def AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key):
michael@0:     """Given optional port and ssh key values, append valid OpenSSH
michael@0:     commandline arguments to the list cmdline if the values are not None."""
michael@0:     if port is not None:
michael@0:         cmdline.append("-P%d" % port)
michael@0:     if ssh_key is not None:
michael@0:         # Don't interpret ~ paths - ssh can handle that on its own
michael@0:         if not ssh_key.startswith('~'):
michael@0:             ssh_key = WindowsPathToMsysPath(ssh_key)
michael@0:         cmdline.extend(["-o", "IdentityFile=%s" % ssh_key])
michael@0: 
michael@0: def DoSSHCommand(command, user, host, port=None, ssh_key=None):
michael@0:     """Execute command on user@host using ssh. Optionally use
michael@0:     port and ssh_key, if provided."""
michael@0:     cmdline = ["ssh"]
michael@0:     AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key)
michael@0:     cmdline.extend(["%s@%s" % (user, host), command])
michael@0:     cmd = Popen(cmdline, stdout=PIPE)
michael@0:     retcode = cmd.wait()
michael@0:     if retcode != 0:
michael@0:         raise Exception("Command %s returned non-zero exit code: %i" % \
michael@0:           (cmdline, retcode))
michael@0:     return cmd.stdout.read().strip()
michael@0: 
michael@0: def DoSCPFile(file, remote_path, user, host, port=None, ssh_key=None):
michael@0:     """Upload file to user@host:remote_path using scp. Optionally use
michael@0:     port and ssh_key, if provided."""
michael@0:     cmdline = ["scp"]
michael@0:     AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key)
michael@0:     cmdline.extend([WindowsPathToMsysPath(file),
michael@0:                     "%s@%s:%s" % (user, host, remote_path)])
michael@0:     check_call(cmdline)
michael@0: 
michael@0: def GetRemotePath(path, local_file, base_path):
michael@0:     """Given a remote path to upload to, a full path to a local file, and an
michael@0:     optional full path that is a base path of the local file, construct the
michael@0:     full remote path to place the file in. If base_path is not None, include
michael@0:     the relative path from base_path to file."""
michael@0:     if base_path is None or not local_file.startswith(base_path):
michael@0:         return path
michael@0:     dir = os.path.dirname(local_file)
michael@0:     # strip base_path + extra slash and make it unixy
michael@0:     dir = dir[len(base_path)+1:].replace('\\','/')
michael@0:     return path + dir
michael@0: 
michael@0: def UploadFiles(user, host, path, files, verbose=False, port=None, ssh_key=None, base_path=None, upload_to_temp_dir=False, post_upload_command=None):
michael@0:     """Upload each file in the list files to user@host:path. Optionally pass
michael@0:     port and ssh_key to the ssh commands. If base_path is not None, upload
michael@0:     files including their path relative to base_path. If upload_to_temp_dir is
michael@0:     True files will be uploaded to a temporary directory on the remote server.
michael@0:     Generally, you should have a post upload command specified in these cases
michael@0:     that can move them around to their correct location(s).
michael@0:     If post_upload_command is not None, execute that command on the remote host
michael@0:     after uploading all files, passing it the upload path, and the full paths to
michael@0:     all files uploaded.
michael@0:     If verbose is True, print status updates while working."""
michael@0:     if upload_to_temp_dir:
michael@0:         path = DoSSHCommand("mktemp -d", user, host, port=port, ssh_key=ssh_key)
michael@0:     if not path.endswith("/"):
michael@0:         path += "/"
michael@0:     if base_path is not None:
michael@0:         base_path = os.path.abspath(base_path)
michael@0:     remote_files = []
michael@0:     try:
michael@0:         for file in files:
michael@0:             file = os.path.abspath(file)
michael@0:             if not os.path.isfile(file):
michael@0:                 raise IOError("File not found: %s" % file)
michael@0:             # first ensure that path exists remotely
michael@0:             remote_path = GetRemotePath(path, file, base_path)
michael@0:             DoSSHCommand("mkdir -p " + remote_path, user, host, port=port, ssh_key=ssh_key)
michael@0:             if verbose:
michael@0:                 print "Uploading " + file
michael@0:             DoSCPFile(file, remote_path, user, host, port=port, ssh_key=ssh_key)
michael@0:             remote_files.append(remote_path + '/' + os.path.basename(file))
michael@0:         if post_upload_command is not None:
michael@0:             if verbose:
michael@0:                 print "Running post-upload command: " + post_upload_command
michael@0:             file_list = '"' + '" "'.join(remote_files) + '"'
michael@0:             DoSSHCommand('%s "%s" %s' % (post_upload_command, path, file_list), user, host, port=port, ssh_key=ssh_key)
michael@0:     finally:
michael@0:         if upload_to_temp_dir:
michael@0:             DoSSHCommand("rm -rf %s" % path, user, host, port=port,
michael@0:                          ssh_key=ssh_key)
michael@0:     if verbose:
michael@0:         print "Upload complete"
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     host = RequireEnvironmentVariable('UPLOAD_HOST')
michael@0:     user = RequireEnvironmentVariable('UPLOAD_USER')
michael@0:     path = OptionalEnvironmentVariable('UPLOAD_PATH')
michael@0:     upload_to_temp_dir = OptionalEnvironmentVariable('UPLOAD_TO_TEMP')
michael@0:     port = OptionalEnvironmentVariable('UPLOAD_PORT')
michael@0:     if port is not None:
michael@0:         port = int(port)
michael@0:     key = OptionalEnvironmentVariable('UPLOAD_SSH_KEY')
michael@0:     post_upload_command = OptionalEnvironmentVariable('POST_UPLOAD_CMD')
michael@0:     if (not path and not upload_to_temp_dir) or (path and upload_to_temp_dir):
michael@0:         print "One (and only one of UPLOAD_PATH or UPLOAD_TO_TEMP must be " + \
michael@0:               "defined."
michael@0:         sys.exit(1)
michael@0:     if sys.platform == 'win32':
michael@0:         if path is not None:
michael@0:             path = FixupMsysPath(path)
michael@0:         if post_upload_command is not None:
michael@0:             post_upload_command = FixupMsysPath(post_upload_command)
michael@0: 
michael@0:     parser = OptionParser(usage="usage: %prog [options] <files>")
michael@0:     parser.add_option("-b", "--base-path",
michael@0:                       action="store", dest="base_path",
michael@0:                       help="Preserve file paths relative to this path when uploading. If unset, all files will be uploaded directly to UPLOAD_PATH.")
michael@0:     (options, args) = parser.parse_args()
michael@0:     if len(args) < 1:
michael@0:         print "You must specify at least one file to upload"
michael@0:         sys.exit(1)
michael@0:     try:
michael@0:         UploadFiles(user, host, path, args, base_path=options.base_path,
michael@0:                     port=port, ssh_key=key, upload_to_temp_dir=upload_to_temp_dir,
michael@0:                     post_upload_command=post_upload_command,
michael@0:                     verbose=True)
michael@0:     except IOError, (strerror):
michael@0:         print strerror
michael@0:         sys.exit(1)
michael@0:     except Exception, (err):
michael@0:         print err
michael@0:         sys.exit(2)