Simplify temporary file creation

Switch to using NamedTemporaryFile and converge the key and
certificate writes to a single function. They were identical
except for the data they were copying in, so we move that to a
higher-level function instead.

Author: sgallagher

src/sscg/__init__.py

@@ -17,52 +17,10 @@
 DEFAULT_KEY_STRENGTH = 2048  # 2048-bit encryption
 DEFAULT_HASH_ALG = "sha256"
 
-def write_certificate(options, cert, destination):
-    """
-    Write out the certificate to a temporary file first, then atomically copy
-    it to the destination path. This will avoid race-condition bugs with
-    checking for file presence and then writing to it. Note: this will clobber
-    the destination path.
-    """
-
-    # Create the temporary file in the same directory as the destination
-    # This ensures that we can atomically move it to the final name.
-
-    try:
-        (fd, fpath) = tempfile.mkstemp(dir=os.path.dirname(destination))
-        if options.debug:
-            print(_("Creating temporary certificate file at {}").format(fpath))
-        f = os.fdopen(fd, "w")
-
-        f.write(crypto.dump_certificate(options.cert_format, cert).decode("UTF-8"))
-        f.close()
-    except IOError:
-        # Something went wrong. Remove the temporary file before failing.
-        print(_("Could not write to {0}. Error: {1}").format(
-              fpath, sys.exc_info()[1]),
-              file=sys.stderr)
-        os.unlink(fpath)
-        raise
 
-    # Now atomically move the temporary file into place.
-    # We use os.rename because this is guaranteed to be atomic if it succeeds
-    # This operation can fail on some flavors of UNIX if the source and
-    # destination are on different filesystems, but this should not be the case.
-    try:
-        if options.debug:
-            print(_("Renaming {} to {}".format(fpath, destination)))
-        os.rename(fpath, destination)
-    except IOError:
-        # Something went wrong. Remove the temporary file before failing.
-        print(_("Could not rename to {0}. Error: {1}").format(
-              destination, sys.exc_info()[1]))
-        os.unlink(fpath)
-        raise
-
-
-def write_certificate_key(options, key, destination, cipher=None, passphrase=None):
+def write_secure_file(options, destination, data):
     """
-    Write out the certificate key to a temporary file first, then atomically
+    Write out the certificate or key to a temporary file first, then atomically
     copy it to the destination path. This will avoid race-condition bugs with
     checking for file presence and then writing to it. Note: this will clobber
     the destination path.
@@ -71,34 +29,26 @@ def write_certificate_key(options, key, destination, cipher=None, passphrase=Non
     # Create the temporary file in the same directory as the destination
     # This ensures that we can atomically move it to the final name.
 
+    f = tempfile.NamedTemporaryFile(dir=os.path.dirname(destination),
+                                    delete=False)
     try:
-        (fd, fpath) = tempfile.mkstemp(dir=os.path.dirname(destination))
-        if options.debug:
-            print(_("Creating temporary keyfile at {}").format(fpath))
-
-        f = os.fdopen(fd, "w")
-
-        f.write(crypto.dump_privatekey(options.cert_format, key, cipher, passphrase).decode("UTF-8"))
+        f.write(data)
+        f.flush()
+    except IOError as e:
         f.close()
-    except IOError:
-        # Something went wrong. Remove the temporary file before failing.
-        print(_("Could not write to {0}. Error: {1}").format(
-              fpath, sys.exc_info()[1]),
-              file=sys.stderr)
-        os.unlink(fpath)
-        raise
+        os.unlink(f.name)
+        raise Exception(_("Could not write to {0}. Error: {1}").format(f.name, e))
 
     # Now atomically move the temporary file into place.
     # We use os.rename because this is guaranteed to be atomic if it succeeds
     # This operation can fail on some flavors of UNIX if the source and
     # destination are on different filesystems, but this should not be the case.
+    if options.debug:
+        print(_("Renaming {} to {}").format(f.name, destination))
+
+    f.close()
     try:
-        if options.debug:
-            print(_("Renaming {} to {}").format(fpath, destination))
-        os.rename(fpath, destination)
-    except IOError:
-        # Something went wrong. Remove the temporary file before failing.
-        print(_("Could not rename to {0}. Error: {1}").format(
-              destination, sys.exc_info()[1]))
-        os.unlink(fpath)
-        raise
+        os.rename(f.name, destination)
+    except IOError as e:
+        os.unlink(f.name)
+        raise Exception(_("Could not rename to {0}. Error: {1}").format(destination, e))

src/sscg/main.py

@@ -6,8 +6,7 @@
 from OpenSSL import crypto
 from socket import gethostname
 from sscg import DEFAULT_CA_CERT, DEFAULT_KEY_STRENGTH, DEFAULT_LIFESPAN, \
-                 DEFAULT_CERT_FORMAT, DEFAULT_HASH_ALG, write_certificate,\
-    write_certificate_key
+                 DEFAULT_CERT_FORMAT, DEFAULT_HASH_ALG, write_secure_file
 from sscg.authority import create_temp_ca
 from sscg.service import create_service_cert
 
@@ -138,13 +137,19 @@ def main():
 
     try:
         # Write out the CA Certificate
-        write_certificate(options, ca_cert, options.ca_file)
+        write_secure_file(options,
+                          options.ca_file,
+                          crypto.dump_certificate(options.cert_format, ca_cert))
 
         # Write out the Service Certificate
-        write_certificate(options, svc_cert, options.cert_file)
+        write_secure_file(options,
+                          options.cert_file,
+                          crypto.dump_certificate(options.cert_format, svc_cert))
 
         # Write out the Service Private Key
-        write_certificate_key(options, svc_key, options.cert_key_file)
+        write_secure_file(options,
+                          options.cert_key_file,
+                          crypto.dump_privatekey(options.cert_format, svc_key))
     except:
         print (_("Error writing certificate files: {}").format(sys.exc_info()[1]),
                file=sys.stderr)