Add a TOTP Status Column (#123)

* Add a TOTP Status Column

This column shows whether the column has (valid) settings to calculate a TOTP.
Double clicking it brings up the Setup TOTP dialog.

* Process PR feedback

* Merge the two column providers.
* Add additional tests
* Double click behavior for both columns is to copy to the clipboard.

* Show Setup TOTP when a TOTP can not be generated on doubleclick

KeeTrayTOTP.Tests/TrayTOTP_ColumnproviderTests.cs

@@ -0,0 +1,141 @@
+using FluentAssertions;
+using KeePass.App.Configuration;
+using KeePass.Forms;
+using KeePass.Plugins;
+using KeePass.UI;
+using KeePassLib.Security;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Moq;
+using System;
+
+namespace KeeTrayTOTP.Tests
+{
+    [TestClass]
+    public class TrayTOTP_ColumnproviderTests : IDisposable
+    {
+        private readonly KeeTrayTOTPExt _plugin;
+        private readonly IPluginHost _pluginHost;
+
+        const string InvalidSeed = "C5CYMIHWQUUZMKUGZHGEOSJSQDE4L===";
+        const string ValidSeed = "JBSWY3DPEHPK3PXP";
+        const string ValidSettings = "30;6";
+
+        public TrayTOTP_ColumnproviderTests()
+        {
+            (_plugin, _pluginHost) = CreateInitializedPlugin();
+        }
+
+        [DataRow(ValidSeed, ValidSettings, "TOTP Enabled")]
+        [DataRow(ValidSeed, ";6", "Error, bad settings!")]
+        [DataRow(ValidSeed, "30", "Error, bad settings!")]
+        [DataRow(ValidSeed, null, "Error, storage!")]
+        [DataRow(InvalidSeed, ValidSettings, "Error, bad seed!")]
+        [DataRow(null, ValidSettings, "Error, storage!")]
+        [DataRow(null, null, "")]
+        [DataTestMethod]
+        public void GetCellDataStatus_ShouldReturnExpectedValues(string seed, string settings, string expected)
+        {
+            var column = new TrayTOTP_ColumnProvider(_plugin);
+            var pwEntry = new KeePassLib.PwEntry(true, true);
+            if (seed != null)
+            {
+                var seedKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSeed_StringName, Localization.Strings.TOTPSeed);
+                pwEntry.Strings.Set(seedKey, new ProtectedString(false, seed));
+            }
+            if (settings != null)
+            {
+                var settingsKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSettings_StringName, Localization.Strings.TOTPSettings);
+                pwEntry.Strings.Set(settingsKey, new ProtectedString(false, settings));
+            }
+
+            var actual = column.GetCellData("TOTP Status", pwEntry);
+
+            actual.Should().Be(expected);
+        }
+
+        [DataRow(ValidSeed, ";6", "Error, bad settings!")]
+        [DataRow(ValidSeed, "30", "Error, bad settings!")]
+        [DataRow(ValidSeed, null, "Error, storage!")]
+        [DataRow(InvalidSeed, ValidSettings, "Error, bad seed!")]
+        [DataRow(null, ValidSettings, "Error, storage!")]
+        [DataRow(null, null, "")]
+        [DataTestMethod]
+        public void GetCellDataCode_ShouldReturnExpectedValues(string seed, string settings, string expected)
+        {
+            var column = new TrayTOTP_ColumnProvider(_plugin);
+            var pwEntry = new KeePassLib.PwEntry(true, true);
+            if (seed != null)
+            {
+                var seedKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSeed_StringName, Localization.Strings.TOTPSeed);
+                pwEntry.Strings.Set(seedKey, new ProtectedString(false, seed));
+            }
+            if (settings != null)
+            {
+                var settingsKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSettings_StringName, Localization.Strings.TOTPSettings);
+                pwEntry.Strings.Set(settingsKey, new ProtectedString(false, settings));
+            }
+
+            var actual = column.GetCellData("TOTP", pwEntry);
+
+            actual.Should().Be(expected);
+        }
+
+        [DataRow(true, @"^\d{6} \(\d{1,2}\)$", DisplayName = "Column timer visible should show a code and validity")]
+        [DataRow(false, @"^\d{6}$", DisplayName = "Column timer invisible should only show a code")]
+        [DataTestMethod]
+        public void GetCellDataCode_WithValidSeedAndSettings_ShouldReturnA6DigitCodeWithDuration(bool showTimer, string regex)
+        {
+            _plugin.PluginHost.CustomConfig.SetBool(KeeTrayTOTPExt.setname_bool_TOTPColumnTimer_Visible, showTimer);
+
+            var column = new TrayTOTP_ColumnProvider(_plugin);
+            var pwEntry = new KeePassLib.PwEntry(true, true);
+            var seedKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSeed_StringName, Localization.Strings.TOTPSeed);
+            pwEntry.Strings.Set(seedKey, new ProtectedString(false, ValidSeed));
+            var settingsKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSettings_StringName, Localization.Strings.TOTPSettings);
+            pwEntry.Strings.Set(settingsKey, new ProtectedString(false, ValidSettings));
+
+            var actual = column.GetCellData("TOTP", pwEntry);
+
+            actual.Should().MatchRegex(regex);
+        }
+
+        [TestMethod]
+        public void GetCellData_WithAnInvalidColumn_ShouldReturnEmptyString()
+        {
+            var column = new TrayTOTP_ColumnProvider(_plugin);
+            var pwEntry = new KeePassLib.PwEntry(true, true);
+            var seedKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSeed_StringName, Localization.Strings.TOTPSeed);
+            pwEntry.Strings.Set(seedKey, new ProtectedString(false, ValidSeed));
+            var settingsKey = _pluginHost.CustomConfig.GetString(KeeTrayTOTPExt.setname_string_TOTPSettings_StringName, Localization.Strings.TOTPSettings);
+            pwEntry.Strings.Set(settingsKey, new ProtectedString(false, ValidSettings));
+
+            var actual = column.GetCellData("InvalidColumnName", pwEntry);
+
+            actual.Should().BeEmpty();
+        }
+
+        private static (KeeTrayTOTPExt, IPluginHost) CreateInitializedPlugin()
+        {
+            var plugin = new KeeTrayTOTPExt();
+            var pluginHost = new Mock<IPluginHost>(MockBehavior.Strict);
+
+            var keepassForm = new MainForm();
+            pluginHost.SetupGet(c => c.MainWindow).Returns(keepassForm);
+
+            var customConfig = new AceCustomConfig();
+            pluginHost.SetupGet(c => c.CustomConfig).Returns(customConfig);
+
+            var columnProviderPool = new ColumnProviderPool();
+            pluginHost.SetupGet(c => c.ColumnProviderPool).Returns(columnProviderPool);
+
+            plugin.Initialize(pluginHost.Object);
+
+            return (plugin, pluginHost.Object);
+        }
+
+        public void Dispose()
+        {
+            _plugin.Terminate();
+        }
+    }
+}

KeeTrayTOTP.Tests/TrayTOTP_PluginTests.cs

@@ -22,6 +22,34 @@ public void InitializePlugin_ShouldNotThrow()
             act.Should().NotThrow();
         }
 
+        [TestMethod]
+        public void InitializePlugin_ShouldAddOneColumnProvider()
+        {
+            var plugin = CreatePluginHostMock(out var host);
+            var numberOfColumnsBeforeInitialize = host.Object.ColumnProviderPool.Count;
+
+            plugin.Initialize(host.Object);
+
+            host.Object.ColumnProviderPool.Should().HaveCount(numberOfColumnsBeforeInitialize + 1);
+        }
+
+        [TestMethod]
+        public void Terminate_ShouldOnlyRemoveOurColumns()
+        {
+            var plugin = CreatePluginHostMock(out var host);
+
+            // add a fake column that does not belong to our plugin
+            var otherColumn = new Mock<ColumnProvider>().Object;
+            host.Object.ColumnProviderPool.Add(otherColumn);
+
+            plugin.Initialize(host.Object);
+            plugin.Terminate();
+
+            // only the fake column should be in the pool after the plugin terminates
+            host.Object.ColumnProviderPool.Should().HaveCount(1);
+            host.Object.ColumnProviderPool.Should().OnlyContain(c => c == otherColumn);
+        }
+
         private static KeeTrayTOTPExt CreatePluginHostMock(out Mock<IPluginHost> host)
         {
             var plugin = new KeeTrayTOTPExt();

KeeTrayTOTP/KeeTrayTOTP.csproj

@@ -103,7 +103,6 @@
     <Compile Include="ShowQR.Designer.cs">
       <DependentUpon>ShowQR.cs</DependentUpon>
     </Compile>
-    <Compile Include="TrayTOTP_CustomColumn.cs" />
     <Compile Include="TrayTOTP_Extensions.cs" />
     <Compile Include="FormHelp.cs">
       <SubType>Form</SubType>
@@ -132,6 +131,7 @@
     <Compile Include="Libraries\TOTPProvider.cs" />
     <Compile Include="TrayTOTP_Plugin.cs" />
     <Compile Include="TrayTOTP_TimeCorrectionCollection.cs" />
+    <Compile Include="TrayTOTP_ColumnProvider.cs" />
   </ItemGroup>
   <ItemGroup>
     <EmbeddedResource Include="FormAbout.resx">

KeeTrayTOTP/Localization/Strings.resx

@@ -120,6 +120,12 @@
   <data name="About" xml:space="preserve">
     <value>About</value>
   </data>
+  <data name="ColumnTOTPCode" xml:space="preserve">
+    <value>TOTP</value>
+  </data>
+  <data name="ColumnTOTPStatus" xml:space="preserve">
+    <value>TOTP Status</value>
+  </data>
   <data name="ConnectionFailed" xml:space="preserve">
     <value>Connection failed!</value>
   </data>
@@ -265,6 +271,9 @@
   <data name="TOTP" xml:space="preserve">
     <value>TOTP</value>
   </data>
+  <data name="TOTPEnabled" xml:space="preserve">
+    <value>TOTP Enabled</value>
+  </data>
   <data name="TOTPSeed" xml:space="preserve">
     <value>TOTP Seed</value>
   </data>

KeeTrayTOTP/TrayTOTP_ColumnProvider.cs

@@ -0,0 +1,112 @@
+using System;
+using System.Windows.Forms;
+
+using KeePass.UI;
+using KeePassLib;
+using KeeTrayTOTP.Libraries;
+
+namespace KeeTrayTOTP
+{
+    /// <summary>
+    /// Provides columns to Keepass showing the TOTP Code / TOTP Status for an entry.
+    /// </summary>
+    internal class TrayTOTP_ColumnProvider : ColumnProvider
+    {
+        private readonly KeeTrayTOTPExt _plugin;
+
+        internal TrayTOTP_ColumnProvider(KeeTrayTOTPExt plugin)
+        {
+            _plugin = plugin;
+        }
+
+        private static readonly string[] _columnName = new[] { Localization.Strings.ColumnTOTPCode, Localization.Strings.ColumnTOTPStatus };
+
+        public override string[] ColumnNames
+        {
+            get { return _columnName; }
+        }
+
+        public override HorizontalAlignment TextAlign
+        {
+            get { return HorizontalAlignment.Left; }
+        }
+
+        /// <summary>
+        /// Tells KeePass what to display in the column.
+        /// </summary>
+        /// <param name="columnName"></param>
+        /// <param name="pe"></param>
+        /// <returns>String displayed in the columns.</returns>
+        public override string GetCellData(string columnName, PwEntry pe)
+        {
+            if (columnName == Localization.Strings.ColumnTOTPCode)
+            {
+                return GetCellDataInternal(pe, GetInnerValueCode);
+            }
+            else if (columnName == Localization.Strings.ColumnTOTPStatus)
+            {
+                return GetCellDataInternal(pe, GetInnerValueStatus);
+            }
+
+            return string.Empty;
+        }
+
+        private string GetCellDataInternal(PwEntry pe, Func<PwEntry, string> innerValueFunc)
+        {
+            var settingsCheck = _plugin.SettingsCheck(pe);
+            var seedCheck = _plugin.SeedCheck(pe);
+
+            if (settingsCheck && seedCheck)
+            {
+                if (_plugin.SettingsValidate(pe))
+                {
+                    if (_plugin.SeedValidate(pe))
+                    {
+                        return innerValueFunc(pe);
+                    }
+                    return Localization.Strings.ErrorBadSeed;
+                }
+                return Localization.Strings.ErrorBadSettings;
+            }
+            return (settingsCheck || seedCheck) ? Localization.Strings.ErrorStorage : string.Empty;
+        }
+
+        private static string GetInnerValueStatus(PwEntry entry)
+        {
+            return Localization.Strings.TOTPEnabled;
+        }
+
+        private string GetInnerValueCode(PwEntry entry)
+        {
+            string[] settings = _plugin.SettingsGet(entry);
+            var totpGenerator = new TOTPProvider(settings, ref _plugin.TimeCorrections);
+            return totpGenerator.GenerateByByte(Base32.Decode(_plugin.SeedGet(entry).ReadString().ExtWithoutSpaces())) + (_plugin.PluginHost.CustomConfig.GetBool(KeeTrayTOTPExt.setname_bool_TOTPColumnTimer_Visible, true) ? totpGenerator.Timer.ToString().ExtWithParenthesis().ExtWithSpaceBefore() : string.Empty);
+        }
+
+        /// <summary>
+        /// Informs KeePass if PerformCellAction must be called when the cell is double clicked.
+        /// </summary>
+        /// <param name="columnName">Column Name.</param>
+        public override bool SupportsCellAction(string columnName)
+        {
+            return true;
+        }
+
+        /// <summary>
+        /// Happens when a cell of the column is double-clicked.
+        /// </summary>
+        /// <param name="columnName">Column's name.</param>
+        /// <param name="pe">Entry associated with the clicked cell.</param>
+        public override void PerformCellAction(string columnName, PwEntry pe)
+        {
+            if (!_plugin.CanGenerateTOTP(pe))
+            {
+                UIUtil.ShowDialogAndDestroy(new SetupTOTP(_plugin, pe));
+            }
+            else if (_plugin.PluginHost.CustomConfig.GetBool(KeeTrayTOTPExt.setname_bool_TOTPColumnCopy_Enable, true))
+            {
+                _plugin.TOTPCopyToClipboard(pe);
+            }
+        }
+    }
+}