CVE-2011-2185

@@ -8,6 +8,7 @@
 import tempfile
 import re
 import os
+from StringIO import StringIO
 
 from fabric.api import *
 
@@ -65,9 +66,8 @@ def upload_template(filename, destination, context=None, use_jinja=False,
     directory by default, or from ``template_dir`` if given.
 
     The resulting rendered file will be uploaded to the remote file path
-    ``destination`` (which should include the desired remote filename.) If the
-    destination file already exists, it will be renamed with a ``.bak``
-    extension unless ``backup=False`` is specified.
+    ``destination``.  If the destination file already exists, it will be
+    renamed with a ``.bak`` extension unless ``backup=False`` is specified.
 
     By default, the file will be copied to ``destination`` as the logged-in
     user; specify ``use_sudo=True`` to use `sudo` instead.
@@ -79,15 +79,14 @@ def upload_template(filename, destination, context=None, use_jinja=False,
     Alternately, you may use the ``mode`` kwarg to specify an exact mode, in
     the same vein as ``os.chmod`` or the Unix ``chmod`` command.
     """
-    basename = os.path.basename(filename)
-
-    # This temporary file should not be automatically deleted on close, as we
-    # need it there to upload it (Windows locks the file for reading while
-    # open).
-    tempdir = tempfile.mkdtemp()
-    tempfile_name = os.path.join(tempdir, basename)
-    output = open(tempfile_name, "w+b")
-    # Init
+    func = use_sudo and sudo or run
+    # Normalize destination to be an actual filename, due to using StringIO
+    with settings(hide('everything'), warn_only=True):
+        if func('test -d %s' % destination).succeeded:
+            sep = "" if destination.endswith('/') else "/"
+            destination += sep + os.path.basename(filename)
+
+    # Process template
     text = None
     if use_jinja:
         try:
@@ -101,32 +100,19 @@ def upload_template(filename, destination, context=None, use_jinja=False,
             text = inputfile.read()
         if context:
             text = text % context
-    output.write(text)
-    output.close()
 
-    # Back up any original file 
-    func = use_sudo and sudo or run
-    # Back up any original file (need to do figure out ultimate destination)
-    if backup:
-        to_backup = destination
-        with settings(hide('everything'), warn_only=True):
-            # Is destination a directory?
-            if func('test -f %s' % to_backup).failed:
-                # If so, tack on the filename to get "real" destination
-                to_backup = destination + '/' + basename
-        if exists(to_backup):
-            func("cp %s %s.bak" % (to_backup, to_backup))
+    # Back up original file
+    if backup and exists(destination):
+        func("cp %s{,.bak}" % destination)
 
     # Upload the file.
     put(
-        local_path=tempfile_name,
+        local_path=StringIO(text),
         remote_path=destination,
         use_sudo=use_sudo,
         mirror_local_mode=mirror_local_mode,
         mode=mode
     )
-    output.close()
-    os.remove(tempfile_name)
 
 
 def sed(filename, before, after, limit='', use_sudo=False, backup='.bak',

@@ -1,16 +1,36 @@
 from __future__ import with_statement
 
+from fabric.api import hide, get, show
 from fabric.contrib.files import upload_template
 
-from utils import FabricTest
+from utils import FabricTest, eq_contents
 from server import server
 
+
 class TestContrib(FabricTest):
-    @server()
+    # Make sure it knows / is a directory.
+    # This is in lieu of starting down the "actual honest to god fake operating
+    # system" road...:(
+    @server(responses={'test -d /': ""})
     def test_upload_template_uses_correct_remote_filename(self):
         """
         upload_template() shouldn't munge final remote filename
         """
         template = self.mkfile('template.txt', 'text')
-        upload_template(template, '/')
-        assert self.exists_remotely('/template.txt')
+        with hide('everything'):
+            upload_template(template, '/')
+            assert self.exists_remotely('/template.txt')
+
+    @server()
+    def test_upload_template_handles_file_destination(self):
+        """
+        upload_template() should work OK with file and directory destinations
+        """
+        template = self.mkfile('template.txt', '%(varname)s')
+        local = self.path('result.txt')
+        remote = '/configfile.txt'
+        var = 'foobar'
+        with hide('everything'):
+            upload_template(template, remote, {'varname': var})
+            get(remote, local)
+        eq_contents(local, var)

@@ -3,12 +3,15 @@
 """
 
 from os import getcwd, sep
+import os.path
 from datetime import datetime
+from tempfile import mkdtemp
 
 from fabric.network import needs_host
 from fabric.operations import local, run, put
 from fabric.state import env, output
 
+__all__ = ['rsync_project', 'upload_project']
 
 @needs_host
 def rsync_project(remote_dir, local_dir=None, exclude=(), delete=False,
@@ -100,23 +103,44 @@ def rsync_project(remote_dir, local_dir=None, exclude=(), delete=False,
     return local(cmd)
 
 
-def upload_project():
+def upload_project(local_dir=None, remote_dir=""):
     """
     Upload the current project to a remote system, tar/gzipping during the move.
 
-    This function makes use of the ``/tmp/`` directory and the ``tar`` and
-    ``gzip`` programs/libraries; thus it will not work too well on Win32
-    systems unless one is using Cygwin or something similar.
+    This function makes use of the ``tar`` and ``gzip`` programs/libraries, thus
+    it will not work too well on Win32 systems unless one is using Cygwin or
+    something similar.
+
+    ``upload_project`` will attempt to clean up the local and remote tarfiles
+    when it finishes executing, even in the event of a failure.
+
+    :param local_dir: default current working directory, the project folder to
+        upload
 
-    ``upload_project`` will attempt to clean up the tarfiles when it finishes
-    executing.
     """
-    tar_file = "/tmp/fab.%s.tar" % datetime.utcnow().strftime(
-        '%Y_%m_%d_%H-%M-%S')
-    cwd_name = getcwd().split(sep)[-1]
-    tgz_name = cwd_name + ".tar.gz"
-    local("tar -czf %s ." % tar_file)
-    put(tar_file, cwd_name + ".tar.gz")
-    local("rm -f " + tar_file)
-    run("tar -xzf " + tgz_name)
-    run("rm -f " + tgz_name)
+    if not local_dir:
+        local_dir = os.getcwd()
+    else:
+        # Remove final '/' in local_dir so that basename() works
+        local_dir = local_dir[:-1] if local_dir[-1] == os.sep else local_dir
+
+    local_path, local_name = os.path.split(local_dir)
+
+    tar_file = "%s.tar.gz" % local_name
+    target_tar = os.path.join(remote_dir, tar_file)
+
+    tmp_folder = mkdtemp()
+    try:
+        tar_path = os.path.join(tmp_folder, tar_file)
+        local("tar -czf %s -C %s %s" % (tar_path, local_path, local_name))
+        put(tar_path, target_tar)
+        
+        try:
+            run("tar -xzf %s" % tar_file)
+
+        finally:
+            run("rm -f %s" % tar_file)
+
+    finally:
+        local("rm -rf %s" % tmp_folder)
+

@@ -0,0 +1,161 @@
+import unittest
+import os
+
+import fudge
+from fudge.inspector import arg
+
+from fabric.contrib import project
+
+
+class UploadProjectTestCase(unittest.TestCase):
+    """Test case for :func: `fabric.contrib.project.upload_project`."""
+
+    fake_tmp = "testtempfolder"
+
+
+    def setUp(self):
+        fudge.clear_expectations()
+
+        # We need to mock out run, local, and put
+
+        self.fake_run = fudge.Fake('project.run', callable=True)
+        self.patched_run = fudge.patch_object(
+                               project,
+                               'run',
+                               self.fake_run
+                           )
+
+        self.fake_local = fudge.Fake('local', callable=True)
+        self.patched_local = fudge.patch_object(
+                                 project,
+                                 'local',
+                                 self.fake_local
+                             )
+
+        self.fake_put = fudge.Fake('put', callable=True)
+        self.patched_put = fudge.patch_object(
+                               project,
+                               'put',
+                               self.fake_put
+                           )
+
+        # We don't want to create temp folders
+        self.fake_mkdtemp = fudge.Fake(
+                                'mkdtemp',
+                                expect_call=True
+                            ).returns(self.fake_tmp)
+        self.patched_mkdtemp = fudge.patch_object(
+                                   project,
+                                   'mkdtemp',
+                                   self.fake_mkdtemp
+                               )
+
+
+    def tearDown(self):
+        self.patched_run.restore()
+        self.patched_local.restore()
+        self.patched_put.restore()
+
+        fudge.clear_expectations()
+
+
+    @fudge.with_fakes
+    def test_temp_folder_is_used(self):
+        """A unique temp folder is used for creating the archive to upload."""
+
+        # Exercise
+        project.upload_project()
+
+
+    @fudge.with_fakes
+    def test_project_is_archived_locally(self):
+        """The project should be archived locally before being uploaded."""
+
+        # local() is called more than once so we need an extra next_call()
+        # otherwise fudge compares the args to the last call to local()
+        self.fake_local.with_args(arg.startswith("tar -czf")).next_call()
+
+        # Exercise
+        project.upload_project()
+
+
+    @fudge.with_fakes
+    def test_current_directory_is_uploaded_by_default(self):
+        """By default the project uploaded is the current working directory."""
+
+        cwd_path, cwd_name = os.path.split(os.getcwd())
+
+        # local() is called more than once so we need an extra next_call()
+        # otherwise fudge compares the args to the last call to local()
+        self.fake_local.with_args(
+            arg.endswith("-C %s %s" % (cwd_path, cwd_name))
+        ).next_call()
+
+        # Exercise
+        project.upload_project()
+
+
+    @fudge.with_fakes
+    def test_path_to_local_project_can_be_specified(self):
+        """It should be possible to specify which local folder to upload."""
+
+        project_path = "path/to/my/project"
+
+        # local() is called more than once so we need an extra next_call()
+        # otherwise fudge compares the args to the last call to local()
+        self.fake_local.with_args(
+            arg.endswith("-C %s %s" % os.path.split(project_path))
+        ).next_call()
+
+        # Exercise
+        project.upload_project(local_dir=project_path)
+
+
+    @fudge.with_fakes
+    def test_path_to_local_project_can_end_in_separator(self):
+        """A local path ending in a separator should be handled correctly."""
+
+        project_path = "path/to/my"
+        base = "project"
+
+        # local() is called more than once so we need an extra next_call()
+        # otherwise fudge compares the args to the last call to local()
+        self.fake_local.with_args(
+            arg.endswith("-C %s %s" % (project_path, base))
+        ).next_call()
+
+        # Exercise
+        project.upload_project(local_dir="%s/%s/" % (project_path, base))
+
+
+    @fudge.with_fakes
+    def test_default_remote_folder_is_home(self):
+        """Project is uploaded to remote home by default."""
+
+        local_dir = "folder"
+
+        # local() is called more than once so we need an extra next_call()
+        # otherwise fudge compares the args to the last call to local()
+        self.fake_put.with_args(
+            "%s/folder.tar.gz" % self.fake_tmp, "folder.tar.gz"
+        ).next_call()
+
+        # Exercise
+        project.upload_project(local_dir=local_dir)
+
+    @fudge.with_fakes
+    def test_path_to_remote_folder_can_be_specified(self):
+        """It should be possible to specify which local folder to upload to."""
+
+        local_dir = "folder"
+        remote_path = "path/to/remote/folder"
+
+        # local() is called more than once so we need an extra next_call()
+        # otherwise fudge compares the args to the last call to local()
+        self.fake_put.with_args(
+            "%s/folder.tar.gz" % self.fake_tmp, "%s/folder.tar.gz" % remote_path
+        ).next_call()
+
+        # Exercise
+        project.upload_project(local_dir=local_dir, remote_dir=remote_path)
+