Update refresh_entry_token to use context managers

HEPCloud · Sep 15, 2022 · 80d9041 · 80d9041
1 parent 62a51f3
commit 80d9041
Showing 1 changed file with 5 additions and 10 deletions.
diff --git a/src/decisionengine_modules/glideinwms/glide_frontend_element.py b/src/decisionengine_modules/glideinwms/glide_frontend_element.py
@@ -2,7 +2,6 @@
 import structlog
 import math
 import os
-import shutil
 import socket
 import sys
 import tempfile
@@ -1305,7 +1304,6 @@ def refresh_entry_token(self, glidein_site, work_dir="/var/lib/gwms-frontend"):
         """
         tkn_file = ""
         tkn_str = ""
-        tmpnm = ""
 
         try:
             tkn_dir = os.path.join(work_dir, "tokens.d")
@@ -1332,7 +1330,6 @@ def refresh_entry_token(self, glidein_site, work_dir="/var/lib/gwms-frontend"):
                 tkn_age = time.time() - os.stat(tkn_file).st_mtime
             if tkn_age > one_hr and os.path.exists(pwd_file):
                 # TODO: scope, duration, identity  should be configurable
-                (fd, tmpnm) = tempfile.mkstemp()
                 scope = "condor:/READ condor:/ADVERTISE_STARTD condor:/ADVERTISE_MASTER"
                 duration = 24 * one_hr
                 identity = f"{glidein_site}@{socket.gethostname()}"
@@ -1343,10 +1340,11 @@ def refresh_entry_token(self, glidein_site, work_dir="/var/lib/gwms-frontend"):
                 self.logger.debug("identity= %s" % identity)
                 tkn_str = token_util.create_and_sign_token(pwd_file, scope=scope, duration=duration, identity=identity)
                 self.logger.debug("tkn_str= %s" % tkn_str)
-                os.write(fd, tkn_str)
-                os.close(fd)
-                shutil.move(tmpnm, tkn_file)
-                os.chmod(tkn_file, 0o600)
+                with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=tkn_dir) as fd:
+                    os.chmod(fd.name, 0o600)
+                    fd.write(tkn_str)
+                    fd.flush()
+                    os.replace(fd.name, tkn_file)
                 self.logger.debug("created token %s" % tkn_file)
             elif os.path.exists(tkn_file):
                 with open(tkn_file) as fbuf:
@@ -1356,9 +1354,6 @@ def refresh_entry_token(self, glidein_site, work_dir="/var/lib/gwms-frontend"):
             self.logger.warning(f"failed to create {tkn_file}: {err}")
             for i in sys.exc_info():
                 self.logger.warning("%s" % i)
-        finally:
-            if os.path.exists(tmpnm):
-                os.remove(tmpnm)
 
         return tkn_str