Split OMEMO plug-in into files, various fixes

14 files changed, 511 insertions, 433 deletions

diff --git a/plugins/omemo/src/account_settings_entry.vala b/plugins/omemo/src/account_settings_entry.vala
new file mode 100644
index 00000000..c6871f6e
--- /dev/null
+++ b/plugins/omemo/src/account_settings_entry.vala
@@ -0,0 +1,23 @@
+namespace Dino.Plugins.Omemo {
+
+public class AccountSettingsEntry : Plugins.AccountSettingsEntry {
+    private Plugin plugin;
+
+    public AccountSettingsEntry(Plugin plugin) {
+        this.plugin = plugin;
+    }
+
+    public override string id { get {
+        return "omemo_identity_key";
+    }}
+
+    public override string name { get {
+        return "OMEMO";
+    }}
+
+    public override Plugins.AccountSettingsWidget get_widget() {
+        return new AccountSettingWidget(plugin);
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/account_settings_widget.vala b/plugins/omemo/src/account_settings_widget.vala
new file mode 100644
index 00000000..87ea0e37
--- /dev/null
+++ b/plugins/omemo/src/account_settings_widget.vala
@@ -0,0 +1,63 @@
+using Gtk;
+using Dino.Entities;
+
+namespace Dino.Plugins.Omemo {
+
+public class AccountSettingWidget : Plugins.AccountSettingsWidget, Box {
+    private Plugin plugin;
+    private Label fingerprint;
+    private Account account;
+
+    public AccountSettingWidget(Plugin plugin) {
+        this.plugin = plugin;
+
+        fingerprint = new Label("...");
+        fingerprint.xalign = 0;
+        Border border = new Button().get_style_context().get_padding(StateFlags.NORMAL);
+        fingerprint.set_padding(border.left + 1, border.top + 1);
+        fingerprint.visible = true;
+        pack_start(fingerprint);
+
+        Button btn = new Button();
+        btn.image = new Image.from_icon_name("view-list-symbolic", IconSize.BUTTON);
+        btn.relief = ReliefStyle.NONE;
+        btn.visible = true;
+        btn.valign = Align.CENTER;
+        btn.clicked.connect(() => { activated(); });
+        pack_start(btn, false);
+    }
+
+    public void set_account(Account account) {
+        this.account = account;
+        try {
+            Qlite.Row? row = plugin.db.identity.row_with(plugin.db.identity.account_id, account.id);
+            if (row == null) {
+                fingerprint.set_markup(@"Own fingerprint\n<span font='8'>Will be generated on first connect</span>");
+            } else {
+                uint8[] arr = Base64.decode(row[plugin.db.identity.identity_key_public_base64]);
+                arr = arr[1:arr.length];
+                string res = "";
+                foreach (uint8 i in arr) {
+                    string s = i.to_string("%x");
+                    if (s.length == 1) s = "0" + s;
+                    res = res + s;
+                    if ((res.length % 9) == 8) {
+                        if (res.length == 35) {
+                            res += "\n";
+                        } else {
+                            res += " ";
+                        }
+                    }
+                }
+                fingerprint.set_markup(@"Own fingerprint\n<span font_family='monospace' font='8'>$res</span>");
+            }
+        } catch (Qlite.DatabaseError e) {
+            fingerprint.set_markup(@"Own fingerprint\n<span font='8'>Database error</span>");
+        }
+    }
+
+    public void deactivate() {
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/bundle.vala b/plugins/omemo/src/bundle.vala
new file mode 100644
index 00000000..211dc29b
--- /dev/null
+++ b/plugins/omemo/src/bundle.vala
@@ -0,0 +1,87 @@
+using Gee;
+using Signal;
+using Xmpp.Core;
+
+namespace Dino.Plugins.Omemo {
+
+public class Bundle {
+    private StanzaNode? node;
+
+    public Bundle(StanzaNode? node) {
+        this.node = node;
+    }
+
+    public int32 signed_pre_key_id { owned get {
+        if (node == null) return -1;
+        string id = node.get_deep_attribute("signedPreKeyPublic", "signedPreKeyId");
+        if (id == null) return -1;
+        return int.parse(id);
+    }}
+
+    public ECPublicKey? signed_pre_key { owned get {
+        if (node == null) return null;
+        string? key = node.get_deep_string_content("signedPreKeyPublic");
+        if (key == null) return null;
+        try {
+            return Plugin.context.decode_public_key(Base64.decode(key));
+        } catch (Error e) {
+            return null;
+        }
+    }}
+
+    public uint8[]? signed_pre_key_signature { owned get {
+        if (node == null) return null;
+        string? sig = node.get_deep_string_content("signedPreKeySignature");
+        if (sig == null) return null;
+        return Base64.decode(sig);
+    }}
+
+    public ECPublicKey? identity_key { owned get {
+        if (node == null) return null;
+        string? key = node.get_deep_string_content("identityKey");
+        if (key == null) return null;
+        try {
+            return Plugin.context.decode_public_key(Base64.decode(key));
+        } catch (Error e) {
+            return null;
+        }
+    }}
+
+    public ArrayList<PreKey> pre_keys { owned get {
+        ArrayList<PreKey> list = new ArrayList<PreKey>();
+        if (node == null || node.get_subnode("prekeys") == null) return list;
+        node.get_deep_subnodes("prekeys", "preKeyPublic")
+                .filter((node) => node.get_attribute("preKeyId") != null)
+                .map<PreKey>(PreKey.create)
+                .foreach((key) => list.add(key));
+        return list;
+    }}
+
+    public class PreKey {
+        private StanzaNode node;
+
+        public static PreKey create(owned StanzaNode node) {
+            return new PreKey(node);
+        }
+
+        public PreKey(StanzaNode node) {
+            this.node = node;
+        }
+
+        public int32 key_id { owned get {
+            return int.parse(node.get_attribute("preKeyId") ?? "-1");
+        }}
+
+        public ECPublicKey? key { owned get {
+            string? key = node.get_string_content();
+            if (key == null) return null;
+            try {
+                return Plugin.context.decode_public_key(Base64.decode(key));
+            } catch (Error e) {
+                return null;
+            }
+        }}
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/database.vala b/plugins/omemo/src/database.vala
index 1216ca84..db530c69 100644
--- a/plugins/omemo/src/database.vala
+++ b/plugins/omemo/src/database.vala
@@ -4,7 +4,7 @@ using Qlite;
 
 using Dino.Entities;
 
-namespace Dino.Omemo {
+namespace Dino.Plugins.Omemo {
 
 public class Database : Qlite.Database {
     private const int VERSION = 0;
@@ -63,7 +63,7 @@ public class Database : Qlite.Database {
     public PreKeyTable pre_key { get; private set; }
     public SessionTable session { get; private set; }
 
-    public Database(string fileName) {
+    public Database(string fileName) throws DatabaseError {
         base(fileName, VERSION);
         identity = new IdentityTable(this);
         signed_pre_key = new SignedPreKeyTable(this);
diff --git a/plugins/omemo/src/encrypt_status.vala b/plugins/omemo/src/encrypt_status.vala
new file mode 100644
index 00000000..c6b45ac6
--- /dev/null
+++ b/plugins/omemo/src/encrypt_status.vala
@@ -0,0 +1,17 @@
+namespace Dino.Plugins.Omemo {
+
+public class EncryptStatus {
+    public bool encrypted { get; internal set; }
+    public int other_devices { get; internal set; }
+    public int other_success { get; internal set; }
+    public int other_lost { get; internal set; }
+    public int other_unknown { get; internal set; }
+    public int other_failure { get; internal set; }
+    public int own_devices { get; internal set; }
+    public int own_success { get; internal set; }
+    public int own_lost { get; internal set; }
+    public int own_unknown { get; internal set; }
+    public int own_failure { get; internal set; }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/encryption_list_entry.vala b/plugins/omemo/src/encryption_list_entry.vala
new file mode 100644
index 00000000..753ffe67
--- /dev/null
+++ b/plugins/omemo/src/encryption_list_entry.vala
@@ -0,0 +1,23 @@
+namespace Dino.Plugins.Omemo {
+
+public class EncryptionListEntry : Plugins.EncryptionListEntry, Object {
+    private Plugin plugin;
+
+    public EncryptionListEntry(Plugin plugin) {
+        this.plugin = plugin;
+    }
+
+    public Entities.Encryption encryption { get {
+        return Entities.Encryption.OMEMO;
+    }}
+
+    public string name { get {
+        return "OMEMO";
+    }}
+
+    public bool can_encrypt(Entities.Conversation conversation) {
+        return Manager.get_instance(plugin.app.stream_interaction).can_encrypt(conversation);
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/manager.vala b/plugins/omemo/src/manager.vala
index 69a69d9c..e5db631e 100644
--- a/plugins/omemo/src/manager.vala
+++ b/plugins/omemo/src/manager.vala
@@ -4,7 +4,7 @@ using Qlite;
 using Xmpp;
 using Gee;
 
-namespace Dino.Omemo {
+namespace Dino.Plugins.Omemo {
 
 public class Manager : StreamInteractionModule, Object {
     public const string id = "omemo_manager";
@@ -31,7 +31,7 @@ public class Manager : StreamInteractionModule, Object {
 
     private void on_pre_message_send(Entities.Message message, Xmpp.Message.Stanza message_stanza, Conversation conversation) {
         if (message.encryption == Encryption.OMEMO) {
-            Module module = Module.get_module(stream_interactor.get_stream(conversation.account));
+            StreamModule module = stream_interactor.get_stream(conversation.account).get_module(StreamModule.IDENTITY);
             EncryptStatus status = module.encrypt(message_stanza, conversation.account.bare_jid.to_string());
             if (status.other_failure > 0 || (status.other_lost == status.other_devices && status.other_devices > 0)) {
                 message.marked = Entities.Message.Marked.WONTSEND;
@@ -63,9 +63,9 @@ public class Manager : StreamInteractionModule, Object {
     }
 
     private void on_account_added(Account account) {
-        stream_interactor.module_manager.get_module(account, Module.IDENTITY).store_created.connect((context, store) => on_store_created(account, context, store));
-        stream_interactor.module_manager.get_module(account, Module.IDENTITY).device_list_loaded.connect(() => on_device_list_loaded(account));
-        stream_interactor.module_manager.get_module(account, Module.IDENTITY).session_started.connect((jid, device_id) => on_session_started(account, jid));
+        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).store_created.connect((store) => on_store_created(account, store));
+        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).device_list_loaded.connect(() => on_device_list_loaded(account));
+        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).session_started.connect((jid, device_id) => on_session_started(account, jid));
     }
 
     private void on_session_started(Account account, string jid) {
@@ -96,7 +96,7 @@ public class Manager : StreamInteractionModule, Object {
         }
     }
 
-    private void on_store_created(Account account, Context context, Store store) {
+    private void on_store_created(Account account, Store store) {
         Qlite.Row? row = null;
         try {
             row = db.identity.row_with(db.identity.account_id, account.id);
@@ -107,19 +107,19 @@ public class Manager : StreamInteractionModule, Object {
 
         if (row == null) {
             // OMEMO not yet initialized, starting with empty base
-            store.identity_key_store.local_registration_id = Random.int_range(1, int32.MAX);
+            try {
+                store.identity_key_store.local_registration_id = Random.int_range(1, int32.MAX);
 
-            Signal.ECKeyPair key_pair = context.generate_key_pair();
-            store.identity_key_store.identity_key_private = key_pair.private.serialize();
-            store.identity_key_store.identity_key_public = key_pair.public.serialize();
+                Signal.ECKeyPair key_pair = Plugin.context.generate_key_pair();
+                store.identity_key_store.identity_key_private = key_pair.private.serialize();
+                store.identity_key_store.identity_key_public = key_pair.public.serialize();
 
-            try {
                 identity_id = (int) db.identity.insert().or("REPLACE")
-                .value(db.identity.account_id, account.id)
-                .value(db.identity.device_id, (int) store.local_registration_id)
-                .value(db.identity.identity_key_private_base64, Base64.encode(store.identity_key_store.identity_key_private))
-                .value(db.identity.identity_key_public_base64, Base64.encode(store.identity_key_store.identity_key_public))
-                .perform();
+                        .value(db.identity.account_id, account.id)
+                        .value(db.identity.device_id, (int) store.local_registration_id)
+                        .value(db.identity.identity_key_private_base64, Base64.encode(store.identity_key_store.identity_key_private))
+                        .value(db.identity.identity_key_public_base64, Base64.encode(store.identity_key_store.identity_key_public))
+                        .perform();
             } catch (Error e) {
                 // Ignore error
             }
@@ -139,118 +139,9 @@ public class Manager : StreamInteractionModule, Object {
         }
     }
 
-    private class BackedSignedPreKeyStore : SimpleSignedPreKeyStore {
-        private Database db;
-        private int identity_id;
-
-        public BackedSignedPreKeyStore(Database db, int identity_id) {
-            this.db = db;
-            this.identity_id = identity_id;
-            init();
-        }
-
-        private void init() {
-            foreach (Row row in db.signed_pre_key.select().with(db.signed_pre_key.identity_id, "=", identity_id)) {
-                store_signed_pre_key(row[db.signed_pre_key.signed_pre_key_id], Base64.decode(row[db.signed_pre_key.record_base64]));
-            }
-
-            signed_pre_key_stored.connect(on_signed_pre_key_stored);
-            signed_pre_key_deleted.connect(on_signed_pre_key_deleted);
-        }
-
-        public void on_signed_pre_key_stored(SignedPreKeyStore.Key key) {
-            db.signed_pre_key.insert().or("REPLACE")
-                .value(db.signed_pre_key.identity_id, identity_id)
-                .value(db.signed_pre_key.signed_pre_key_id, (int) key.key_id)
-                .value(db.signed_pre_key.record_base64, Base64.encode(key.record))
-                .perform();
-        }
-
-        public void on_signed_pre_key_deleted(SignedPreKeyStore.Key key) {
-            db.signed_pre_key.delete()
-                .with(db.signed_pre_key.identity_id, "=", identity_id)
-                .with(db.signed_pre_key.signed_pre_key_id, "=", (int) key.key_id)
-                .perform();
-        }
-    }
-
-    private class BackedPreKeyStore : SimplePreKeyStore {
-        private Database db;
-        private int identity_id;
-
-        public BackedPreKeyStore(Database db, int identity_id) {
-            this.db = db;
-            this.identity_id = identity_id;
-            init();
-        }
-
-        private void init() {
-            foreach (Row row in db.pre_key.select().with(db.pre_key.identity_id, "=", identity_id)) {
-                store_pre_key(row[db.pre_key.pre_key_id], Base64.decode(row[db.pre_key.record_base64]));
-            }
-
-            pre_key_stored.connect(on_pre_key_stored);
-            pre_key_deleted.connect(on_pre_key_deleted);
-        }
-
-        public void on_pre_key_stored(PreKeyStore.Key key) {
-            db.pre_key.insert().or("REPLACE")
-                .value(db.pre_key.identity_id, identity_id)
-                .value(db.pre_key.pre_key_id, (int) key.key_id)
-                .value(db.pre_key.record_base64, Base64.encode(key.record))
-                .perform();
-        }
-
-        public void on_pre_key_deleted(PreKeyStore.Key key) {
-            db.pre_key.delete()
-                .with(db.pre_key.identity_id, "=", identity_id)
-                .with(db.pre_key.pre_key_id, "=", (int) key.key_id)
-                .perform();
-        }
-    }
-
-    private class BackedSessionStore : SimpleSessionStore {
-        private Database db;
-        private int identity_id;
-
-        public BackedSessionStore(Database db, int identity_id) {
-            this.db = db;
-            this.identity_id = identity_id;
-            init();
-        }
-
-        private void init() {
-            Address addr = new Address();
-            foreach (Row row in db.session.select().with(db.session.identity_id, "=", identity_id)) {
-                addr.name = row[db.session.address_name];
-                addr.device_id = row[db.session.device_id];
-                store_session(addr, Base64.decode(row[db.session.record_base64]));
-            }
-
-            session_stored.connect(on_session_stored);
-            session_removed.connect(on_session_deleted);
-        }
-
-        public void on_session_stored(SessionStore.Session session) {
-            db.session.insert().or("REPLACE")
-                .value(db.session.identity_id, identity_id)
-                .value(db.session.address_name, session.name)
-                .value(db.session.device_id, session.device_id)
-                .value(db.session.record_base64, Base64.encode(session.record))
-                .perform();
-        }
-
-        public void on_session_deleted(SessionStore.Session session) {
-            db.session.delete()
-                .with(db.session.identity_id, "=", identity_id)
-                .with(db.session.address_name, "=", session.name)
-                .with(db.session.device_id, "=", session.device_id)
-                .perform();
-        }
-    }
 
-    public bool con_encrypt(Entities.Conversation conversation) {
-        return true; // TODO
+    public bool can_encrypt(Entities.Conversation conversation) {
+        return stream_interactor.get_stream(conversation.account).get_module(StreamModule.IDENTITY).is_known_address(conversation.counterpart.bare_jid.to_string());
     }
 
     internal string get_id() {
diff --git a/plugins/omemo/src/message_flag.vala b/plugins/omemo/src/message_flag.vala
new file mode 100644
index 00000000..cea1e9b2
--- /dev/null
+++ b/plugins/omemo/src/message_flag.vala
@@ -0,0 +1,23 @@
+using Xmpp;
+
+namespace Dino.Plugins.Omemo {
+
+public class MessageFlag : Message.MessageFlag {
+    public const string id = "omemo";
+
+    public bool decrypted = false;
+
+    public static MessageFlag? get_flag(Message.Stanza message) {
+        return (MessageFlag) message.get_flag(NS_URI, id);
+    }
+
+    public override string get_ns() {
+        return NS_URI;
+    }
+
+    public override string get_id() {
+        return id;
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/plugin.vala b/plugins/omemo/src/plugin.vala
index a062640b..04e02625 100644
--- a/plugins/omemo/src/plugin.vala
+++ b/plugins/omemo/src/plugin.vala
@@ -1,130 +1,34 @@
-using Xmpp;
+namespace Dino.Plugins.Omemo {
 
-namespace Dino.Omemo {
+public class Plugin : RootInterface, Object {
+    public static Signal.Context context;
 
-    public class EncryptionListEntry : Plugins.EncryptionListEntry, Object {
-        private Plugin plugin;
+    public Dino.Application app;
+    public Database db;
+    public EncryptionListEntry list_entry;
+    public AccountSettingsEntry settings_entry;
 
-        public EncryptionListEntry(Plugin plugin) {
-            this.plugin = plugin;
-        }
-
-        public Entities.Encryption encryption { get {
-            return Entities.Encryption.OMEMO;
-        }}
-
-        public string name { get {
-            return "OMEMO";
-        }}
-
-        public bool can_encrypt(Entities.Conversation conversation) {
-            return Manager.get_instance(plugin.app.stream_interaction).con_encrypt(conversation);
-        }
-    }
-
-    public class AccountSettingsEntry : Plugins.AccountSettingsEntry {
-        private Plugin plugin;
-
-        public AccountSettingsEntry(Plugin plugin) {
-            this.plugin = plugin;
-        }
-
-        public override string id { get {
-            return "omemo_identity_key";
-        }}
-
-        public override string name { get {
-            return "OMEMO";
-        }}
-
-        public override Plugins.AccountSettingsWidget get_widget() {
-            return new AccountSettingWidget(plugin);
-        }
-    }
-
-    public class AccountSettingWidget : Plugins.AccountSettingsWidget, Gtk.Box {
-        private Plugin plugin;
-        private Gtk.Label fingerprint;
-        private Entities.Account account;
-
-        public AccountSettingWidget(Plugin plugin) {
-            this.plugin = plugin;
-
-            fingerprint = new Gtk.Label("...");
-            fingerprint.xalign = 0;
-            Gtk.Border border = new Gtk.Button().get_style_context().get_padding(Gtk.StateFlags.NORMAL);
-            fingerprint.set_padding(border.left + 1, border.top + 1);
-            fingerprint.visible = true;
-            pack_start(fingerprint);
-
-            Gtk.Button btn = new Gtk.Button();
-            btn.image = new Gtk.Image.from_icon_name("view-list-symbolic", Gtk.IconSize.BUTTON);
-            btn.relief = Gtk.ReliefStyle.NONE;
-            btn.visible = true;
-            btn.valign = Gtk.Align.CENTER;
-            btn.clicked.connect(() => { activated(); });
-            pack_start(btn, false);
-        }
-
-        public void set_account(Entities.Account account) {
-            this.account = account;
-            try {
-                Qlite.Row? row = plugin.db.identity.row_with(plugin.db.identity.account_id, account.id);
-                if (row == null) {
-                    fingerprint.set_markup(@"Own fingerprint\n<span font='8'>Will be generated on first connect</span>");
-                } else {
-                    uint8[] arr = Base64.decode(row[plugin.db.identity.identity_key_public_base64]);
-                    arr = arr[1:arr.length];
-                    string res = "";
-                    foreach (uint8 i in arr) {
-                        string s = i.to_string("%x");
-                        if (s.length == 1) s = "0" + s;
-                        res = res + s;
-                        if ((res.length % 9) == 8) {
-                            if (res.length == 35) {
-                                res += "\n";
-                            } else {
-                                res += " ";
-                            }
-                        }
-                    }
-                    fingerprint.set_markup(@"Own fingerprint\n<span font_family='monospace' font='8'>$res</span>");
-                }
-            } catch (Qlite.DatabaseError e) {
-                fingerprint.set_markup(@"Own fingerprint\n<span font='8'>Database error</span>");
-            }
-        }
-
-        public void deactivate() {
-        }
-    }
-
-    public class Plugin : Plugins.RootInterface, Object {
-        public Dino.Application app;
-        public Database db;
-        public EncryptionListEntry list_entry;
-        public AccountSettingsEntry settings_entry;
-
-        public void registered(Dino.Application app) {
+    public void registered(Dino.Application app) {
+        try {
+            context = new Signal.Context(false);
             this.app = app;
             this.db = new Database("omemo.db");
             this.list_entry = new EncryptionListEntry(this);
             this.settings_entry = new AccountSettingsEntry(this);
-            app.plugin_registry.register_encryption_list_entry(list_entry);
-            app.plugin_registry.register_account_settings_entry(settings_entry);
-            app.stream_interaction.module_manager.initialize_account_modules.connect((account, list) => {
-                list.add(new Module());
+            this.app.plugin_registry.register_encryption_list_entry(list_entry);
+            this.app.plugin_registry.register_account_settings_entry(settings_entry);
+            this.app.stream_interaction.module_manager.initialize_account_modules.connect((account, list) => {
+                list.add(new StreamModule());
             });
-            Manager.start(app.stream_interaction, db);
-        }
-
-        public void shutdown() {
-            // Nothing to do
+            Manager.start(this.app.stream_interaction, db);
+        } catch (Error e) {
+            print(@"Error initializing OMEMO: $(e.message)\n");
         }
     }
 
+    public void shutdown() {
+        // Nothing to do
+    }
 }
 
-public Type register_plugin(Module module) {
-    return typeof (Dino.Omemo.Plugin);
-}
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/pre_key_store.vala b/plugins/omemo/src/pre_key_store.vala
new file mode 100644
index 00000000..0fd78ffc
--- /dev/null
+++ b/plugins/omemo/src/pre_key_store.vala
@@ -0,0 +1,53 @@
+using Signal;
+using Qlite;
+
+namespace Dino.Plugins.Omemo {
+
+private class BackedPreKeyStore : SimplePreKeyStore {
+    private Database db;
+    private int identity_id;
+
+    public BackedPreKeyStore(Database db, int identity_id) {
+        this.db = db;
+        this.identity_id = identity_id;
+        init();
+    }
+
+    private void init() {
+        try {
+            foreach (Row row in db.pre_key.select().with(db.pre_key.identity_id, "=", identity_id)) {
+                store_pre_key(row[db.pre_key.pre_key_id], Base64.decode(row[db.pre_key.record_base64]));
+            }
+        } catch (Error e) {
+            print(@"OMEMO: Error while initializing pre key store: $(e.message)\n");
+        }
+
+        pre_key_stored.connect(on_pre_key_stored);
+        pre_key_deleted.connect(on_pre_key_deleted);
+    }
+
+    public void on_pre_key_stored(PreKeyStore.Key key) {
+        try {
+            db.pre_key.insert().or("REPLACE")
+                    .value(db.pre_key.identity_id, identity_id)
+                    .value(db.pre_key.pre_key_id, (int) key.key_id)
+                    .value(db.pre_key.record_base64, Base64.encode(key.record))
+                    .perform();
+        } catch (Error e) {
+            print(@"OMEMO: Error while updating pre key store: $(e.message)\n");
+        }
+    }
+
+    public void on_pre_key_deleted(PreKeyStore.Key key) {
+        try {
+            db.pre_key.delete()
+                    .with(db.pre_key.identity_id, "=", identity_id)
+                    .with(db.pre_key.pre_key_id, "=", (int) key.key_id)
+                    .perform();
+        } catch (Error e) {
+            print(@"OMEMO: Error while updating pre key store: $(e.message)\n");
+        }
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/register_plugin.vala b/plugins/omemo/src/register_plugin.vala
new file mode 100644
index 00000000..0d0e1c3e
--- /dev/null
+++ b/plugins/omemo/src/register_plugin.vala
@@ -0,0 +1,3 @@
+public Type register_plugin(Module module) {
+    return typeof (Dino.Plugins.Omemo.Plugin);
+}
diff --git a/plugins/omemo/src/session_store.vala b/plugins/omemo/src/session_store.vala
new file mode 100644
index 00000000..f70e16ea
--- /dev/null
+++ b/plugins/omemo/src/session_store.vala
@@ -0,0 +1,58 @@
+using Signal;
+using Qlite;
+
+namespace Dino.Plugins.Omemo {
+
+private class BackedSessionStore : SimpleSessionStore {
+    private Database db;
+    private int identity_id;
+
+    public BackedSessionStore(Database db, int identity_id) {
+        this.db = db;
+        this.identity_id = identity_id;
+        init();
+    }
+
+    private void init() {
+        try {
+            Address addr = new Address();
+            foreach (Row row in db.session.select().with(db.session.identity_id, "=", identity_id)) {
+                addr.name = row[db.session.address_name];
+                addr.device_id = row[db.session.device_id];
+                store_session(addr, Base64.decode(row[db.session.record_base64]));
+            }
+        } catch (Error e) {
+            print(@"OMEMO: Error while initializing session store: $(e.message)\n");
+        }
+
+        session_stored.connect(on_session_stored);
+        session_removed.connect(on_session_deleted);
+    }
+
+    public void on_session_stored(SessionStore.Session session) {
+        try {
+            db.session.insert().or("REPLACE")
+                    .value(db.session.identity_id, identity_id)
+                    .value(db.session.address_name, session.name)
+                    .value(db.session.device_id, session.device_id)
+                    .value(db.session.record_base64, Base64.encode(session.record))
+                    .perform();
+        } catch (Error e) {
+            print(@"OMEMO: Error while updating session store: $(e.message)\n");
+        }
+    }
+
+    public void on_session_deleted(SessionStore.Session session) {
+        try {
+            db.session.delete()
+                    .with(db.session.identity_id, "=", identity_id)
+                    .with(db.session.address_name, "=", session.name)
+                    .with(db.session.device_id, "=", session.device_id)
+                    .perform();
+        } catch (Error e) {
+            print(@"OMEMO: Error while updating session store: $(e.message)\n");
+        }
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/signed_pre_key_store.vala b/plugins/omemo/src/signed_pre_key_store.vala
new file mode 100644
index 00000000..44d8b3b4
--- /dev/null
+++ b/plugins/omemo/src/signed_pre_key_store.vala
@@ -0,0 +1,54 @@
+using Qlite;
+using Signal;
+
+namespace Dino.Plugins.Omemo {
+
+private class BackedSignedPreKeyStore : SimpleSignedPreKeyStore {
+    private Database db;
+    private int identity_id;
+
+    public BackedSignedPreKeyStore(Database db, int identity_id) {
+        this.db = db;
+        this.identity_id = identity_id;
+        init();
+    }
+
+    private void init() {
+        try {
+            foreach (Row row in db.signed_pre_key.select().with(db.signed_pre_key.identity_id, "=", identity_id)) {
+                store_signed_pre_key(row[db.signed_pre_key.signed_pre_key_id], Base64.decode(row[db.signed_pre_key.record_base64]));
+            }
+        } catch (Error e) {
+            print(@"OMEMO: Error while initializing signed pre key store: $(e.message)\n");
+        }
+
+        signed_pre_key_stored.connect(on_signed_pre_key_stored);
+        signed_pre_key_deleted.connect(on_signed_pre_key_deleted);
+    }
+
+    public void on_signed_pre_key_stored(SignedPreKeyStore.Key key) {
+        try {
+            db.signed_pre_key.insert().or("REPLACE")
+                    .value(db.signed_pre_key.identity_id, identity_id)
+                    .value(db.signed_pre_key.signed_pre_key_id, (int) key.key_id)
+                    .value(db.signed_pre_key.record_base64, Base64.encode(key.record))
+                    .perform();
+        } catch (Error e) {
+            print(@"OMEMO: Error while updating signed pre key store: $(e.message)\n");
+        }
+
+    }
+
+    public void on_signed_pre_key_deleted(SignedPreKeyStore.Key key) {
+        try {
+            db.signed_pre_key.delete()
+                    .with(db.signed_pre_key.identity_id, "=", identity_id)
+                    .with(db.signed_pre_key.signed_pre_key_id, "=", (int) key.key_id)
+                    .perform();
+        } catch (Error e) {
+            print(@"OMEMO: Error while updating signed pre key store: $(e.message)\n");
+        }
+    }
+}
+
+}
\ No newline at end of file
diff --git a/plugins/omemo/src/module.vala b/plugins/omemo/src/stream_module.vala
index 728251f0..546da102 100644
--- a/plugins/omemo/src/module.vala
+++ b/plugins/omemo/src/stream_module.vala
@@ -4,7 +4,7 @@ using Xmpp.Core;
 using Xmpp.Xep;
 using Signal;
 
-namespace Dino.Omemo {
+namespace Dino.Plugins.Omemo {
 
 private const string NS_URI = "eu.siacs.conversations.axolotl";
 private const string NODE_DEVICELIST = NS_URI + ".devicelist";
@@ -13,36 +13,23 @@ private const string NODE_VERIFICATION = NS_URI + ".verification";
 
 private const int NUM_KEYS_TO_PUBLISH = 100;
 
-public class Module : XmppStreamModule {
-    private const string ID = "axolotl_module";
-    public static ModuleIdentity<Module> IDENTITY = new ModuleIdentity<Module>(NS_URI, ID);
+public class StreamModule : XmppStreamModule {
+    private const string ID = "omemo_module";
+    public static ModuleIdentity<StreamModule> IDENTITY = new ModuleIdentity<StreamModule>(NS_URI, ID);
 
     private Store store;
-    internal static Context context;
     private bool device_list_loading = false;
     private bool device_list_modified = false;
     private Map<string, ArrayList<int32>> device_lists = new HashMap<string, ArrayList<int32>>();
     private Map<string, ArrayList<int32>> ignored_devices = new HashMap<string, ArrayList<int32>>();
 
-    public signal void store_created(Context context, Store store);
+    public signal void store_created(Store store);
     public signal void device_list_loaded();
     public signal void session_started(string jid, int device_id);
 
-    public Module() {
-        lock(context) {
-            if (context == null) {
-                try {
-                    context = new Context(true);
-                } catch (Error e) {
-                    print(@"Error initializing axolotl: $(e.message)\n");
-                }
-            }
-        }
-    }
-
     public EncryptStatus encrypt(Message.Stanza message, string self_bare_jid) {
         EncryptStatus status = new EncryptStatus();
-        if (context == null) return status;
+        if (Plugin.context == null) return status;
         try {
             string name = get_bare_jid(message.to);
             if (device_lists.get(name) == null || device_lists.get(self_bare_jid) == null) return status;
@@ -51,9 +38,9 @@ public class Module : XmppStreamModule {
             if (status.other_devices == 0) return status;
 
             uint8[] key = new uint8[16];
-            context.randomize(key);
+            Plugin.context.randomize(key);
             uint8[] iv = new uint8[16];
-            context.randomize(iv);
+            Plugin.context.randomize(iv);
 
             uint8[] ciphertext = aes_encrypt(Cipher.AES_GCM_NOPADDING, key, iv, message.body.data);
 
@@ -106,7 +93,7 @@ public class Module : XmppStreamModule {
             message.body = "[This message is OMEMO encrypted]";
             status.encrypted = true;
         } catch (Error e) {
-            print(@"Axolotl error while encrypting message: $(e.message)\n");
+            print(@"Signal error while encrypting message: $(e.message)\n");
         }
         return status;
     }
@@ -122,13 +109,13 @@ public class Module : XmppStreamModule {
     }
 
     public override void attach(XmppStream stream) {
-        if (context == null) return;
+        if (Plugin.context == null) return;
         Message.Module.require(stream);
         Pubsub.Module.require(stream);
         stream.get_module(Message.Module.IDENTITY).pre_received_message.connect(on_pre_received_message);
         stream.get_module(Pubsub.Module.IDENTITY).add_filtered_notification(stream, NODE_DEVICELIST, on_devicelist, this);
-        this.store = context.create_store();
-        store_created(context, store);
+        this.store = Plugin.context.create_store();
+        store_created(store);
     }
 
     private void on_pre_received_message(XmppStream stream, Message.Stanza message) {
@@ -148,11 +135,11 @@ public class Module : XmppStreamModule {
                     address.name = get_bare_jid(message.from);
                     address.device_id = header.get_attribute_int("sid");
                     if (key_node.get_attribute_bool("prekey")) {
-                        PreKeySignalMessage msg = context.deserialize_pre_key_signal_message(Base64.decode(key_node.get_string_content()));
+                        PreKeySignalMessage msg = Plugin.context.deserialize_pre_key_signal_message(Base64.decode(key_node.get_string_content()));
                         SessionCipher cipher = store.create_session_cipher(address);
                         key = cipher.decrypt_pre_key_signal_message(msg);
                     } else {
-                        SignalMessage msg = context.deserialize_signal_message(Base64.decode(key_node.get_string_content()));
+                        SignalMessage msg = Plugin.context.deserialize_signal_message(Base64.decode(key_node.get_string_content()));
                         SessionCipher cipher = store.create_session_cipher(address);
                         key = cipher.decrypt_signal_message(msg);
                     }
@@ -175,7 +162,7 @@ public class Module : XmppStreamModule {
                         flag.decrypted = true;
                     }
                 } catch (Error e) {
-                    print(@"Axolotl error while decrypting message: $(e.message)\n");
+                    print(@"Signal error while decrypting message: $(e.message)\n");
                 }
             }
         }
@@ -246,8 +233,12 @@ public class Module : XmppStreamModule {
         foreach(int32 device_id in device_lists[bare_jid]) {
             if (!is_ignored_device(bare_jid, device_id)) {
                 address.device_id = device_id;
-                if (!store.contains_session(address)) {
-                    start_session_with(stream, bare_jid, device_id);
+                try {
+                    if (!store.contains_session(address)) {
+                        start_session_with(stream, bare_jid, device_id);
+                    }
+                } catch (Error e) {
+                    // Ignore
                 }
             }
         }
@@ -259,6 +250,10 @@ public class Module : XmppStreamModule {
         stream.get_module(Pubsub.Module.IDENTITY).request(stream, bare_jid, @"$NODE_BUNDLES:$device_id", on_other_bundle_result, Tuple.create(store, device_id));
     }
 
+    public bool is_known_address(string name) {
+        return device_lists.has_key(name);
+    }
+
     public void ignore_device(string jid, int32 device_id) {
         if (device_id <= 0) return;
         lock (ignored_devices) {
@@ -313,11 +308,11 @@ public class Module : XmppStreamModule {
                     fail = true;
                 }
                 address.device_id = 0; // TODO: Hack to have address obj live longer
-                get_module(stream).session_started(jid, device_id);
+                stream.get_module(IDENTITY).session_started(jid, device_id);
             }
         }
         if (fail) {
-            get_module(stream).ignore_device(jid, device_id);
+            stream.get_module(IDENTITY).ignore_device(jid, device_id);
         }
     }
 
@@ -347,49 +342,53 @@ public class Module : XmppStreamModule {
             signed_pre_key = bundle.signed_pre_key;
         }
 
-        // Validate IdentityKey
-        if (store.identity_key_pair.public.compare(identity_key) != 0) {
-            changed = true;
-        }
-        identity_key_pair = store.identity_key_pair;
-
-        // Validate signedPreKeyRecord + ID
-        if (signed_pre_key_id == -1 || !store.contains_signed_pre_key(signed_pre_key_id) || store.load_signed_pre_key(signed_pre_key_id).key_pair.public.compare(signed_pre_key) != 0) {
-            signed_pre_key_id = Random.int_range(1, int32.MAX); // TODO: No random, use ordered number
-            signed_pre_key_record = context.generate_signed_pre_key(identity_key_pair, signed_pre_key_id);
-            store.store_signed_pre_key(signed_pre_key_record);
-            changed = true;
-        } else {
-            signed_pre_key_record = store.load_signed_pre_key(signed_pre_key_id);
-        }
+        try {
+            // Validate IdentityKey
+            if (store.identity_key_pair.public.compare(identity_key) != 0) {
+                changed = true;
+            }
+            identity_key_pair = store.identity_key_pair;
+
+            // Validate signedPreKeyRecord + ID
+            if (signed_pre_key_id == -1 || !store.contains_signed_pre_key(signed_pre_key_id) || store.load_signed_pre_key(signed_pre_key_id).key_pair.public.compare(signed_pre_key) != 0) {
+                signed_pre_key_id = Random.int_range(1, int32.MAX); // TODO: No random, use ordered number
+                signed_pre_key_record = Plugin.context.generate_signed_pre_key(identity_key_pair, signed_pre_key_id);
+                store.store_signed_pre_key(signed_pre_key_record);
+                changed = true;
+            } else {
+                signed_pre_key_record = store.load_signed_pre_key(signed_pre_key_id);
+            }
 
-        // Validate PreKeys
-        Set<PreKeyRecord> pre_key_records = new HashSet<PreKeyRecord>();
-        foreach (var entry in keys.entries) {
-            if (store.contains_pre_key(entry.key)) {
-                PreKeyRecord record = store.load_pre_key(entry.key);
-                if (record.key_pair.public.compare(entry.value) == 0) {
-                    pre_key_records.add(record);
+            // Validate PreKeys
+            Set<PreKeyRecord> pre_key_records = new HashSet<PreKeyRecord>();
+            foreach (var entry in keys.entries) {
+                if (store.contains_pre_key(entry.key)) {
+                    PreKeyRecord record = store.load_pre_key(entry.key);
+                    if (record.key_pair.public.compare(entry.value) == 0) {
+                        pre_key_records.add(record);
+                    }
                 }
             }
-        }
-        int new_keys = NUM_KEYS_TO_PUBLISH - pre_key_records.size;
-        if (new_keys > 0) {
-            int32 next_id = Random.int_range(1, int32.MAX); // TODO: No random, use ordered number
-            Set<PreKeyRecord> new_records = context.generate_pre_keys((uint)next_id, (uint)new_keys);
-            pre_key_records.add_all(new_records);
-            foreach (PreKeyRecord record in new_records) {
-                store.store_pre_key(record);
+            int new_keys = NUM_KEYS_TO_PUBLISH - pre_key_records.size;
+            if (new_keys > 0) {
+                int32 next_id = Random.int_range(1, int32.MAX); // TODO: No random, use ordered number
+                Set<PreKeyRecord> new_records = Plugin.context.generate_pre_keys((uint)next_id, (uint)new_keys);
+                pre_key_records.add_all(new_records);
+                foreach (PreKeyRecord record in new_records) {
+                    store.store_pre_key(record);
+                }
+                changed = true;
             }
-            changed = true;
-        }
 
-        if (changed) {
-            publish_bundles(stream, signed_pre_key_record, identity_key_pair, pre_key_records, (int32) store.local_registration_id);
+            if (changed) {
+                publish_bundles(stream, signed_pre_key_record, identity_key_pair, pre_key_records, (int32) store.local_registration_id);
+            }
+        } catch (Error e) {
+            print(@"Unexpected error while publishing bundle: $(e.message)\n");
         }
     }
 
-    public static void publish_bundles(XmppStream stream, SignedPreKeyRecord signed_pre_key_record, IdentityKeyPair identity_key_pair, Set<PreKeyRecord> pre_key_records, int32 device_id) {
+    public static void publish_bundles(XmppStream stream, SignedPreKeyRecord signed_pre_key_record, IdentityKeyPair identity_key_pair, Set<PreKeyRecord> pre_key_records, int32 device_id) throws Error {
         ECKeyPair tmp;
         StanzaNode bundle = new StanzaNode.build("bundle", NS_URI)
                 .add_self_xmlns()
@@ -415,10 +414,6 @@ public class Module : XmppStreamModule {
 
     }
 
-    public static Module? get_module(XmppStream stream) {
-        return (Module?) stream.get_module(IDENTITY);
-    }
-
     public override string get_ns() {
         return NS_URI;
     }
@@ -428,120 +423,4 @@ public class Module : XmppStreamModule {
     }
 }
 
-public class MessageFlag : Message.MessageFlag {
-    public const string id = "axolotl";
-
-    public bool decrypted = false;
-
-    public static MessageFlag? get_flag(Message.Stanza message) {
-        return (MessageFlag) message.get_flag(NS_URI, id);
-    }
-
-    public override string get_ns() {
-        return NS_URI;
-    }
-
-    public override string get_id() {
-        return id;
-    }
-}
-
-internal class Bundle {
-    private StanzaNode? node;
-
-    public Bundle(StanzaNode? node) {
-        this.node = node;
-    }
-
-    public int32 signed_pre_key_id { owned get {
-        if (node == null) return -1;
-        string id = node.get_deep_attribute("signedPreKeyPublic", "signedPreKeyId");
-        if (id == null) return -1;
-        return id.to_int();
-    }}
-
-    public ECPublicKey? signed_pre_key { owned get {
-        if (node == null) return null;
-        string? key = node.get_deep_string_content("signedPreKeyPublic");
-        if (key == null) return null;
-        try {
-            return Module.context.decode_public_key(Base64.decode(key));
-        } catch (Error e) {
-            return null;
-        }
-    }}
-
-    public uint8[] signed_pre_key_signature { owned get {
-        if (node == null) return null;
-        string? sig = node.get_deep_string_content("signedPreKeySignature");
-        if (sig == null) return null;
-        try {
-            return Base64.decode(sig);
-        } catch (Error e) {
-            return null;
-        }
-    }}
-
-    public ECPublicKey? identity_key { owned get {
-        if (node == null) return null;
-        string? key = node.get_deep_string_content("identityKey");
-        if (key == null) return null;
-        try {
-            return Module.context.decode_public_key(Base64.decode(key));
-        } catch (Error e) {
-            return null;
-        }
-    }}
-
-    public ArrayList<PreKey> pre_keys { owned get {
-        if (node == null || node.get_subnode("prekeys") == null) return null;
-        ArrayList<PreKey> list = new ArrayList<PreKey>();
-        node.get_deep_subnodes("prekeys", "preKeyPublic")
-                .filter((node) => node.get_attribute("preKeyId") != null)
-                .map<PreKey>(PreKey.create)
-                .foreach((key) => list.add(key));
-        return list;
-    }}
-
-    internal class PreKey {
-        private StanzaNode node;
-
-        public static PreKey create(owned StanzaNode node) {
-            return new PreKey(node);
-        }
-
-        public PreKey(StanzaNode node) {
-            this.node = node;
-        }
-
-        public int32 key_id { owned get {
-            return (node.get_attribute("preKeyId") ?? "-1").to_int();
-        }}
-
-        public ECPublicKey? key { owned get {
-            string? key = node.get_string_content();
-            if (key == null) return null;
-            try {
-                return Module.context.decode_public_key(Base64.decode(key));
-            } catch (Error e) {
-                return null;
-            }
-        }}
-    }
-}
-
-public class EncryptStatus {
-    public bool encrypted { get; internal set; }
-    public int other_devices { get; internal set; }
-    public int other_success { get; internal set; }
-    public int other_lost { get; internal set; }
-    public int other_unknown { get; internal set; }
-    public int other_failure { get; internal set; }
-    public int own_devices { get; internal set; }
-    public int own_success { get; internal set; }
-    public int own_lost { get; internal set; }
-    public int own_unknown { get; internal set; }
-    public int own_failure { get; internal set; }
-}
-
 }
\ No newline at end of file