build/upload.py

   1 #!/usr/bin/python
   2 #
   3 # This Source Code Form is subject to the terms of the Mozilla Public
   4 # License, v. 2.0. If a copy of the MPL was not distributed with this
   5 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   6 #
   7 # When run directly, this script expects the following environment variables
   8 # to be set:
   9 # UPLOAD_HOST    : host to upload files to
  10 # UPLOAD_USER    : username on that host
  11 # UPLOAD_PATH    : path on that host to put the files in
  12 #
  13 # And will use the following optional environment variables if set:
  14 # UPLOAD_SSH_KEY : path to a ssh private key to use
  15 # UPLOAD_PORT    : port to use for ssh
  16 # POST_UPLOAD_CMD: a commandline to run on the remote host after uploading.
  17 #                  UPLOAD_PATH and the full paths of all files uploaded will
  18 #                  be appended to the commandline.
  19 #
  20 # All files to be uploaded should be passed as commandline arguments to this
  21 # script. The script takes one other parameter, --base-path, which you can use
  22 # to indicate that files should be uploaded including their paths relative
  23 # to the base path.
  24
  25 import sys, os
  26 from optparse import OptionParser
  27 from subprocess import check_call, check_output
  28 import redo
  29
  30 def RequireEnvironmentVariable(v):
  31     """Return the value of the environment variable named v, or print
  32     an error and exit if it's unset (or empty)."""
  33     if not v in os.environ or os.environ[v] == "":
  34         print "Error: required environment variable %s not set" % v
  35         sys.exit(1)
  36     return os.environ[v]
  37
  38 def OptionalEnvironmentVariable(v):
  39     """Return the value of the environment variable named v, or None
  40     if it's unset (or empty)."""
  41     if v in os.environ and os.environ[v] != "":
  42         return os.environ[v]
  43     return None
  44
  45 def FixupMsysPath(path):
  46     """MSYS helpfully translates absolute pathnames in environment variables
  47     and commandline arguments into Windows native paths. This sucks if you're
  48     trying to pass an absolute path on a remote server. This function attempts
  49     to un-mangle such paths."""
  50     if 'OSTYPE' in os.environ and os.environ['OSTYPE'] == 'msys':
  51         # sort of awful, find out where our shell is (should be in msys/bin)
  52         # and strip the first part of that path out of the other path
  53         if 'SHELL' in os.environ:
  54             sh = os.environ['SHELL']
  55             msys = sh[:sh.find('/bin')]
  56             if path.startswith(msys):
  57                 path = path[len(msys):]
  58     return path
  59
  60 def WindowsPathToMsysPath(path):
  61     """Translate a Windows pathname to an MSYS pathname.
  62     Necessary because we call out to ssh/scp, which are MSYS binaries
  63     and expect MSYS paths."""
  64     if sys.platform != 'win32':
  65         return path
  66     (drive, path) = os.path.splitdrive(os.path.abspath(path))
  67     return "/" + drive[0] + path.replace('\\','/')
  68
  69 def AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key):
  70     """Given optional port and ssh key values, append valid OpenSSH
  71     commandline arguments to the list cmdline if the values are not None."""
  72     if port is not None:
  73         cmdline.append("-P%d" % port)
  74     if ssh_key is not None:
  75         # Don't interpret ~ paths - ssh can handle that on its own
  76         if not ssh_key.startswith('~'):
  77             ssh_key = WindowsPathToMsysPath(ssh_key)
  78         cmdline.extend(["-o", "IdentityFile=%s" % ssh_key])
  79
  80 def DoSSHCommand(command, user, host, port=None, ssh_key=None):
  81     """Execute command on user@host using ssh. Optionally use
  82     port and ssh_key, if provided."""
  83     cmdline = ["ssh"]
  84     AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key)
  85     cmdline.extend(["%s@%s" % (user, host), command])
  86
  87     with redo.retrying(check_output, sleeptime=10) as f:
  88         output = f(cmdline).strip()
  89         return output
  90
  91     raise Exception("Command %s returned non-zero exit code" % cmdline)
  92
  93 def DoSCPFile(file, remote_path, user, host, port=None, ssh_key=None):
  94     """Upload file to user@host:remote_path using scp. Optionally use
  95     port and ssh_key, if provided."""
  96     cmdline = ["scp"]
  97     AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key)
  98     cmdline.extend([WindowsPathToMsysPath(file),
  99                     "%s@%s:%s" % (user, host, remote_path)])
 100     with redo.retrying(check_call, sleeptime=10) as f:
 101         f(cmdline)
 102         return
 103
 104     raise Exception("Command %s returned non-zero exit code" % cmdline)
 105
 106 def GetRemotePath(path, local_file, base_path):
 107     """Given a remote path to upload to, a full path to a local file, and an
 108     optional full path that is a base path of the local file, construct the
 109     full remote path to place the file in. If base_path is not None, include
 110     the relative path from base_path to file."""
 111     if base_path is None or not local_file.startswith(base_path):
 112         return path
 113     dir = os.path.dirname(local_file)
 114     # strip base_path + extra slash and make it unixy
 115     dir = dir[len(base_path)+1:].replace('\\','/')
 116     return path + dir
 117
 118 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):
 119     """Upload each file in the list files to user@host:path. Optionally pass
 120     port and ssh_key to the ssh commands. If base_path is not None, upload
 121     files including their path relative to base_path. If upload_to_temp_dir is
 122     True files will be uploaded to a temporary directory on the remote server.
 123     Generally, you should have a post upload command specified in these cases
 124     that can move them around to their correct location(s).
 125     If post_upload_command is not None, execute that command on the remote host
 126     after uploading all files, passing it the upload path, and the full paths to
 127     all files uploaded.
 128     If verbose is True, print status updates while working."""
 129     if upload_to_temp_dir:
 130         path = DoSSHCommand("mktemp -d", user, host, port=port, ssh_key=ssh_key)
 131     if not path.endswith("/"):
 132         path += "/"
 133     if base_path is not None:
 134         base_path = os.path.abspath(base_path)
 135     remote_files = []
 136     try:
 137         for file in files:
 138             file = os.path.abspath(file)
 139             if not os.path.isfile(file):
 140                 raise IOError("File not found: %s" % file)
 141             # first ensure that path exists remotely
 142             remote_path = GetRemotePath(path, file, base_path)
 143             DoSSHCommand("mkdir -p " + remote_path, user, host, port=port, ssh_key=ssh_key)
 144             if verbose:
 145                 print "Uploading " + file
 146             DoSCPFile(file, remote_path, user, host, port=port, ssh_key=ssh_key)
 147             remote_files.append(remote_path + '/' + os.path.basename(file))
 148         if post_upload_command is not None:
 149             if verbose:
 150                 print "Running post-upload command: " + post_upload_command
 151             file_list = '"' + '" "'.join(remote_files) + '"'
 152             DoSSHCommand('%s "%s" %s' % (post_upload_command, path, file_list), user, host, port=port, ssh_key=ssh_key)
 153     finally:
 154         if upload_to_temp_dir:
 155             DoSSHCommand("rm -rf %s" % path, user, host, port=port,
 156                          ssh_key=ssh_key)
 157     if verbose:
 158         print "Upload complete"
 159
 160 if __name__ == '__main__':
 161     host = RequireEnvironmentVariable('UPLOAD_HOST')
 162     user = RequireEnvironmentVariable('UPLOAD_USER')
 163     path = OptionalEnvironmentVariable('UPLOAD_PATH')
 164     upload_to_temp_dir = OptionalEnvironmentVariable('UPLOAD_TO_TEMP')
 165     port = OptionalEnvironmentVariable('UPLOAD_PORT')
 166     if port is not None:
 167         port = int(port)
 168     key = OptionalEnvironmentVariable('UPLOAD_SSH_KEY')
 169     post_upload_command = OptionalEnvironmentVariable('POST_UPLOAD_CMD')
 170     if (not path and not upload_to_temp_dir) or (path and upload_to_temp_dir):
 171         print "One (and only one of UPLOAD_PATH or UPLOAD_TO_TEMP must be " + \
 172               "defined."
 173         sys.exit(1)
 174     if sys.platform == 'win32':
 175         if path is not None:
 176             path = FixupMsysPath(path)
 177         if post_upload_command is not None:
 178             post_upload_command = FixupMsysPath(post_upload_command)
 179
 180     parser = OptionParser(usage="usage: %prog [options] <files>")
 181     parser.add_option("-b", "--base-path",
 182                       action="store", dest="base_path",
 183                       help="Preserve file paths relative to this path when uploading. If unset, all files will be uploaded directly to UPLOAD_PATH.")
 184     (options, args) = parser.parse_args()
 185     if len(args) < 1:
 186         print "You must specify at least one file to upload"
 187         sys.exit(1)
 188     try:
 189         UploadFiles(user, host, path, args, base_path=options.base_path,
 190                     port=port, ssh_key=key, upload_to_temp_dir=upload_to_temp_dir,
 191                     post_upload_command=post_upload_command,
 192                     verbose=True)
 193     except IOError, (strerror):
 194         print strerror
 195         sys.exit(1)
 196     except Exception, (err):
 197         print err
 198         sys.exit(2)