Back to index

enigmail  1.4.3
upload.py
Go to the documentation of this file.
00001 #!/usr/bin/python
00002 #
00003 # ***** BEGIN LICENSE BLOCK *****
00004 # Version: MPL 1.1/GPL 2.0/LGPL 2.1
00005 #
00006 # The contents of this file are subject to the Mozilla Public License Version
00007 # 1.1 (the "License"); you may not use this file except in compliance with
00008 # the License. You may obtain a copy of the License at
00009 # http://www.mozilla.org/MPL/
00010 #
00011 # Software distributed under the License is distributed on an "AS IS" basis,
00012 # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
00013 # for the specific language governing rights and limitations under the
00014 # License.
00015 #
00016 # The Original Code is mozilla.org code.
00017 #
00018 # The Initial Developer of the Original Code is
00019 # The Mozilla Foundation
00020 # Portions created by the Initial Developer are Copyright (C) 2008
00021 # the Initial Developer. All Rights Reserved.
00022 #
00023 # Contributor(s):
00024 #  Ted Mielczarek <ted.mielczarek@gmail.com>
00025 #
00026 # Alternatively, the contents of this file may be used under the terms of
00027 # either the GNU General Public License Version 2 or later (the "GPL"), or
00028 # the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
00029 # in which case the provisions of the GPL or the LGPL are applicable instead
00030 # of those above. If you wish to allow use of your version of this file only
00031 # under the terms of either the GPL or the LGPL, and not to allow others to
00032 # use your version of this file under the terms of the MPL, indicate your
00033 # decision by deleting the provisions above and replace them with the notice
00034 # and other provisions required by the GPL or the LGPL. If you do not delete
00035 # the provisions above, a recipient may use your version of this file under
00036 # the terms of any one of the MPL, the GPL or the LGPL.
00037 #
00038 # ***** END LICENSE BLOCK *****
00039 #
00040 # When run directly, this script expects the following environment variables
00041 # to be set:
00042 # UPLOAD_HOST    : host to upload files to
00043 # UPLOAD_USER    : username on that host
00044 # UPLOAD_PATH    : path on that host to put the files in
00045 #
00046 # And will use the following optional environment variables if set:
00047 # UPLOAD_SSH_KEY : path to a ssh private key to use
00048 # UPLOAD_PORT    : port to use for ssh
00049 # POST_UPLOAD_CMD: a commandline to run on the remote host after uploading.
00050 #                  UPLOAD_PATH and the full paths of all files uploaded will
00051 #                  be appended to the commandline.
00052 #
00053 # All files to be uploaded should be passed as commandline arguments to this
00054 # script. The script takes one other parameter, --base-path, which you can use
00055 # to indicate that files should be uploaded including their paths relative
00056 # to the base path.
00057 
00058 import sys, os
00059 from optparse import OptionParser
00060 from subprocess import PIPE, Popen, check_call
00061 
00062 def RequireEnvironmentVariable(v):
00063     """Return the value of the environment variable named v, or print
00064     an error and exit if it's unset (or empty)."""
00065     if not v in os.environ or os.environ[v] == "":
00066         print "Error: required environment variable %s not set" % v
00067         sys.exit(1)
00068     return os.environ[v]
00069 
00070 def OptionalEnvironmentVariable(v):
00071     """Return the value of the environment variable named v, or None
00072     if it's unset (or empty)."""
00073     if v in os.environ and os.environ[v] != "":
00074         return os.environ[v]
00075     return None
00076 
00077 def FixupMsysPath(path):
00078     """MSYS helpfully translates absolute pathnames in environment variables
00079     and commandline arguments into Windows native paths. This sucks if you're
00080     trying to pass an absolute path on a remote server. This function attempts
00081     to un-mangle such paths."""
00082     if 'OSTYPE' in os.environ and os.environ['OSTYPE'] == 'msys':
00083         # sort of awful, find out where our shell is (should be in msys/bin)
00084         # and strip the first part of that path out of the other path
00085         if 'SHELL' in os.environ:
00086             sh = os.environ['SHELL']
00087             msys = sh[:sh.find('/bin')]
00088             if path.startswith(msys):
00089                 path = path[len(msys):]
00090     return path
00091 
00092 def WindowsPathToMsysPath(path):
00093     """Translate a Windows pathname to an MSYS pathname.
00094     Necessary because we call out to ssh/scp, which are MSYS binaries
00095     and expect MSYS paths."""
00096     if sys.platform != 'win32':
00097         return path
00098     (drive, path) = os.path.splitdrive(os.path.abspath(path))
00099     return "/" + drive[0] + path.replace('\\','/')
00100 
00101 def AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key):
00102     """Given optional port and ssh key values, append valid OpenSSH
00103     commandline arguments to the list cmdline if the values are not None."""
00104     if port is not None:
00105         cmdline.append("-P%d" % port)
00106     if ssh_key is not None:
00107         # Don't interpret ~ paths - ssh can handle that on its own
00108         if not ssh_key.startswith('~'):
00109             ssh_key = WindowsPathToMsysPath(ssh_key)
00110         cmdline.extend(["-o", "IdentityFile=%s" % ssh_key])
00111 
00112 def DoSSHCommand(command, user, host, port=None, ssh_key=None):
00113     """Execute command on user@host using ssh. Optionally use
00114     port and ssh_key, if provided."""
00115     cmdline = ["ssh"]
00116     AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key)
00117     cmdline.extend(["%s@%s" % (user, host), command])
00118     cmd = Popen(cmdline, stdout=PIPE)
00119     retcode = cmd.wait()
00120     if retcode != 0:
00121         raise Exception("Command %s returned non-zero exit code: %i" % \
00122           (cmdline, retcode))
00123     return cmd.stdout.read().strip()
00124 
00125 def DoSCPFile(file, remote_path, user, host, port=None, ssh_key=None):
00126     """Upload file to user@host:remote_path using scp. Optionally use
00127     port and ssh_key, if provided."""
00128     cmdline = ["scp"]
00129     AppendOptionalArgsToSSHCommandline(cmdline, port, ssh_key)
00130     cmdline.extend([WindowsPathToMsysPath(file),
00131                     "%s@%s:%s" % (user, host, remote_path)])
00132     check_call(cmdline)
00133 
00134 def GetRemotePath(path, local_file, base_path):
00135     """Given a remote path to upload to, a full path to a local file, and an
00136     optional full path that is a base path of the local file, construct the
00137     full remote path to place the file in. If base_path is not None, include
00138     the relative path from base_path to file."""
00139     if base_path is None or not local_file.startswith(base_path):
00140         return path
00141     dir = os.path.dirname(local_file)
00142     # strip base_path + extra slash and make it unixy
00143     dir = dir[len(base_path)+1:].replace('\\','/')
00144     return path + dir
00145 
00146 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):
00147     """Upload each file in the list files to user@host:path. Optionally pass
00148     port and ssh_key to the ssh commands. If base_path is not None, upload
00149     files including their path relative to base_path. If upload_to_temp_dir is
00150     True files will be uploaded to a temporary directory on the remote server.
00151     Generally, you should have a post upload command specified in these cases
00152     that can move them around to their correct location(s).
00153     If post_upload_command is not None, execute that command on the remote host
00154     after uploading all files, passing it the upload path, and the full paths to
00155     all files uploaded.
00156     If verbose is True, print status updates while working."""
00157     if upload_to_temp_dir:
00158         path = DoSSHCommand("mktemp -d", user, host, port=port, ssh_key=ssh_key)
00159     if not path.endswith("/"):
00160         path += "/"
00161     if base_path is not None:
00162         base_path = os.path.abspath(base_path)
00163     remote_files = []
00164     try:
00165         for file in files:
00166             file = os.path.abspath(file)
00167             if not os.path.isfile(file):
00168                 raise IOError("File not found: %s" % file)
00169             # first ensure that path exists remotely
00170             remote_path = GetRemotePath(path, file, base_path)
00171             DoSSHCommand("mkdir -p " + remote_path, user, host, port=port, ssh_key=ssh_key)
00172             if verbose:
00173                 print "Uploading " + file
00174             DoSCPFile(file, remote_path, user, host, port=port, ssh_key=ssh_key)
00175             remote_files.append(remote_path + '/' + os.path.basename(file))
00176         if post_upload_command is not None:
00177             if verbose:
00178                 print "Running post-upload command: " + post_upload_command
00179             file_list = '"' + '" "'.join(remote_files) + '"'
00180             DoSSHCommand('%s "%s" %s' % (post_upload_command, path, file_list), user, host, port=port, ssh_key=ssh_key)
00181     finally:
00182         if upload_to_temp_dir:
00183             DoSSHCommand("rm -rf %s" % path, user, host, port=port,
00184                          ssh_key=ssh_key)
00185     if verbose:
00186         print "Upload complete"
00187 
00188 if __name__ == '__main__':
00189     host = RequireEnvironmentVariable('UPLOAD_HOST')
00190     user = RequireEnvironmentVariable('UPLOAD_USER')
00191     path = OptionalEnvironmentVariable('UPLOAD_PATH')
00192     upload_to_temp_dir = OptionalEnvironmentVariable('UPLOAD_TO_TEMP')
00193     port = OptionalEnvironmentVariable('UPLOAD_PORT')
00194     if port is not None:
00195         port = int(port)
00196     key = OptionalEnvironmentVariable('UPLOAD_SSH_KEY')
00197     post_upload_command = OptionalEnvironmentVariable('POST_UPLOAD_CMD')
00198     if (not path and not upload_to_temp_dir) or (path and upload_to_temp_dir):
00199         print "One (and only one of UPLOAD_PATH or UPLOAD_TO_TEMP must be " + \
00200               "defined."
00201         sys.exit(1)
00202     if sys.platform == 'win32':
00203         if path is not None:
00204             path = FixupMsysPath(path)
00205         if post_upload_command is not None:
00206             post_upload_command = FixupMsysPath(post_upload_command)
00207 
00208     parser = OptionParser(usage="usage: %prog [options] <files>")
00209     parser.add_option("-b", "--base-path",
00210                       action="store", dest="base_path",
00211                       help="Preserve file paths relative to this path when uploading. If unset, all files will be uploaded directly to UPLOAD_PATH.")
00212     (options, args) = parser.parse_args()
00213     if len(args) < 1:
00214         print "You must specify at least one file to upload"
00215         sys.exit(1)
00216     try:
00217         UploadFiles(user, host, path, args, base_path=options.base_path,
00218                     port=port, ssh_key=key, upload_to_temp_dir=upload_to_temp_dir,
00219                     post_upload_command=post_upload_command,
00220                     verbose=True)
00221     except IOError, (strerror):
00222         print strerror
00223         sys.exit(1)
00224     except Exception, (err):
00225         print err
00226         sys.exit(2)