From d67cc7fe1145113efe7a9980c4f96b614574ed5e Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Sat, 28 Jan 2023 16:20:01 -0500 Subject: [PATCH] Respect standard TMPDIR/TEMP/TMP environment variables If any standard temp directory environment variables are set, then don't use the output directory as the temp directory. Issue: #22 Signed-off-by: Andrew Gunnerson --- avbroot.py | 4 ++-- avbroot/util.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/avbroot.py b/avbroot.py index b691041..01ebbfc 100644 --- a/avbroot.py +++ b/avbroot.py @@ -211,7 +211,7 @@ def patch_subcommand(args): # Set default temp directory to the output directory because this is the # only way to control where external libraries put their temp files - tempfile.tempdir = os.path.dirname(os.path.abspath(output)) + util.set_default_temp_dir(os.path.dirname(os.path.abspath(output))) # Decrypt keys to temp directory with tempfile.TemporaryDirectory(dir=util.tmpfs_path()) as key_dir: @@ -257,7 +257,7 @@ def patch_subcommand(args): def extract_subcommand(args): # Set default temp directory to the output directory because this is the # only way to control where external libraries put their temp files - tempfile.tempdir = os.path.abspath(args.directory) + util.set_default_temp_dir(os.path.abspath(args.directory)) with zipfile.ZipFile(args.input, 'r') as z: info = z.getinfo(PATH_PAYLOAD) diff --git a/avbroot/util.py b/avbroot/util.py index d8ade89..9cfe523 100644 --- a/avbroot/util.py +++ b/avbroot/util.py @@ -57,6 +57,19 @@ def tmpfs_path(): return None +def set_default_temp_dir(path): + ''' + Set the default temporary directory for use if the user hasn't overridden + it with standard environment variables. + ''' + + if any(os.getenv(k) for k in ('TMPDIR', 'TEMP', 'TMP')): + # Rely on tempfile's default handling + return + + tempfile.tempdir = path + + @contextlib.contextmanager def open_output_file(path): '''