3 files changed, 76 insertions, 24 deletions

diff --git a/framework/src/domain/composercontroller.cpp b/framework/src/domain/composercontroller.cpp
index fb3857a5..630f9e10 100644
--- a/framework/src/domain/composercontroller.cpp
+++ b/framework/src/domain/composercontroller.cpp
@@ -38,6 +38,8 @@
 #include "mime/mailtemplates.h"
 #include "mime/mailcrypto.h"
 
+Q_DECLARE_METATYPE(GpgME::Key);
+
 class IdentitySelector : public Selector {
 public:
     IdentitySelector(ComposerController &controller) : Selector(new IdentitiesModel), mController(controller)
@@ -105,6 +107,18 @@ public:
         item->setData(addressee, ComposerController::AddresseeNameRole);
         item->setData(false, ComposerController::KeyFoundRole);
         appendRow(QList<QStandardItem*>() << item);
+        findKey(addressee, item);
+    }
+
+    void findKey(const QString &addressee, QStandardItem *item)
+    {
+        auto keys = MailCrypto::findKeys(QStringList{} << addressee, false, MailCrypto::OPENPGP);
+        if (item) {
+            if (!keys.empty()) {
+                item->setData(true, ComposerController::KeyFoundRole);
+                item->setData(QVariant::fromValue(keys.front()), ComposerController::KeyRole);
+            }
+        }
     }
 
     void remove(const QString &addressee)
@@ -118,6 +132,30 @@ public:
         }
     }
 
+    bool foundAllKeys()
+    {
+        std::vector<GpgME::Key> keys;
+        auto root = invisibleRootItem();
+        for (int row = 0; row < root->rowCount(); row++) {
+            auto item = root->child(row, 0);
+            if (!item->data(ComposerController::KeyFoundRole).toBool()) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    std::vector<GpgME::Key> getKeys()
+    {
+        std::vector<GpgME::Key> keys;
+        auto root = invisibleRootItem();
+        for (int row = 0; row < root->rowCount(); row++) {
+            auto item = root->child(row, 0);
+            keys.push_back(item->data(ComposerController::KeyRole).value<GpgME::Key>());
+        }
+        return keys;
+    }
+
     void setStringList(const QStringList &list)
     {
         clear();
@@ -479,8 +517,21 @@ KMime::Message::Ptr ComposerController::assembleMessage()
     if (getSign()) {
         signingKeys = mPersonalKeys;
     }
+    std::vector<GpgME::Key> encryptionKeys;
+    if (getEncrypt()) {
+        if (!mToModel->foundAllKeys() || !mCcModel->foundAllKeys() || !mBccModel->foundAllKeys()) {
+            qWarning() << "Can't encrypt with missing keys";
+            return nullptr;
+        }
+        auto toKeys = mToModel->getKeys();
+        encryptionKeys.insert(std::end(encryptionKeys), std::begin(toKeys), std::end(toKeys));
+        auto ccKeys = mCcModel->getKeys();
+        encryptionKeys.insert(std::end(encryptionKeys), std::begin(ccKeys), std::end(ccKeys));
+        auto bccKeys = mBccModel->getKeys();
+        encryptionKeys.insert(std::end(encryptionKeys), std::begin(bccKeys), std::end(bccKeys));
+    }
 
-    return MailTemplates::createMessage(mExistingMessage, mToModel->stringList(), mCcModel->stringList(), mBccModel->stringList(), getIdentity(), getSubject(), getBody(), getHtmlBody(), attachments, signingKeys);
+    return MailTemplates::createMessage(mExistingMessage, mToModel->stringList(), mCcModel->stringList(), mBccModel->stringList(), getIdentity(), getSubject(), getBody(), getHtmlBody(), attachments, signingKeys, encryptionKeys);
 }
 
 void ComposerController::updateSendAction()
diff --git a/framework/src/domain/mime/mailtemplates.cpp b/framework/src/domain/mime/mailtemplates.cpp
index bbe079d8..5f5b35df 100644
--- a/framework/src/domain/mime/mailtemplates.cpp
+++ b/framework/src/domain/mime/mailtemplates.cpp
@@ -958,7 +958,7 @@ static KMime::Types::Mailbox::List stringListToMailboxes(const QStringList &list
     return mailboxes;
 }
 
-KMime::Message::Ptr MailTemplates::createMessage(KMime::Message::Ptr existingMessage, const QStringList &to, const QStringList &cc, const QStringList &bcc, const KMime::Types::Mailbox &from, const QString &subject, const QString &body, bool htmlBody, const QList<Attachment> &attachments, const std::vector<GpgME::Key> &signingKeys)
+KMime::Message::Ptr MailTemplates::createMessage(KMime::Message::Ptr existingMessage, const QStringList &to, const QStringList &cc, const QStringList &bcc, const KMime::Types::Mailbox &from, const QString &subject, const QString &body, bool htmlBody, const QList<Attachment> &attachments, const std::vector<GpgME::Key> &signingKeys, const std::vector<GpgME::Key> &encryptionKeys)
 {
     auto mail = existingMessage;
     if (!mail) {
@@ -995,41 +995,42 @@ KMime::Message::Ptr MailTemplates::createMessage(KMime::Message::Ptr existingMes
     }
     mail->assemble();
 
-    KMime::Content *bodyPart;
-    if (!attachments.isEmpty()) {
-        bodyPart = new KMime::Content;
-        bodyPart->contentType(true)->setMimeType("multipart/mixed");
-        bodyPart->contentType()->setBoundary(KMime::multiPartBoundary());
-        bodyPart->contentTransferEncoding()->setEncoding(KMime::Headers::CE7Bit);
-        bodyPart->setPreamble("This is a multi-part message in MIME format.\n");
-        bodyPart->addContent(createBodyPart(body.toUtf8(), htmlBody));
-        for (const auto &attachment : attachments) {
-            bodyPart->addContent(createAttachmentPart(attachment.data, attachment.filename, attachment.isInline, attachment.mimeType, attachment.name));
+    std::unique_ptr<KMime::Content> bodyPart{[&] {
+        if (!attachments.isEmpty()) {
+            auto bodyPart = new KMime::Content;
+            bodyPart->contentType(true)->setMimeType("multipart/mixed");
+            bodyPart->contentType()->setBoundary(KMime::multiPartBoundary());
+            bodyPart->contentTransferEncoding()->setEncoding(KMime::Headers::CE7Bit);
+            bodyPart->setPreamble("This is a multi-part message in MIME format.\n");
+            bodyPart->addContent(createBodyPart(body.toUtf8(), htmlBody));
+            for (const auto &attachment : attachments) {
+                bodyPart->addContent(createAttachmentPart(attachment.data, attachment.filename, attachment.isInline, attachment.mimeType, attachment.name));
+            }
+            return bodyPart;
+        } else {
+            return createBodyPart(body.toUtf8(), htmlBody);
         }
-    } else {
-        bodyPart = createBodyPart(body.toUtf8(), htmlBody);
-    }
+    }()};
     bodyPart->assemble();
 
-    KMime::Content *signedResult = nullptr;
-    if (!signingKeys.empty()) {
-        signedResult = MailCrypto::sign(bodyPart, signingKeys);
-        if (!signedResult) {
+    QByteArray bodyData;
+    if (!signingKeys.empty() || !encryptionKeys.empty()) {
+        auto result = MailCrypto::processCrypto(bodyPart.get(), signingKeys, encryptionKeys, MailCrypto::OPENPGP);
+        if (!result) {
             qWarning() << "Signing failed";
             return {};
         }
+        bodyData = result->encodedContent();
     } else {
         if (!bodyPart->contentType(false)) {
             bodyPart->contentType(true)->setMimeType("text/plain");
             bodyPart->assemble();
         }
+        bodyData = bodyPart->encodedContent();
     }
 
-    const QByteArray allData = mail->head() + (signedResult ? signedResult->encodedContent() : bodyPart->encodedContent());
-    delete bodyPart;
     KMime::Message::Ptr resultMessage(new KMime::Message);
-    resultMessage->setContent(allData);
+    resultMessage->setContent(mail->head() + bodyData);
     resultMessage->parse(); // Not strictly necessary.
-
     return resultMessage;
 }
diff --git a/framework/src/domain/mime/mailtemplates.h b/framework/src/domain/mime/mailtemplates.h
index 51db2ba0..21efb5a0 100644
--- a/framework/src/domain/mime/mailtemplates.h
+++ b/framework/src/domain/mime/mailtemplates.h
@@ -37,5 +37,5 @@ namespace MailTemplates
     void reply(const KMime::Message::Ptr &origMsg, const std::function<void(const KMime::Message::Ptr &result)> &callback, const KMime::Types::AddrSpecList &me = {});
     QString plaintextContent(const KMime::Message::Ptr &origMsg);
     QString body(const KMime::Message::Ptr &msg, bool &isHtml);
-    KMime::Message::Ptr createMessage(KMime::Message::Ptr existingMessage, const QStringList &to, const QStringList &cc, const QStringList &bcc, const KMime::Types::Mailbox &from, const QString &subject, const QString &body, bool htmlBody, const QList<Attachment> &attachments, const std::vector<GpgME::Key> &signingKeys = {});
+    KMime::Message::Ptr createMessage(KMime::Message::Ptr existingMessage, const QStringList &to, const QStringList &cc, const QStringList &bcc, const KMime::Types::Mailbox &from, const QString &subject, const QString &body, bool htmlBody, const QList<Attachment> &attachments, const std::vector<GpgME::Key> &signingKeys = {}, const std::vector<GpgME::Key> &encryptionKeys = {});
 };