#12

Handle secure desktop.
sirAndros · Feb 18, 2019 · efe4725 · efe4725
1 parent 7431de9
commit efe4725
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 14 deletions.
diff --git a/KeePassWinHelloExt.cs b/KeePassWinHelloExt.cs
@@ -53,7 +53,7 @@ private void OnWindowAdded(object sender, GwmWindowEventArgs e)
             var keyPromptForm = e.Form as KeyPromptForm;
             if (keyPromptForm != null)
             {
-                _keyManager.OnKeyPrompt(keyPromptForm);
+                _keyManager.OnKeyPrompt(keyPromptForm, _host.MainWindow);
                 return;
             }
 

diff --git a/KeyManagement/KeyManager.cs b/KeyManagement/KeyManager.cs
@@ -2,6 +2,8 @@
 using System.Diagnostics;
 using System.Linq;
 using System.Reflection;
+using System.Threading;
+using System.Threading.Tasks;
 using System.Windows.Forms;
 using KeePass.Forms;
 using KeePassLib.Keys;
@@ -13,31 +15,54 @@ class KeyManager
     {
         private readonly KeyCipher _keyCipher;
         private readonly KeyStorage _keyStorage;
+        private bool _isSecureDesktopSettingChanged = false;
 
         public KeyManager(IntPtr windowHandle)
         {
             _keyStorage = new KeyStorage();
             _keyCipher = new KeyCipher(Settings.ConfirmationMessage, windowHandle);
         }
 
-        public void OnKeyPrompt(KeyPromptForm keyPromptForm)
+        public void OnKeyPrompt(KeyPromptForm keyPromptForm, MainForm mainWindow)
         {
-            if (keyPromptForm.SecureDesktopMode)
-                return;
-
             if (!Settings.Instance.Enabled)
                 return;
 
-            CompositeKey compositeKey;
-            if (ExtractCompositeKey(GetDbPath(keyPromptForm), out compositeKey))
+            string dbPath = GetDbPath(keyPromptForm);
+            if (keyPromptForm.SecureDesktopMode)
             {
-                SetCompositeKey(keyPromptForm, compositeKey);
-                // Remove flushing
-                keyPromptForm.Visible = false;
-                keyPromptForm.Opacity = 0;
-
-                keyPromptForm.DialogResult = DialogResult.OK;
-                keyPromptForm.Close();
+                if (IsKeyForDataBaseExist(dbPath))
+                {
+                    var dbFile = GetIoInfo(keyPromptForm);
+                    CloseFormWithResult(keyPromptForm, DialogResult.Cancel);
+                    Task.Factory.StartNew(() =>
+                    {
+                        KeePass.Program.Config.Security.MasterKeyOnSecureDesktop = false;
+                        _isSecureDesktopSettingChanged = true;
+                        Thread.Yield();
+                        ReOpenKeyPromptForm(mainWindow, dbFile);
+                    })
+                    .ContinueWith(_ =>
+                    {
+                        KeePass.Program.Config.Security.MasterKeyOnSecureDesktop = true;
+                        _isSecureDesktopSettingChanged = false;
+                    });
+                }
+            }
+            else
+            {
+                CompositeKey compositeKey;
+                if (ExtractCompositeKey(dbPath, out compositeKey))
+                {
+                    SetCompositeKey(keyPromptForm, compositeKey);
+                    CloseFormWithResult(keyPromptForm, DialogResult.OK);
+                }
+                else if (_isSecureDesktopSettingChanged)
+                {
+                    var dbFile = GetIoInfo(keyPromptForm);
+                    CloseFormWithResult(keyPromptForm, DialogResult.Cancel);
+                    Task.Factory.StartNew(() => ReOpenKeyPromptForm(mainWindow, dbFile));
+                }
             }
         }
 
@@ -68,6 +93,28 @@ public void OnDBClosing(object sender, FileClosingEventArgs e)
             }
         }
 
+        private static void CloseFormWithResult(KeyPromptForm keyPromptForm, DialogResult result)
+        {
+            // Remove flushing
+            keyPromptForm.Visible = false;
+            keyPromptForm.Opacity = 0;
+
+            keyPromptForm.DialogResult = result;
+            keyPromptForm.Close();
+        }
+
+        private static void ReOpenKeyPromptForm(MainForm mainWindow, IOConnectionInfo dbFile)
+        {
+            Action action = () => mainWindow.OpenDatabase(dbFile, null, false);
+            mainWindow.Invoke(action);
+        }
+
+        private bool IsKeyForDataBaseExist(string dbPath)
+        {
+            return !String.IsNullOrEmpty(dbPath)
+                && _keyStorage.ContainsKey(dbPath);
+        }
+
         private bool ExtractCompositeKey(string dbPath, out CompositeKey compositeKey)
         {
             compositeKey = null;

diff --git a/KeyManagement/KeyStorage.cs b/KeyManagement/KeyStorage.cs
@@ -40,6 +40,13 @@ public void Remove(string dbPath)
             _keys.Remove(dbPath);
         }
 
+        public bool ContainsKey(string dbPath)
+        {
+            Data data;
+            return _keys.TryGetValue(dbPath, out data)
+                && data.IsValid();
+        }
+
         public bool TryGetValue(string dbPath, out ProtectedKey protectedKey)
         {
             Data data;