3 files changed, 28 insertions, 22 deletions

diff --git a/src/ui/model/GpgKeyTableProxyModel.cpp b/src/ui/model/GpgKeyTableProxyModel.cpp
index 746f3344..90bf16d3 100644
--- a/src/ui/model/GpgKeyTableProxyModel.cpp
+++ b/src/ui/model/GpgKeyTableProxyModel.cpp
@@ -28,9 +28,10 @@
 
 #include "GpgKeyTableProxyModel.h"
 
-#include "core/function/gpg/GpgKeyGetter.h"
+#include "core/function/gpg/GpgAbstractKeyGetter.h"
 #include "core/model/CacheObject.h"
 #include "core/model/GpgKey.h"
+#include "core/model/GpgKeyTableModel.h"
 #include "core/struct/cache_object/AllFavoriteKeyPairsCO.h"
 #include "core/utils/GpgUtils.h"
 
@@ -57,34 +58,31 @@ GpgKeyTableProxyModel::GpgKeyTableProxyModel(
 }
 
 auto GpgKeyTableProxyModel::filterAcceptsRow(
-    int source_row, const QModelIndex &sourceParent) const -> bool {
-  auto index = sourceModel()->index(source_row, 6, sourceParent);
-  auto key_id = sourceModel()->data(index).toString();
+    int sourceRow, const QModelIndex &sourceParent) const -> bool {
+  auto index = sourceModel()->index(sourceRow, 0, sourceParent);
 
-  auto key =
-      GpgKeyGetter::GetInstance(model_->GetGpgContextChannel()).GetKey(key_id);
-  LOG_D() << "get key: " << key_id
-          << "from channel: " << model_->GetGpgContextChannel();
-  assert(key.IsGood());
+  auto *i = index.isValid()
+                ? static_cast<GpgKeyTableItem *>(index.internalPointer())
+                : nullptr;
+  assert(i != nullptr);
+
+  auto *key = i->Key();
+  assert(key->IsGood());
 
   if (!(display_mode_ & GpgKeyTableDisplayMode::kPRIVATE_KEY) &&
-      key.IsPrivateKey()) {
+      key->IsPrivateKey()) {
     return false;
   }
 
   if (!(display_mode_ & GpgKeyTableDisplayMode::kPUBLIC_KEY) &&
-      !key.IsPrivateKey()) {
+      !key->IsPrivateKey()) {
     return false;
   }
 
   if (!custom_filter_(key)) return false;
 
-  if (display_mode_ & GpgKeyTableDisplayMode::kFAVORITES) {
-    LOG_D() << "kFAVORITES Mode" << "key id" << key_id << "favorite_key_ids_"
-            << favorite_key_ids_;
-  }
   if (display_mode_ & GpgKeyTableDisplayMode::kFAVORITES &&
-      !favorite_key_ids_.contains(key_id)) {
+      !favorite_key_ids_.contains(key->ID())) {
     return false;
   }
 
@@ -92,11 +90,14 @@ auto GpgKeyTableProxyModel::filterAcceptsRow(
 
   QStringList infos;
   for (int column = 0; column < sourceModel()->columnCount(); ++column) {
-    auto index = sourceModel()->index(source_row, column, sourceParent);
+    auto index = sourceModel()->index(sourceRow, column, sourceParent);
     infos << sourceModel()->data(index).toString();
 
-    for (const auto &uid : key.UIDs()) {
-      infos << uid.GetUID();
+    if (key->KeyType() == GpgAbstractKeyType::kGPG_KEY) {
+      auto *k = dynamic_cast<GpgKey *>(key);
+      for (const auto &uid : k->UIDs()) {
+        infos << uid.GetUID();
+      }
     }
   }
 
@@ -197,4 +198,8 @@ void GpgKeyTableProxyModel::slot_update_favorites_cache() {
   }
 }
 
+void GpgKeyTableProxyModel::SetFilter(const KeyFilter &filter) {
+  this->custom_filter_ = filter;
+  invalidateFilter();
+}
 }  // namespace GpgFrontend::UI
\ No newline at end of file
diff --git a/src/ui/model/GpgKeyTableProxyModel.h b/src/ui/model/GpgKeyTableProxyModel.h
index bad1aa97..5b9c0336 100644
--- a/src/ui/model/GpgKeyTableProxyModel.h
+++ b/src/ui/model/GpgKeyTableProxyModel.h
@@ -38,7 +38,7 @@ namespace GpgFrontend::UI {
 class GpgKeyTableProxyModel : public QSortFilterProxyModel {
   Q_OBJECT
  public:
-  using KeyFilter = std::function<bool(const GpgKey &)>;
+  using KeyFilter = std::function<bool(const GpgAbstractKey *)>;
 
   explicit GpgKeyTableProxyModel(QSharedPointer<GpgKeyTableModel> model,
                                  GpgKeyTableDisplayMode display_mode,
@@ -47,6 +47,8 @@ class GpgKeyTableProxyModel : public QSortFilterProxyModel {
 
   void SetSearchKeywords(const QString &keywords);
 
+  void SetFilter(const KeyFilter &filter);
+
   void ResetGpgKeyTableModel(QSharedPointer<GpgKeyTableModel> model);
 
  protected:
diff --git a/src/ui/model/GpgKeyTreeProxyModel.cpp b/src/ui/model/GpgKeyTreeProxyModel.cpp
index 2f46f548..bff18dd6 100644
--- a/src/ui/model/GpgKeyTreeProxyModel.cpp
+++ b/src/ui/model/GpgKeyTreeProxyModel.cpp
@@ -28,7 +28,6 @@
 
 #include "GpgKeyTreeProxyModel.h"
 
-#include "core/function/gpg/GpgKeyGetter.h"
 #include "core/model/CacheObject.h"
 #include "core/model/GpgKey.h"
 #include "core/model/GpgKeyTreeModel.h"
@@ -83,7 +82,7 @@ auto GpgKeyTreeProxyModel::filterAcceptsRow(
     auto index = sourceModel()->index(source_row, column, sourceParent);
     infos << sourceModel()->data(index).toString();
 
-    if (!key->IsSubKey()) {
+    if (key->KeyType() != GpgAbstractKeyType::kGPG_SUBKEY) {
       for (const auto &uid : dynamic_cast<const GpgKey *>(key)->UIDs()) {
         infos << uid.GetUID();
       }