Add utils.tempdir() context manager for easy temp dirs

Fixes bug 883323 (and others)

Users of tempfile.mkdtemp() need to make sure the directory is cleaned
up when it's done being used. Unfortunately, not all of the code does
so at all, or safely (by using a try/finally block).

Change-Id: I270109d83efec4f8b3dd954021493f4d96c6ab79

nova/auth/manager.py

@@ -24,9 +24,7 @@
 """
 
 import os
-import shutil
 import string  # pylint: disable=W0402
-import tempfile
 import uuid
 import zipfile
 
@@ -767,45 +765,44 @@ def get_credentials(self, user, project=None, use_dmz=True):
         pid = Project.safe_id(project)
         private_key, signed_cert = crypto.generate_x509_cert(user.id, pid)
 
-        tmpdir = tempfile.mkdtemp()
-        zf = os.path.join(tmpdir, "temp.zip")
-        zippy = zipfile.ZipFile(zf, 'w')
-        if use_dmz and FLAGS.region_list:
-            regions = {}
-            for item in FLAGS.region_list:
-                region, _sep, region_host = item.partition("=")
-                regions[region] = region_host
-        else:
-            regions = {'nova': FLAGS.ec2_host}
-        for region, host in regions.iteritems():
-            rc = self.__generate_rc(user,
-                                    pid,
-                                    use_dmz,
-                                    host)
-            zippy.writestr(FLAGS.credential_rc_file % region, rc)
-
-        zippy.writestr(FLAGS.credential_key_file, private_key)
-        zippy.writestr(FLAGS.credential_cert_file, signed_cert)
-
-        (vpn_ip, vpn_port) = self.get_project_vpn_data(project)
-        if vpn_ip:
-            configfile = open(FLAGS.vpn_client_template, "r")
-            s = string.Template(configfile.read())
-            configfile.close()
-            config = s.substitute(keyfile=FLAGS.credential_key_file,
-                                  certfile=FLAGS.credential_cert_file,
-                                  ip=vpn_ip,
-                                  port=vpn_port)
-            zippy.writestr(FLAGS.credential_vpn_file, config)
-        else:
-            LOG.warn(_("No vpn data for project %s"), pid)
-
-        zippy.writestr(FLAGS.ca_file, crypto.fetch_ca(pid))
-        zippy.close()
-        with open(zf, 'rb') as f:
-            read_buffer = f.read()
+        with utils.tempdir() as tmpdir:
+            zf = os.path.join(tmpdir, "temp.zip")
+            zippy = zipfile.ZipFile(zf, 'w')
+            if use_dmz and FLAGS.region_list:
+                regions = {}
+                for item in FLAGS.region_list:
+                    region, _sep, region_host = item.partition("=")
+                    regions[region] = region_host
+            else:
+                regions = {'nova': FLAGS.ec2_host}
+            for region, host in regions.iteritems():
+                rc = self.__generate_rc(user,
+                                        pid,
+                                        use_dmz,
+                                        host)
+                zippy.writestr(FLAGS.credential_rc_file % region, rc)
+
+            zippy.writestr(FLAGS.credential_key_file, private_key)
+            zippy.writestr(FLAGS.credential_cert_file, signed_cert)
+
+            (vpn_ip, vpn_port) = self.get_project_vpn_data(project)
+            if vpn_ip:
+                configfile = open(FLAGS.vpn_client_template, "r")
+                s = string.Template(configfile.read())
+                configfile.close()
+                config = s.substitute(keyfile=FLAGS.credential_key_file,
+                                      certfile=FLAGS.credential_cert_file,
+                                      ip=vpn_ip,
+                                      port=vpn_port)
+                zippy.writestr(FLAGS.credential_vpn_file, config)
+            else:
+                LOG.warn(_("No vpn data for project %s"), pid)
+
+            zippy.writestr(FLAGS.ca_file, crypto.fetch_ca(pid))
+            zippy.close()
+            with open(zf, 'rb') as f:
+                read_buffer = f.read()
 
-        shutil.rmtree(tmpdir)
         return read_buffer
 
     def get_environment_rc(self, user, project=None, use_dmz=True):

nova/cloudpipe/pipelib.py

@@ -65,36 +65,39 @@ def __init__(self):
 
     def get_encoded_zip(self, project_id):
         # Make a payload.zip
-        tmpfolder = tempfile.mkdtemp()
-        filename = "payload.zip"
-        zippath = os.path.join(tmpfolder, filename)
-        z = zipfile.ZipFile(zippath, "w", zipfile.ZIP_DEFLATED)
-        shellfile = open(FLAGS.boot_script_template, "r")
-        s = string.Template(shellfile.read())
-        shellfile.close()
-        boot_script = s.substitute(cc_dmz=FLAGS.ec2_dmz_host,
-                                   cc_port=FLAGS.ec2_port,
-                                   dmz_net=FLAGS.dmz_net,
-                                   dmz_mask=FLAGS.dmz_mask,
-                                   num_vpn=FLAGS.cnt_vpn_clients)
-        # genvpn, sign csr
-        crypto.generate_vpn_files(project_id)
-        z.writestr('autorun.sh', boot_script)
-        crl = os.path.join(crypto.ca_folder(project_id), 'crl.pem')
-        z.write(crl, 'crl.pem')
-        server_key = os.path.join(crypto.ca_folder(project_id), 'server.key')
-        z.write(server_key, 'server.key')
-        ca_crt = os.path.join(crypto.ca_path(project_id))
-        z.write(ca_crt, 'ca.crt')
-        server_crt = os.path.join(crypto.ca_folder(project_id), 'server.crt')
-        z.write(server_crt, 'server.crt')
-        z.close()
-        zippy = open(zippath, "r")
-        # NOTE(vish): run instances expects encoded userdata, it is decoded
-        # in the get_metadata_call. autorun.sh also decodes the zip file,
-        # hence the double encoding.
-        encoded = zippy.read().encode("base64").encode("base64")
-        zippy.close()
+        with utils.tempdir() as tmpdir:
+            filename = "payload.zip"
+            zippath = os.path.join(tmpdir, filename)
+            z = zipfile.ZipFile(zippath, "w", zipfile.ZIP_DEFLATED)
+            shellfile = open(FLAGS.boot_script_template, "r")
+            s = string.Template(shellfile.read())
+            shellfile.close()
+            boot_script = s.substitute(cc_dmz=FLAGS.ec2_dmz_host,
+                                       cc_port=FLAGS.ec2_port,
+                                       dmz_net=FLAGS.dmz_net,
+                                       dmz_mask=FLAGS.dmz_mask,
+                                       num_vpn=FLAGS.cnt_vpn_clients)
+            # genvpn, sign csr
+            crypto.generate_vpn_files(project_id)
+            z.writestr('autorun.sh', boot_script)
+            crl = os.path.join(crypto.ca_folder(project_id), 'crl.pem')
+            z.write(crl, 'crl.pem')
+            server_key = os.path.join(crypto.ca_folder(project_id),
+                                      'server.key')
+            z.write(server_key, 'server.key')
+            ca_crt = os.path.join(crypto.ca_path(project_id))
+            z.write(ca_crt, 'ca.crt')
+            server_crt = os.path.join(crypto.ca_folder(project_id),
+                                      'server.crt')
+            z.write(server_crt, 'server.crt')
+            z.close()
+            zippy = open(zippath, "r")
+            # NOTE(vish): run instances expects encoded userdata, it is decoded
+            # in the get_metadata_call. autorun.sh also decodes the zip file,
+            # hence the double encoding.
+            encoded = zippy.read().encode("base64").encode("base64")
+            zippy.close()
+
         return encoded
 
     def launch_vpn_instance(self, project_id, user_id):

nova/compat/flagfile.py

@@ -175,6 +175,8 @@ def handle_flagfiles_managed(args):
              # Do stuff
         # Any temporary fils have been removed
     '''
+    # NOTE(johannes): Would be nice to use utils.tempdir(), but it
+    # causes an import loop
     tempdir = tempfile.mkdtemp(prefix='nova-conf-')
     try:
         yield handle_flagfiles(args, tempdir=tempdir)

nova/crypto.py

@@ -27,9 +27,7 @@
 import base64
 import hashlib
 import os
-import shutil
 import string
-import tempfile
 
 import Crypto.Cipher.AES
 
@@ -127,36 +125,26 @@ def _generate_fingerprint(public_key_file):
 
 
 def generate_fingerprint(public_key):
-    tmpdir = tempfile.mkdtemp()
-    try:
-        pubfile = os.path.join(tmpdir, 'temp.pub')
-        with open(pubfile, 'w') as f:
-            f.write(public_key)
-        return _generate_fingerprint(pubfile)
-    except exception.ProcessExecutionError:
-        raise exception.InvalidKeypair()
-    finally:
+    with utils.tempdir() as tmpdir:
         try:
-            shutil.rmtree(tmpdir)
-        except IOError, e:
-            LOG.debug(_('Could not remove tmpdir: %s'), str(e))
+            pubfile = os.path.join(tmpdir, 'temp.pub')
+            with open(pubfile, 'w') as f:
+                f.write(public_key)
+            return _generate_fingerprint(pubfile)
+        except exception.ProcessExecutionError:
+            raise exception.InvalidKeypair()
 
 
 def generate_key_pair(bits=1024):
     # what is the magic 65537?
 
-    tmpdir = tempfile.mkdtemp()
-    keyfile = os.path.join(tmpdir, 'temp')
-    utils.execute('ssh-keygen', '-q', '-b', bits, '-N', '',
-                  '-t', 'rsa', '-f', keyfile)
-    fingerprint = _generate_fingerprint('%s.pub' % (keyfile))
-    private_key = open(keyfile).read()
-    public_key = open(keyfile + '.pub').read()
-
-    try:
-        shutil.rmtree(tmpdir)
-    except OSError, e:
-        LOG.debug(_('Could not remove tmpdir: %s'), str(e))
+    with utils.tempdir() as tmpdir:
+        keyfile = os.path.join(tmpdir, 'temp')
+        utils.execute('ssh-keygen', '-q', '-b', bits, '-N', '',
+                      '-t', 'rsa', '-f', keyfile)
+        fingerprint = _generate_fingerprint('%s.pub' % (keyfile))
+        private_key = open(keyfile).read()
+        public_key = open(keyfile + '.pub').read()
 
     return (private_key, public_key, fingerprint)
 
@@ -233,19 +221,15 @@ def _user_cert_subject(user_id, project_id):
 def generate_x509_cert(user_id, project_id, bits=1024):
     """Generate and sign a cert for user in project."""
     subject = _user_cert_subject(user_id, project_id)
-    tmpdir = tempfile.mkdtemp()
-    keyfile = os.path.abspath(os.path.join(tmpdir, 'temp.key'))
-    csrfile = os.path.join(tmpdir, 'temp.csr')
-    utils.execute('openssl', 'genrsa', '-out', keyfile, str(bits))
-    utils.execute('openssl', 'req', '-new', '-key', keyfile, '-out', csrfile,
-                  '-batch', '-subj', subject)
-    private_key = open(keyfile).read()
-    csr = open(csrfile).read()
 
-    try:
-        shutil.rmtree(tmpdir)
-    except OSError, e:
-        LOG.debug(_('Could not remove tmpdir: %s'), str(e))
+    with utils.tempdir() as tmpdir:
+        keyfile = os.path.abspath(os.path.join(tmpdir, 'temp.key'))
+        csrfile = os.path.join(tmpdir, 'temp.csr')
+        utils.execute('openssl', 'genrsa', '-out', keyfile, str(bits))
+        utils.execute('openssl', 'req', '-new', '-key', keyfile, '-out',
+                      csrfile, '-batch', '-subj', subject)
+        private_key = open(keyfile).read()
+        csr = open(csrfile).read()
 
     (serial, signed_csr) = sign_csr(csr, project_id)
     fname = os.path.join(ca_folder(project_id), 'newcerts/%s.pem' % serial)
@@ -298,26 +282,30 @@ def sign_csr(csr_text, project_id=None):
 
 
 def _sign_csr(csr_text, ca_folder):
-    tmpfolder = tempfile.mkdtemp()
-    inbound = os.path.join(tmpfolder, 'inbound.csr')
-    outbound = os.path.join(tmpfolder, 'outbound.csr')
-    csrfile = open(inbound, 'w')
-    csrfile.write(csr_text)
-    csrfile.close()
-    LOG.debug(_('Flags path: %s'), ca_folder)
-    start = os.getcwd()
-    # Change working dir to CA
-    if not os.path.exists(ca_folder):
-        os.makedirs(ca_folder)
-    os.chdir(ca_folder)
-    utils.execute('openssl', 'ca', '-batch', '-out', outbound, '-config',
-                  './openssl.cnf', '-infiles', inbound)
-    out, _err = utils.execute('openssl', 'x509', '-in', outbound,
-                              '-serial', '-noout')
-    serial = string.strip(out.rpartition('=')[2])
-    os.chdir(start)
-    with open(outbound, 'r') as crtfile:
-        return (serial, crtfile.read())
+    with utils.tempdir() as tmpdir:
+        inbound = os.path.join(tmpdir, 'inbound.csr')
+        outbound = os.path.join(tmpdir, 'outbound.csr')
+
+        with open(inbound, 'w') as csrfile:
+            csrfile.write(csr_text)
+
+        LOG.debug(_('Flags path: %s'), ca_folder)
+        start = os.getcwd()
+
+        # Change working dir to CA
+        if not os.path.exists(ca_folder):
+            os.makedirs(ca_folder)
+
+        os.chdir(ca_folder)
+        utils.execute('openssl', 'ca', '-batch', '-out', outbound, '-config',
+                      './openssl.cnf', '-infiles', inbound)
+        out, _err = utils.execute('openssl', 'x509', '-in', outbound,
+                                  '-serial', '-noout')
+        serial = string.strip(out.rpartition('=')[2])
+        os.chdir(start)
+
+        with open(outbound, 'r') as crtfile:
+            return (serial, crtfile.read())
 
 
 def _build_cipher(key, iv):

nova/tests/test_crypto.py

@@ -17,8 +17,6 @@
 """
 
 import os
-import shutil
-import tempfile
 
 import mox
 
@@ -50,9 +48,8 @@ def test_encrypt_decrypt(self):
 
 class X509Test(test.TestCase):
     def test_can_generate_x509(self):
-        tmpdir = tempfile.mkdtemp()
-        self.flags(ca_path=tmpdir)
-        try:
+        with utils.tempdir() as tmpdir:
+            self.flags(ca_path=tmpdir)
             crypto.ensure_ca_filesystem()
             _key, cert_str = crypto.generate_x509_cert('fake', 'fake')
 
@@ -70,14 +67,10 @@ def test_can_generate_x509(self):
                     project_cert_file, '-verbose', signed_cert_file)
             self.assertFalse(err)
 
-        finally:
-            shutil.rmtree(tmpdir)
-
     def test_encrypt_decrypt_x509(self):
-        tmpdir = tempfile.mkdtemp()
-        self.flags(ca_path=tmpdir)
-        project_id = "fake"
-        try:
+        with utils.tempdir() as tmpdir:
+            self.flags(ca_path=tmpdir)
+            project_id = "fake"
             crypto.ensure_ca_filesystem()
             cert = crypto.fetch_ca(project_id)
             public_key = os.path.join(tmpdir, "public.pem")
@@ -92,8 +85,6 @@ def test_encrypt_decrypt_x509(self):
                                      process_input=text)
             dec = crypto.decrypt_text(project_id, enc)
             self.assertEqual(text, dec)
-        finally:
-            shutil.rmtree(tmpdir)
 
 
 class RevokeCertsTest(test.TestCase):