From e6069fa183aa4b8aa6c1e8afdff151f8d74c3fc5 Mon Sep 17 00:00:00 2001
From: Samuel Hand <samuel.hand@openmailbox.org>
Date: Sun, 29 Jul 2018 13:31:57 +0100
Subject: [PATCH] Mark messages from rejected and unknown devices

---
 plugins/omemo/src/stream_module.vala |  1 -
 plugins/omemo/src/trust_manager.vala | 39 ++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/plugins/omemo/src/stream_module.vala b/plugins/omemo/src/stream_module.vala
index 25e3685f..f086bd4c 100644
--- a/plugins/omemo/src/stream_module.vala
+++ b/plugins/omemo/src/stream_module.vala
@@ -33,7 +33,6 @@ public class StreamModule : XmppStreamModule {
         store_created(store);
         received_pipeline_listener = new ReceivedPipelineListener(store);
         stream.get_module(MessageModule.IDENTITY).received_pipeline.connect(received_pipeline_listener);
-        print("Adding filtered notification\n");
         stream.get_module(Pubsub.Module.IDENTITY).add_filtered_notification(stream, NODE_DEVICELIST, (stream, jid, id, node) => on_devicelist(stream, jid, id, node));
     }
 
diff --git a/plugins/omemo/src/trust_manager.vala b/plugins/omemo/src/trust_manager.vala
index 7e17e5ea..ca505c3f 100644
--- a/plugins/omemo/src/trust_manager.vala
+++ b/plugins/omemo/src/trust_manager.vala
@@ -10,10 +10,14 @@ public class TrustManager {
 
     private StreamInteractor stream_interactor;
     private Database db;
+    private ReceivedMessageListener received_message_listener;
 
     public TrustManager(StreamInteractor stream_interactor, Database db) {
         this.stream_interactor = stream_interactor;
         this.db = db;
+
+        received_message_listener = new ReceivedMessageListener(stream_interactor, db);
+        stream_interactor.get_module(MessageProcessor.IDENTITY).received_pipeline.connect(received_message_listener);
     }
 
     private StanzaNode create_encrypted_key(uint8[] key, Address address, Store store) throws GLib.Error {
@@ -124,6 +128,41 @@ public class TrustManager {
         }
         return devices;
     }
+
+    private class ReceivedMessageListener : MessageListener {
+        public string[] after_actions_const = new string[]{ };
+        public override string action_group { get { return "DECRYPT"; } }
+        public override string[] after_actions { get { return after_actions_const; } }
+
+        private StreamInteractor stream_interactor;
+        private Database db;
+
+        public ReceivedMessageListener(StreamInteractor stream_interactor, Database db) {
+            this.stream_interactor = stream_interactor;
+            this.db = db;
+        }
+
+        public override async bool run(Entities.Message message, Xmpp.MessageStanza stanza, Conversation conversation) {
+            MessageFlag? flag = MessageFlag.get_flag(stanza);
+            if(flag != null && ((!)flag).decrypted) {
+                StanzaNode header = stanza.stanza.get_subnode("encrypted", "eu.siacs.conversations.axolotl").get_subnode("header");
+                Jid jid = message.from;
+                if(conversation.type_ == Conversation.Type.GROUPCHAT) {
+                    jid = stream_interactor.get_module(MucManager.IDENTITY).get_real_jid(jid, conversation.account);
+                }
+                Database.IdentityMetaTable.TrustLevel trust_level = (Database.IdentityMetaTable.TrustLevel) db.identity_meta.with_address(conversation.account.id, jid.bare_jid.to_string()).with(db.identity_meta.device_id, "=", header.get_attribute_int("sid")).single()[db.identity_meta.trust_level];
+                if (trust_level == Database.IdentityMetaTable.TrustLevel.UNTRUSTED) {
+                    message.body = "OMEMO message from a rejected device";
+                    message.marked = Message.Marked.WONTSEND;
+                }
+                if (trust_level == Database.IdentityMetaTable.TrustLevel.UNKNOWN) {
+                    message.body = "OMEMO message from an unknown device: "+message.body;
+                    message.marked = Message.Marked.WONTSEND;
+                }
+            }
+            return false;
+        }
+    }
 }
 
 }