From e6cbdef4ab51a3db8a198c05f171f9dbb9462b43 Mon Sep 17 00:00:00 2001
From: Mark Grassi <markgras@buffalo.edu>
Date: Tue, 19 Jul 2022 20:32:03 -0400
Subject: [PATCH] Implement code review comments.

---
 utils/apparmor/aa.py       | 11 +++++++++--
 utils/apparmor/easyprof.py |  3 ++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/utils/apparmor/aa.py b/utils/apparmor/aa.py
index a5c70f3ca..15111f50d 100644
--- a/utils/apparmor/aa.py
+++ b/utils/apparmor/aa.py
@@ -2164,8 +2164,15 @@ def write_profile(profile, is_attachment=False):
 
     serialize_options = {'METADATA': True, 'is_attachment': is_attachment}
     profile_string = serialize_profile(split_to_merged(aa), profile, serialize_options)
-    with open(prof_filename, 'w') as file:  # Keep this context simple: don't want to corrupt file.
-        file.write(profile_string)
+    with NamedTemporaryFile('w', suffix='~', delete=False, dir=profile_dir) as newprof:
+        if os.path.exists(prof_filename):
+            shutil.copymode(prof_filename, newprof.name)
+        else:
+            # permission_600 = stat.S_IRUSR | stat.S_IWUSR    # Owner read and write
+            # os.chmod(newprof.name, permission_600)
+            pass
+        newprof.write(profile_string)
+    os.rename(newprof.name, prof_filename)
 
     if profile in changed:
         changed.pop(profile)
diff --git a/utils/apparmor/easyprof.py b/utils/apparmor/easyprof.py
index a103e0309..d27dbdfc3 100644
--- a/utils/apparmor/easyprof.py
+++ b/utils/apparmor/easyprof.py
@@ -677,8 +677,9 @@ class AppArmorEasyProfile:
 
             if not isinstance(policy, bytes):
                 policy = policy.encode('utf-8')
-            with open(out_fn, "wb") as f:
+            with NamedTemporaryFile('wb', prefix='aa-easyprof', suffix='~', delete=False) as f:
                 f.write(policy)
+            os.rename(f.name, out_fn)
 
     def gen_manifest(self, params):
         '''Take params list and output a JSON file'''