Add encrypted file download

4 files changed, 203 insertions, 165 deletions

diff --git a/src/timeline/TimelineModel.cpp b/src/timeline/TimelineModel.cpp
index b904dfd7..f606b603 100644
--- a/src/timeline/TimelineModel.cpp
+++ b/src/timeline/TimelineModel.cpp
@@ -3,11 +3,15 @@
 #include <algorithm>
 #include <type_traits>
 
+#include <QFileDialog>
+#include <QMimeDatabase>
 #include <QRegularExpression>
+#include <QStandardPaths>
 
 #include "ChatPage.h"
 #include "Logging.h"
 #include "MainWindow.h"
+#include "MxcImageProvider.h"
 #include "Olm.h"
 #include "TimelineViewManager.h"
 #include "Utils.h"
@@ -89,16 +93,41 @@ eventFormattedBody(const mtx::events::RoomEvent<T> &e)
 }
 
 template<class T>
+boost::optional<mtx::crypto::EncryptedFile>
+eventEncryptionInfo(const mtx::events::Event<T> &)
+{
+        return boost::none;
+}
+
+template<class T>
+auto
+eventEncryptionInfo(const mtx::events::RoomEvent<T> &e) -> std::enable_if_t<
+  std::is_same<decltype(e.content.file), boost::optional<mtx::crypto::EncryptedFile>>::value,
+  boost::optional<mtx::crypto::EncryptedFile>>
+{
+        return e.content.file;
+}
+
+template<class T>
 QString
 eventUrl(const mtx::events::Event<T> &)
 {
         return "";
 }
+
+QString
+eventUrl(const mtx::events::StateEvent<mtx::events::state::Avatar> &e)
+{
+        return QString::fromStdString(e.content.url);
+}
+
 template<class T>
 auto
 eventUrl(const mtx::events::RoomEvent<T> &e)
   -> std::enable_if_t<std::is_same<decltype(e.content.url), std::string>::value, QString>
 {
+        if (e.content.file)
+                return QString::fromStdString(e.content.file->url);
         return QString::fromStdString(e.content.url);
 }
 
@@ -1342,3 +1371,158 @@ TimelineModel::addPendingMessage(mtx::events::collections::TimelineEvents event)
         if (!isProcessingPending)
                 emit nextPendingMessage();
 }
+
+void
+TimelineModel::saveMedia(QString eventId) const
+{
+        mtx::events::collections::TimelineEvents event = events.value(eventId);
+
+        if (auto e = boost::get<mtx::events::EncryptedEvent<mtx::events::msg::Encrypted>>(&event)) {
+                event = decryptEvent(*e).event;
+        }
+
+        QString mxcUrl =
+          boost::apply_visitor([](const auto &e) -> QString { return eventUrl(e); }, event);
+        QString originalFilename =
+          boost::apply_visitor([](const auto &e) -> QString { return eventFilename(e); }, event);
+        QString mimeType =
+          boost::apply_visitor([](const auto &e) -> QString { return eventMimeType(e); }, event);
+
+        using EncF = boost::optional<mtx::crypto::EncryptedFile>;
+        EncF encryptionInfo =
+          boost::apply_visitor([](const auto &e) -> EncF { return eventEncryptionInfo(e); }, event);
+
+        qml_mtx_events::EventType eventType = boost::apply_visitor(
+          [](const auto &e) -> qml_mtx_events::EventType { return toRoomEventType(e); }, event);
+
+        QString dialogTitle;
+        if (eventType == qml_mtx_events::EventType::ImageMessage) {
+                dialogTitle = tr("Save image");
+        } else if (eventType == qml_mtx_events::EventType::VideoMessage) {
+                dialogTitle = tr("Save video");
+        } else if (eventType == qml_mtx_events::EventType::AudioMessage) {
+                dialogTitle = tr("Save audio");
+        } else {
+                dialogTitle = tr("Save file");
+        }
+
+        QString filterString = QMimeDatabase().mimeTypeForName(mimeType).filterString();
+
+        auto filename = QFileDialog::getSaveFileName(
+          manager_->getWidget(), dialogTitle, originalFilename, filterString);
+
+        if (filename.isEmpty())
+                return;
+
+        const auto url = mxcUrl.toStdString();
+
+        http::client()->download(
+          url,
+          [filename, url, encryptionInfo](const std::string &data,
+                                          const std::string &,
+                                          const std::string &,
+                                          mtx::http::RequestErr err) {
+                  if (err) {
+                          nhlog::net()->warn("failed to retrieve image {}: {} {}",
+                                             url,
+                                             err->matrix_error.error,
+                                             static_cast<int>(err->status_code));
+                          return;
+                  }
+
+                  try {
+                          auto temp = data;
+                          if (encryptionInfo)
+                                  temp = mtx::crypto::to_string(
+                                    mtx::crypto::decrypt_file(temp, encryptionInfo.value()));
+
+                          QFile file(filename);
+
+                          if (!file.open(QIODevice::WriteOnly))
+                                  return;
+
+                          file.write(QByteArray(temp.data(), (int)temp.size()));
+                          file.close();
+                  } catch (const std::exception &e) {
+                          nhlog::ui()->warn("Error while saving file to: {}", e.what());
+                  }
+          });
+}
+
+void
+TimelineModel::cacheMedia(QString eventId)
+{
+        mtx::events::collections::TimelineEvents event = events.value(eventId);
+
+        if (auto e = boost::get<mtx::events::EncryptedEvent<mtx::events::msg::Encrypted>>(&event)) {
+                event = decryptEvent(*e).event;
+        }
+
+        QString mxcUrl =
+          boost::apply_visitor([](const auto &e) -> QString { return eventUrl(e); }, event);
+        QString mimeType =
+          boost::apply_visitor([](const auto &e) -> QString { return eventMimeType(e); }, event);
+
+        using EncF = boost::optional<mtx::crypto::EncryptedFile>;
+        EncF encryptionInfo =
+          boost::apply_visitor([](const auto &e) -> EncF { return eventEncryptionInfo(e); }, event);
+
+        // If the message is a link to a non mxcUrl, don't download it
+        if (!mxcUrl.startsWith("mxc://")) {
+                emit mediaCached(mxcUrl, mxcUrl);
+                return;
+        }
+
+        QString suffix = QMimeDatabase().mimeTypeForName(mimeType).preferredSuffix();
+
+        const auto url = mxcUrl.toStdString();
+        QFileInfo filename(QString("%1/media_cache/%2.%3")
+                             .arg(QStandardPaths::writableLocation(QStandardPaths::CacheLocation))
+                             .arg(QString(mxcUrl).remove("mxc://"))
+                             .arg(suffix));
+        if (QDir::cleanPath(filename.path()) != filename.path()) {
+                nhlog::net()->warn("mxcUrl '{}' is not safe, not downloading file", url);
+                return;
+        }
+
+        QDir().mkpath(filename.path());
+
+        if (filename.isReadable()) {
+                emit mediaCached(mxcUrl, filename.filePath());
+                return;
+        }
+
+        http::client()->download(
+          url,
+          [this, mxcUrl, filename, url, encryptionInfo](const std::string &data,
+                                                        const std::string &,
+                                                        const std::string &,
+                                                        mtx::http::RequestErr err) {
+                  if (err) {
+                          nhlog::net()->warn("failed to retrieve image {}: {} {}",
+                                             url,
+                                             err->matrix_error.error,
+                                             static_cast<int>(err->status_code));
+                          return;
+                  }
+
+                  try {
+                          auto temp = data;
+                          if (encryptionInfo)
+                                  temp = mtx::crypto::to_string(
+                                    mtx::crypto::decrypt_file(temp, encryptionInfo.value()));
+
+                          QFile file(filename.filePath());
+
+                          if (!file.open(QIODevice::WriteOnly))
+                                  return;
+
+                          file.write(QByteArray(temp.data(), temp.size()));
+                          file.close();
+                  } catch (const std::exception &e) {
+                          nhlog::ui()->warn("Error while saving file to: {}", e.what());
+                  }
+
+                  emit mediaCached(mxcUrl, filename.filePath());
+          });
+}
diff --git a/src/timeline/TimelineModel.h b/src/timeline/TimelineModel.h
index e7842b99..f52091e6 100644
--- a/src/timeline/TimelineModel.h
+++ b/src/timeline/TimelineModel.h
@@ -159,6 +159,8 @@ public:
         Q_INVOKABLE void redactEvent(QString id);
         Q_INVOKABLE int idToIndex(QString id) const;
         Q_INVOKABLE QString indexToId(int index) const;
+        Q_INVOKABLE void cacheMedia(QString eventId);
+        Q_INVOKABLE void saveMedia(QString eventId) const;
 
         void addEvents(const mtx::responses::Timeline &events);
         template<class T>
@@ -185,6 +187,7 @@ signals:
         void eventRedacted(QString id);
         void nextPendingMessage();
         void newMessageToSend(mtx::events::collections::TimelineEvents event);
+        void mediaCached(QString mxcUrl, QString cacheUrl);
 
 private:
         DecryptionResult decryptEvent(
diff --git a/src/timeline/TimelineViewManager.cpp b/src/timeline/TimelineViewManager.cpp
index 2a88c882..6430a426 100644
--- a/src/timeline/TimelineViewManager.cpp
+++ b/src/timeline/TimelineViewManager.cpp
@@ -1,11 +1,8 @@
 #include "TimelineViewManager.h"
 
-#include <QFileDialog>
 #include <QMetaType>
-#include <QMimeDatabase>
 #include <QPalette>
 #include <QQmlContext>
-#include <QStandardPaths>
 
 #include "ChatPage.h"
 #include "ColorImageProvider.h"
@@ -124,146 +121,24 @@ TimelineViewManager::setHistoryView(const QString &room_id)
 }
 
 void
-TimelineViewManager::openImageOverlay(QString mxcUrl,
-                                      QString originalFilename,
-                                      QString mimeType,
-                                      qml_mtx_events::EventType eventType) const
+TimelineViewManager::openImageOverlay(QString mxcUrl, QString eventId) const
 {
         QQuickImageResponse *imgResponse =
           imgProvider->requestImageResponse(mxcUrl.remove("mxc://"), QSize());
-        connect(imgResponse,
-                &QQuickImageResponse::finished,
-                this,
-                [this, mxcUrl, originalFilename, mimeType, eventType, imgResponse]() {
-                        if (!imgResponse->errorString().isEmpty()) {
-                                nhlog::ui()->error("Error when retrieving image for overlay: {}",
-                                                   imgResponse->errorString().toStdString());
-                                return;
-                        }
-                        auto pixmap = QPixmap::fromImage(imgResponse->textureFactory()->image());
-
-                        auto imgDialog = new dialogs::ImageOverlay(pixmap);
-                        imgDialog->show();
-                        connect(imgDialog,
-                                &dialogs::ImageOverlay::saving,
-                                this,
-                                [this, mxcUrl, originalFilename, mimeType, eventType]() {
-                                        saveMedia(mxcUrl, originalFilename, mimeType, eventType);
-                                });
-                });
-}
-
-void
-TimelineViewManager::saveMedia(QString mxcUrl,
-                               QString originalFilename,
-                               QString mimeType,
-                               qml_mtx_events::EventType eventType) const
-{
-        QString dialogTitle;
-        if (eventType == qml_mtx_events::EventType::ImageMessage) {
-                dialogTitle = tr("Save image");
-        } else if (eventType == qml_mtx_events::EventType::VideoMessage) {
-                dialogTitle = tr("Save video");
-        } else if (eventType == qml_mtx_events::EventType::AudioMessage) {
-                dialogTitle = tr("Save audio");
-        } else {
-                dialogTitle = tr("Save file");
-        }
-
-        QString filterString = QMimeDatabase().mimeTypeForName(mimeType).filterString();
-
-        auto filename =
-          QFileDialog::getSaveFileName(container, dialogTitle, originalFilename, filterString);
-
-        if (filename.isEmpty())
-                return;
-
-        const auto url = mxcUrl.toStdString();
-
-        http::client()->download(
-          url,
-          [filename, url](const std::string &data,
-                          const std::string &,
-                          const std::string &,
-                          mtx::http::RequestErr err) {
-                  if (err) {
-                          nhlog::net()->warn("failed to retrieve image {}: {} {}",
-                                             url,
-                                             err->matrix_error.error,
-                                             static_cast<int>(err->status_code));
-                          return;
-                  }
-
-                  try {
-                          QFile file(filename);
-
-                          if (!file.open(QIODevice::WriteOnly))
-                                  return;
-
-                          file.write(QByteArray(data.data(), (int)data.size()));
-                          file.close();
-                  } catch (const std::exception &e) {
-                          nhlog::ui()->warn("Error while saving file to: {}", e.what());
-                  }
-          });
-}
-
-void
-TimelineViewManager::cacheMedia(QString mxcUrl, QString mimeType)
-{
-        // If the message is a link to a non mxcUrl, don't download it
-        if (!mxcUrl.startsWith("mxc://")) {
-                emit mediaCached(mxcUrl, mxcUrl);
-                return;
-        }
-
-        QString suffix = QMimeDatabase().mimeTypeForName(mimeType).preferredSuffix();
-
-        const auto url = mxcUrl.toStdString();
-        QFileInfo filename(QString("%1/media_cache/%2.%3")
-                             .arg(QStandardPaths::writableLocation(QStandardPaths::CacheLocation))
-                             .arg(QString(mxcUrl).remove("mxc://"))
-                             .arg(suffix));
-        if (QDir::cleanPath(filename.path()) != filename.path()) {
-                nhlog::net()->warn("mxcUrl '{}' is not safe, not downloading file", url);
-                return;
-        }
-
-        QDir().mkpath(filename.path());
-
-        if (filename.isReadable()) {
-                emit mediaCached(mxcUrl, filename.filePath());
-                return;
-        }
+        connect(imgResponse, &QQuickImageResponse::finished, this, [this, eventId, imgResponse]() {
+                if (!imgResponse->errorString().isEmpty()) {
+                        nhlog::ui()->error("Error when retrieving image for overlay: {}",
+                                           imgResponse->errorString().toStdString());
+                        return;
+                }
+                auto pixmap = QPixmap::fromImage(imgResponse->textureFactory()->image());
 
-        http::client()->download(
-          url,
-          [this, mxcUrl, filename, url](const std::string &data,
-                                        const std::string &,
-                                        const std::string &,
-                                        mtx::http::RequestErr err) {
-                  if (err) {
-                          nhlog::net()->warn("failed to retrieve image {}: {} {}",
-                                             url,
-                                             err->matrix_error.error,
-                                             static_cast<int>(err->status_code));
-                          return;
-                  }
-
-                  try {
-                          QFile file(filename.filePath());
-
-                          if (!file.open(QIODevice::WriteOnly))
-                                  return;
-
-                          file.write(QByteArray(data.data(), data.size()));
-                          file.close();
-                  } catch (const std::exception &e) {
-                          nhlog::ui()->warn("Error while saving file to: {}", e.what());
-                  }
-
-                  emit mediaCached(mxcUrl, filename.filePath());
-          });
+                auto imgDialog = new dialogs::ImageOverlay(pixmap);
+                imgDialog->show();
+                connect(imgDialog, &dialogs::ImageOverlay::saving, timeline_, [this, eventId]() {
+                        timeline_->saveMedia(eventId);
+                });
+        });
 }
 
 void
@@ -401,3 +276,4 @@ TimelineViewManager::queueVideoMessage(const QString &roomid,
         video.url           = url.toStdString();
         models.value(roomid)->sendMessage(video);
 }
+
diff --git a/src/timeline/TimelineViewManager.h b/src/timeline/TimelineViewManager.h
index 0bc58e68..1cb0de44 100644
--- a/src/timeline/TimelineViewManager.h
+++ b/src/timeline/TimelineViewManager.h
@@ -35,38 +35,13 @@ public:
 
         Q_INVOKABLE TimelineModel *activeTimeline() const { return timeline_; }
         Q_INVOKABLE bool isInitialSync() const { return isInitialSync_; }
-        void openImageOverlay(QString mxcUrl,
-                              QString originalFilename,
-                              QString mimeType,
-                              qml_mtx_events::EventType eventType) const;
-        void saveMedia(QString mxcUrl,
-                       QString originalFilename,
-                       QString mimeType,
-                       qml_mtx_events::EventType eventType) const;
-        Q_INVOKABLE void cacheMedia(QString mxcUrl, QString mimeType);
-        // Qml can only pass enum as int
-        Q_INVOKABLE void openImageOverlay(QString mxcUrl,
-                                          QString originalFilename,
-                                          QString mimeType,
-                                          int eventType) const
-        {
-                openImageOverlay(
-                  mxcUrl, originalFilename, mimeType, (qml_mtx_events::EventType)eventType);
-        }
-        Q_INVOKABLE void saveMedia(QString mxcUrl,
-                                   QString originalFilename,
-                                   QString mimeType,
-                                   int eventType) const
-        {
-                saveMedia(mxcUrl, originalFilename, mimeType, (qml_mtx_events::EventType)eventType);
-        }
+        Q_INVOKABLE void openImageOverlay(QString mxcUrl, QString eventId) const;
 
 signals:
         void clearRoomMessageCount(QString roomid);
         void updateRoomsLastMessage(QString roomid, const DescInfo &info);
         void activeTimelineChanged(TimelineModel *timeline);
         void initialSyncChanged(bool isInitialSync);
-        void mediaCached(QString mxcUrl, QString cacheUrl);
 
 public slots:
         void updateReadReceipts(const QString &room_id, const std::vector<QString> &event_ids);