1 files changed, 93 insertions, 0 deletions
diff --git a/cloudinit/temp_utils.py b/cloudinit/temp_utils.py
new file mode 100644
index 00000000..0355f19d
--- /dev/null
+++ b/cloudinit/temp_utils.py
@@ -0,0 +1,93 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+import contextlib
+import errno
+import os
+import shutil
+import tempfile
+
+_TMPDIR = None
+_ROOT_TMPDIR = "/run/cloud-init/tmp"
+
+
+def _tempfile_dir_arg(odir=None):
+    """Return the proper 'dir' argument for tempfile functions.
+
+    When root, cloud-init will use /run/cloud-init/tmp to avoid
+    any cleaning that a distro boot might do on /tmp (such as
+    systemd-tmpfiles-clean).
+
+    If the caller of this function (mkdtemp or mkstemp) was provided
+    with a 'dir' argument, then that is respected.
+
+    @param odir: original 'dir' arg to 'mkdtemp' or other."""
+
+    if odir is not None:
+        return odir
+
+    global _TMPDIR
+    if _TMPDIR:
+        return _TMPDIR
+
+    if os.getuid() == 0:
+        tdir = _ROOT_TMPDIR
+    else:
+        tdir = os.environ.get('TMPDIR', '/tmp')
+    if not os.path.isdir(tdir):
+        os.makedirs(tdir)
+        os.chmod(tdir, 0o1777)
+
+    _TMPDIR = tdir
+    return tdir
+
+
+def ExtendedTemporaryFile(**kwargs):
+    kwargs['dir'] = _tempfile_dir_arg(kwargs.pop('dir', None))
+    fh = tempfile.NamedTemporaryFile(**kwargs)
+    # Replace its unlink with a quiet version
+    # that does not raise errors when the
+    # file to unlink has been unlinked elsewhere..
+
+    def _unlink_if_exists(path):
+        try:
+            os.unlink(path)
+        except OSError as e:
+            if e.errno != errno.ENOENT:
+                raise e
+
+    fh.unlink = _unlink_if_exists
+
+    # Add a new method that will unlink
+    # right 'now' but still lets the exit
+    # method attempt to remove it (which will
+    # not throw due to our del file being quiet
+    # about files that are not there)
+    def unlink_now():
+        fh.unlink(fh.name)
+
+    setattr(fh, 'unlink_now', unlink_now)
+    return fh
+
+
+@contextlib.contextmanager
+def tempdir(**kwargs):
+    # This seems like it was only added in python 3.2
+    # Make it since its useful...
+    # See: http://bugs.python.org/file12970/tempdir.patch
+    tdir = mkdtemp(**kwargs)
+    try:
+        yield tdir
+    finally:
+        shutil.rmtree(tdir)
+
+
+def mkdtemp(**kwargs):
+    kwargs['dir'] = _tempfile_dir_arg(kwargs.pop('dir', None))
+    return tempfile.mkdtemp(**kwargs)
+
+
+def mkstemp(**kwargs):
+    kwargs['dir'] = _tempfile_dir_arg(kwargs.pop('dir', None))
+    return tempfile.mkstemp(**kwargs)
+
+# vi: ts=4 expandtab