From c6f9ccaa25d73461cb502e43b1b0dd2b9e98909f Mon Sep 17 00:00:00 2001
From: Bartosz Myrcha <bartosz.myrcha@intel.com>
Date: Thu, 1 Feb 2024 22:55:11 -0800
Subject: [PATCH] Store token in protected file (#1585)

Signed-off-by: bmyrcha <bartosz.myrcha@intel.com>
---
 neural_insights/web/configuration.py | 75 ++++++++++++++++++----------
 1 file changed, 49 insertions(+), 26 deletions(-)

diff --git a/neural_insights/web/configuration.py b/neural_insights/web/configuration.py
index fabb92f92bf..7279dd4ddd3 100644
--- a/neural_insights/web/configuration.py
+++ b/neural_insights/web/configuration.py
@@ -226,36 +226,59 @@ def dump_token_to_file(self) -> None:
         """Dump token to file."""
         token_filepath = os.path.join(WORKDIR_LOCATION, "token")
         os.makedirs(os.path.dirname(token_filepath), exist_ok=True)
-        with open(token_filepath, "w") as token_file:
-            token_file.write(self.token)
 
         if sys.platform == "win32":
-            import ntsecuritycon as con  # pylint: disable=import-error
-            import win32api  # pylint: disable=import-error
-            import win32security  # pylint: disable=import-error
-
-            user, _, _ = win32security.LookupAccountName("", win32api.GetUserName())
-            security_descriptor = win32security.GetFileSecurity(
-                token_filepath,
-                win32security.DACL_SECURITY_INFORMATION,
-            )
-            dacl = win32security.ACL()
-            dacl.AddAccessAllowedAce(
-                win32security.ACL_REVISION,
-                con.FILE_GENERIC_READ | con.FILE_GENERIC_WRITE,
-                user,
-            )
-            security_descriptor.SetSecurityDescriptorDacl(1, dacl, 0)
-            win32security.SetFileSecurity(
-                token_filepath,
-                win32security.DACL_SECURITY_INFORMATION,
-                security_descriptor,
-            )
-        else:
-            os.chown(token_filepath, uid=os.geteuid(), gid=os.getgid())
-            os.chmod(token_filepath, 0o600)
+            self.create_secured_token_file_win(token_filepath)
+
+        try:
+            token_file = os.open(token_filepath, flags=os.O_WRONLY | os.O_CREAT, mode=0o600)
+            os.write(token_file, self.token.encode())
+        except Exception as err:
+            raise err
+        finally:
+            os.close(token_file)
+
         log.debug(f"Token has been dumped to {token_filepath}.")
 
+    @staticmethod
+    def create_secured_token_file_win(token_filepath: str):
+        """Create secured file on Windows OS."""
+        import ntsecuritycon as con  # pylint: disable=import-error
+        import win32api  # pylint: disable=import-error
+        import win32file  # pylint: disable=import-error
+        import win32security  # pylint: disable=import-error
+
+        username = win32api.GetUserName()
+        os.makedirs(os.path.dirname(token_filepath), exist_ok=True)
+
+        if os.path.exists(token_filepath):
+            os.remove(token_filepath)
+
+        security_descriptor = win32security.SECURITY_DESCRIPTOR()
+        user_sid, _, _ = win32security.LookupAccountName("", username)
+
+        access_rights = con.FILE_ALL_ACCESS
+
+        dacl = win32security.ACL()
+        dacl.AddAccessAllowedAce(win32security.ACL_REVISION, access_rights, user_sid)
+
+        security_descriptor.SetSecurityDescriptorDacl(1, dacl, 0)
+
+        security_attributes = win32security.SECURITY_ATTRIBUTES()
+        security_attributes.SECURITY_DESCRIPTOR = security_descriptor
+
+        handle = win32file.CreateFile(
+            token_filepath,
+            win32file.GENERIC_WRITE,
+            win32file.FILE_SHARE_READ,
+            security_attributes,
+            win32file.CREATE_NEW,
+            win32file.FILE_ATTRIBUTE_NORMAL,
+            None,
+        )
+
+        win32file.CloseHandle(handle)
+
     def _ensure_valid_port(self, port: int) -> None:
         """Validate if proposed port number is allowed by TCP/IP."""
         if port < 1: