From 0cc2c83525498050ff41f81cafff7a92f4d24162 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sami=20V=C3=A4nttinen?= <sami.vanttinen@protonmail.com>
Date: Mon, 17 Aug 2020 13:17:58 +0300
Subject: [PATCH] Add command for retrieving the current TOTP (#5278)

---
 src/browser/BrowserAction.cpp  | 33 +++++++++++++++++++++++++++++++++
 src/browser/BrowserAction.h    |  1 +
 src/browser/BrowserService.cpp | 26 +++++++++++++++++++++++++-
 src/browser/BrowserService.h   |  1 +
 4 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/browser/BrowserAction.cpp b/src/browser/BrowserAction.cpp
index 361dc2a9c..8e0c26909 100644
--- a/src/browser/BrowserAction.cpp
+++ b/src/browser/BrowserAction.cpp
@@ -107,6 +107,8 @@ QJsonObject BrowserAction::handleAction(const QJsonObject& json)
         return handleGetDatabaseGroups(json, action);
     } else if (action.compare("create-new-group", Qt::CaseSensitive) == 0) {
         return handleCreateNewGroup(json, action);
+    } else if (action.compare("get-totp", Qt::CaseSensitive) == 0) {
+        return handleGetTotp(json, action);
     }
 
     // Action was not recognized
@@ -465,6 +467,37 @@ QJsonObject BrowserAction::handleCreateNewGroup(const QJsonObject& json, const Q
     return buildResponse(action, message, newNonce);
 }
 
+QJsonObject BrowserAction::handleGetTotp(const QJsonObject& json, const QString& action)
+{
+    const QString nonce = json.value("nonce").toString();
+    const QString encrypted = json.value("message").toString();
+
+    if (!m_associated) {
+        return getErrorReply(action, ERROR_KEEPASS_ASSOCIATION_FAILED);
+    }
+
+    const QJsonObject decrypted = decryptMessage(encrypted, nonce);
+    if (decrypted.isEmpty()) {
+        return getErrorReply(action, ERROR_KEEPASS_CANNOT_DECRYPT_MESSAGE);
+    }
+
+    QString command = decrypted.value("action").toString();
+    if (command.isEmpty() || command.compare("get-totp", Qt::CaseSensitive) != 0) {
+        return getErrorReply(action, ERROR_KEEPASS_INCORRECT_ACTION);
+    }
+
+    const QString uuid = decrypted.value("uuid").toString();
+
+    // Get the current TOTP
+    const auto totp = browserService()->getCurrentTotp(uuid);
+    const QString newNonce = incrementNonce(nonce);
+
+    QJsonObject message = buildMessage(newNonce);
+    message["totp"] = totp;
+
+    return buildResponse(action, message, newNonce);
+}
+
 QJsonObject BrowserAction::getErrorReply(const QString& action, const int errorCode) const
 {
     QJsonObject response;
diff --git a/src/browser/BrowserAction.h b/src/browser/BrowserAction.h
index c65409dd8..06d6a131a 100644
--- a/src/browser/BrowserAction.h
+++ b/src/browser/BrowserAction.h
@@ -41,6 +41,7 @@ private:
     QJsonObject handleLockDatabase(const QJsonObject& json, const QString& action);
     QJsonObject handleGetDatabaseGroups(const QJsonObject& json, const QString& action);
     QJsonObject handleCreateNewGroup(const QJsonObject& json, const QString& action);
+    QJsonObject handleGetTotp(const QJsonObject& json, const QString& action);
 
     QJsonObject buildMessage(const QString& nonce) const;
     QJsonObject buildResponse(const QString& action, const QJsonObject& message, const QString& nonce);
diff --git a/src/browser/BrowserService.cpp b/src/browser/BrowserService.cpp
index 01e9a4281..eb752996c 100644
--- a/src/browser/BrowserService.cpp
+++ b/src/browser/BrowserService.cpp
@@ -215,7 +215,6 @@ QJsonObject BrowserService::getDatabaseGroups()
 
 QJsonObject BrowserService::createNewGroup(const QString& groupName)
 {
-
     auto db = getDatabase();
     if (!db) {
         return {};
@@ -284,6 +283,31 @@ QJsonObject BrowserService::createNewGroup(const QString& groupName)
     return result;
 }
 
+QString BrowserService::getCurrentTotp(const QString& uuid)
+{
+    QList<QSharedPointer<Database>> databases;
+    if (browserSettings()->searchInAllDatabases()) {
+        for (auto dbWidget : getMainWindow()->getOpenDatabases()) {
+            auto db = dbWidget->database();
+            if (db) {
+                databases << db;
+            }
+        }
+    } else {
+        databases << getDatabase();
+    }
+
+    auto entryUuid = Tools::hexToUuid(uuid);
+    for (const auto& db : databases) {
+        auto entry = db->rootGroup()->findEntryByUuid(entryUuid, true);
+        if (entry) {
+            return entry->totp();
+        }
+    }
+
+    return {};
+}
+
 QString BrowserService::storeKey(const QString& key)
 {
     auto db = getDatabase();
diff --git a/src/browser/BrowserService.h b/src/browser/BrowserService.h
index 77635cfe1..f52b502b3 100644
--- a/src/browser/BrowserService.h
+++ b/src/browser/BrowserService.h
@@ -58,6 +58,7 @@ public:
 
     QJsonObject getDatabaseGroups();
     QJsonObject createNewGroup(const QString& groupName);
+    QString getCurrentTotp(const QString& uuid);
 
     void addEntry(const QString& dbid,
                   const QString& login,