SSH Agent: Entry context menu control

10 files changed, 128 insertions, 0 deletions

diff --git a/docs/KEYBINDS.md b/docs/KEYBINDS.md
index 968df2037..24f85f68e 100644
--- a/docs/KEYBINDS.md
+++ b/docs/KEYBINDS.md
@@ -20,6 +20,8 @@ Copy Password                | Ctrl + C
 Trigger AutoType             | Ctrl + Shift - V
 Open URL                     | Ctrl + Shift - U
 Copy URL                     | Ctrl + U
+Add key to SSH Agent         | Ctrl + H
+Remove key from SSH Agent    | Ctrl + Shift + H
 Show Minimized               | Ctrl + M
 Hide Window                  | Ctrl + Shift - M
 Select Next Database Tab     | Ctrl + Tab *OR* Ctrl + PGDN
diff --git a/src/gui/DatabaseWidget.cpp b/src/gui/DatabaseWidget.cpp
index e6cc9f309..a9f628735 100644
--- a/src/gui/DatabaseWidget.cpp
+++ b/src/gui/DatabaseWidget.cpp
@@ -699,6 +699,50 @@ void DatabaseWidget::setClipboardTextAndMinimize(const QString& text)
     }
 }
 
+#ifdef WITH_XC_SSHAGENT
+void DatabaseWidget::addToAgent()
+{
+    Entry* currentEntry = m_entryView->currentEntry();
+    Q_ASSERT(currentEntry);
+    if (!currentEntry) {
+        return;
+    }
+
+    KeeAgentSettings settings;
+    if (!settings.fromEntry(currentEntry)) {
+        return;
+    }
+
+    OpenSSHKey key;
+    if (settings.toOpenSSHKey(currentEntry, key, true)) {
+        SSHAgent::instance()->addIdentity(key, settings);
+    } else {
+        m_messageWidget->showMessage(key.errorString(), MessageWidget::Error);
+    }
+}
+
+void DatabaseWidget::removeFromAgent()
+{
+    Entry* currentEntry = m_entryView->currentEntry();
+    Q_ASSERT(currentEntry);
+    if (!currentEntry) {
+        return;
+    }
+
+    KeeAgentSettings settings;
+    if (!settings.fromEntry(currentEntry)) {
+        return;
+    }
+
+    OpenSSHKey key;
+    if (settings.toOpenSSHKey(currentEntry, key, false)) {
+        SSHAgent::instance()->removeIdentity(key);
+    } else {
+        m_messageWidget->showMessage(key.errorString(), MessageWidget::Error);
+    }
+}
+#endif
+
 void DatabaseWidget::performAutoType()
 {
     auto currentEntry = currentSelectedEntry();
@@ -1625,6 +1669,19 @@ bool DatabaseWidget::currentEntryHasTotp()
     return currentEntry->hasTotp();
 }
 
+#ifdef WITH_XC_SSHAGENT
+bool DatabaseWidget::currentEntryHasSshKey()
+{
+    Entry* currentEntry = m_entryView->currentEntry();
+    Q_ASSERT(currentEntry);
+    if (!currentEntry) {
+        return false;
+    }
+
+    return KeeAgentSettings::inEntry(currentEntry);
+}
+#endif
+
 bool DatabaseWidget::currentEntryHasNotes()
 {
     auto currentEntry = currentSelectedEntry();
diff --git a/src/gui/DatabaseWidget.h b/src/gui/DatabaseWidget.h
index 2b9388f37..d54c63439 100644
--- a/src/gui/DatabaseWidget.h
+++ b/src/gui/DatabaseWidget.h
@@ -110,6 +110,9 @@ public:
     bool currentEntryHasUrl();
     bool currentEntryHasNotes();
     bool currentEntryHasTotp();
+#ifdef WITH_XC_SSHAGENT
+    bool currentEntryHasSshKey();
+#endif
 
     QByteArray entryViewState() const;
     bool setEntryViewState(const QByteArray& state) const;
@@ -169,6 +172,10 @@ public slots:
     void showTotpKeyQrCode();
     void copyTotp();
     void setupTotp();
+#ifdef WITH_XC_SSHAGENT
+    void addToAgent();
+    void removeFromAgent();
+#endif
     void performAutoType();
     void openUrl();
     void downloadSelectedFavicons();
diff --git a/src/gui/MainWindow.cpp b/src/gui/MainWindow.cpp
index 485c21125..d4257a772 100644
--- a/src/gui/MainWindow.cpp
+++ b/src/gui/MainWindow.cpp
@@ -190,9 +190,20 @@ MainWindow::MainWindow()
 
 #ifdef WITH_XC_SSHAGENT
     connect(sshAgent(), SIGNAL(error(QString)), this, SLOT(showErrorMessage(QString)));
+    connect(sshAgent(), SIGNAL(enabledChanged(bool)), this, SLOT(agentEnabled(bool)));
     m_ui->settingsWidget->addSettingsPage(new AgentSettingsPage(m_ui->tabWidget));
+
+    m_entryContextMenu->addSeparator();
+    m_entryContextMenu->addAction(m_ui->actionEntryAddToAgent);
+    m_entryContextMenu->addAction(m_ui->actionEntryRemoveFromAgent);
+
+    m_ui->actionEntryAddToAgent->setIcon(resources()->icon("utilities-terminal"));
+    m_ui->actionEntryRemoveFromAgent->setIcon(resources()->icon("utilities-terminal"));
 #endif
 
+    m_ui->actionEntryAddToAgent->setVisible(false);
+    m_ui->actionEntryRemoveFromAgent->setVisible(false);
+
 #if defined(WITH_XC_KEESHARE)
     KeeShare::init(this);
     m_ui->settingsWidget->addSettingsPage(new SettingsPageKeeShare(m_ui->tabWidget));
@@ -269,6 +280,8 @@ MainWindow::MainWindow()
     m_ui->actionEntryAutoType->setShortcut(Qt::CTRL + Qt::SHIFT + Qt::Key_V);
     m_ui->actionEntryOpenUrl->setShortcut(Qt::CTRL + Qt::SHIFT + Qt::Key_U);
     m_ui->actionEntryCopyURL->setShortcut(Qt::CTRL + Qt::Key_U);
+    m_ui->actionEntryAddToAgent->setShortcut(Qt::CTRL + Qt::Key_H);
+    m_ui->actionEntryRemoveFromAgent->setShortcut(Qt::CTRL + Qt::SHIFT + Qt::Key_H);
 
 #if QT_VERSION >= QT_VERSION_CHECK(5, 10, 0)
     // Qt 5.10 introduced a new "feature" to hide shortcuts in context menus
@@ -285,6 +298,8 @@ MainWindow::MainWindow()
     m_ui->actionEntryAutoType->setShortcutVisibleInContextMenu(true);
     m_ui->actionEntryOpenUrl->setShortcutVisibleInContextMenu(true);
     m_ui->actionEntryCopyURL->setShortcutVisibleInContextMenu(true);
+    m_ui->actionEntryAddToAgent->setShortcutVisibleInContextMenu(true);
+    m_ui->actionEntryRemoveFromAgent->setShortcutVisibleInContextMenu(true);
 #endif
 
     connect(m_ui->menuEntries, SIGNAL(aboutToShow()), SLOT(obtainContextFocusLock()));
@@ -440,6 +455,10 @@ MainWindow::MainWindow()
     m_actionMultiplexer.connect(m_ui->actionEntryAutoType, SIGNAL(triggered()), SLOT(performAutoType()));
     m_actionMultiplexer.connect(m_ui->actionEntryOpenUrl, SIGNAL(triggered()), SLOT(openUrl()));
     m_actionMultiplexer.connect(m_ui->actionEntryDownloadIcon, SIGNAL(triggered()), SLOT(downloadSelectedFavicons()));
+#ifdef WITH_XC_SSHAGENT
+    m_actionMultiplexer.connect(m_ui->actionEntryAddToAgent, SIGNAL(triggered()), SLOT(addToAgent()));
+    m_actionMultiplexer.connect(m_ui->actionEntryRemoveFromAgent, SIGNAL(triggered()), SLOT(removeFromAgent()));
+#endif
 
     m_actionMultiplexer.connect(m_ui->actionGroupNew, SIGNAL(triggered()), SLOT(createGroup()));
     m_actionMultiplexer.connect(m_ui->actionGroupEdit, SIGNAL(triggered()), SLOT(switchToGroupEdit()));
@@ -691,6 +710,14 @@ void MainWindow::setMenuActionState(DatabaseWidget::Mode mode)
             m_ui->actionExportCsv->setEnabled(true);
             m_ui->actionExportHtml->setEnabled(true);
             m_ui->actionDatabaseMerge->setEnabled(m_ui->tabWidget->currentIndex() != -1);
+#ifdef WITH_XC_SSHAGENT
+            bool singleEntryHasSshKey =
+                singleEntrySelected && sshAgent()->isEnabled() && dbWidget->currentEntryHasSshKey();
+            m_ui->actionEntryAddToAgent->setVisible(singleEntryHasSshKey);
+            m_ui->actionEntryAddToAgent->setEnabled(singleEntryHasSshKey);
+            m_ui->actionEntryRemoveFromAgent->setVisible(singleEntryHasSshKey);
+            m_ui->actionEntryRemoveFromAgent->setEnabled(singleEntryHasSshKey);
+#endif
 
             m_searchWidgetAction->setEnabled(true);
 
@@ -1185,6 +1212,12 @@ void MainWindow::releaseContextFocusLock()
     m_contextMenuFocusLock = false;
 }
 
+void MainWindow::agentEnabled(bool enabled)
+{
+    m_ui->actionEntryAddToAgent->setVisible(enabled);
+    m_ui->actionEntryRemoveFromAgent->setVisible(enabled);
+}
+
 void MainWindow::showEntryContextMenu(const QPoint& globalPos)
 {
     bool entrySelected = false;
diff --git a/src/gui/MainWindow.h b/src/gui/MainWindow.h
index 5e2bfef80..ac918f0b0 100644
--- a/src/gui/MainWindow.h
+++ b/src/gui/MainWindow.h
@@ -131,6 +131,7 @@ private slots:
     void toggleUsernamesHidden();
     void obtainContextFocusLock();
     void releaseContextFocusLock();
+    void agentEnabled(bool enabled);
 
 private:
     static void setShortcut(QAction* action, QKeySequence::StandardKey standard, int fallback = 0);
diff --git a/src/gui/MainWindow.ui b/src/gui/MainWindow.ui
index c68d18b83..c65cc987a 100644
--- a/src/gui/MainWindow.ui
+++ b/src/gui/MainWindow.ui
@@ -337,6 +337,9 @@
     <addaction name="separator"/>
     <addaction name="actionEntryOpenUrl"/>
     <addaction name="actionEntryDownloadIcon"/>
+    <addaction name="separator"/>
+    <addaction name="actionEntryAddToAgent"/>
+    <addaction name="actionEntryRemoveFromAgent"/>
    </widget>
    <widget class="QMenu" name="menuGroups">
     <property name="title">
@@ -813,6 +816,16 @@
     <string notr="true">Ctrl+/</string>
    </property>
   </action>
+  <action name="actionEntryAddToAgent">
+   <property name="text">
+    <string>Add key to SSH Agent</string>
+   </property>
+  </action>
+  <action name="actionEntryRemoveFromAgent">
+   <property name="text">
+    <string>Remove key from SSH Agent</string>
+   </property>
+  </action>
  </widget>
  <customwidgets>
   <customwidget>
diff --git a/src/sshagent/KeeAgentSettings.cpp b/src/sshagent/KeeAgentSettings.cpp
index eae8412e0..71265dfea 100644
--- a/src/sshagent/KeeAgentSettings.cpp
+++ b/src/sshagent/KeeAgentSettings.cpp
@@ -300,6 +300,17 @@ QByteArray KeeAgentSettings::toXml() const
 }
 
 /**
+ * Check if an entry has KeeAgent settings configured
+ *
+ * @param entry Entry to check the attachment
+ * @return true if XML document exists
+ */
+bool KeeAgentSettings::inEntry(const Entry* entry)
+{
+    return entry->attachments()->hasKey("KeeAgent.settings");
+}
+
+/**
  * Read settings from an entry as an XML attachment.
  *
  * Sets error string on error.
diff --git a/src/sshagent/KeeAgentSettings.h b/src/sshagent/KeeAgentSettings.h
index 132b7bb5d..3be780055 100644
--- a/src/sshagent/KeeAgentSettings.h
+++ b/src/sshagent/KeeAgentSettings.h
@@ -34,6 +34,7 @@ public:
     bool fromXml(const QByteArray& ba);
     QByteArray toXml() const;
 
+    static bool inEntry(const Entry* entry);
     bool fromEntry(const Entry* entry);
     void toEntry(Entry* entry) const;
     bool keyConfigured() const;
diff --git a/src/sshagent/SSHAgent.cpp b/src/sshagent/SSHAgent.cpp
index 6bed354f2..cf7fd2522 100644
--- a/src/sshagent/SSHAgent.cpp
+++ b/src/sshagent/SSHAgent.cpp
@@ -53,6 +53,8 @@ void SSHAgent::setEnabled(bool enabled)
     }
 
     config()->set("SSHAgent", enabled);
+
+    emit enabledChanged(enabled);
 }
 
 QString SSHAgent::authSockOverride() const
diff --git a/src/sshagent/SSHAgent.h b/src/sshagent/SSHAgent.h
index 44ae88bb4..1800e794a 100644
--- a/src/sshagent/SSHAgent.h
+++ b/src/sshagent/SSHAgent.h
@@ -56,6 +56,7 @@ public:
 
 signals:
     void error(const QString& message);
+    void enabledChanged(bool enabled);
 
 public slots:
     void databaseModeChanged();