Rework support for XEP-0384: OMEMO Encryption

Changes:

    Rework integration tests
    New structure of base integration test classes
    bump dependency on signal-protocol-java from 2.4.0 to 2.6.2
    Introduced CachingOmemoStore implementations
    Use CachingOmemoStore classes in integration tests
    Removed OmemoSession classes (replaced with more logical OmemoRatchet classes)
    Consequently also removed load/storeOmemoSession methods from OmemoStore
    Removed some clutter from KeyUtil classes
    Moved trust decision related code from OmemoStore to TrustCallback
    Require authenticated connection for many functions
    Add async initialization function in OmemoStore
    Refactor omemo test package (/java/org/jivesoftware/smack/omemo -> /java/org/jivesoftware/smackx)
    Remove OmemoStore method isFreshInstallation() as well as defaultDeviceId related stuff
    FileBasedOmemoStore: Add cleaner methods to store/load base data types (Using tryWithResource, only for future releases, once Android API gets bumped)
    Attempt to make OmemoManager thread safe
    new logic for getInstanceFor() deviceId determination
    OmemoManagers encrypt methods now don't throw exceptions when encryption for some devices fails. Instead message gets encrypted when possible and more information about failures gets returned alongside the message itself
    Added OmemoMessage class for that purpose
    Reworked entire OmemoService class
    Use safer logic for creating trust-ignoring messages (like ratchet-update messages)
    Restructure elements/provider in order to prepare for OMEMO namespace bumps
    Remove OmemoManager.regenerate() methods in favor of getInstanceFor(connection, randomDeviceId)
    Removed some unnecessary configuration options
    Prepare for support of more AES message key types
    Simplify session creation
    Where possible, avoid side effects in methods
    Add UntrustedOmemoIdentityException
    Add TrustState enum
    More improved tests

smack-omemo/build.gradle

@@ -9,5 +9,7 @@ dependencies {
 	compile project(":smack-extensions")
 	compile project(":smack-experimental")
 
+	compile "org.bouncycastle:bcprov-jdk15on:1.59"
+
 	testCompile project(path: ":smack-core", configuration: "testRuntime")
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/CachingOmemoStore.java

@@ -0,0 +1,446 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo;
+
+import java.util.Date;
+import java.util.HashMap;
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
+
+import org.jivesoftware.smackx.omemo.exceptions.CorruptedOmemoKeyException;
+import org.jivesoftware.smackx.omemo.internal.OmemoCachedDeviceList;
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+import org.jivesoftware.smackx.omemo.util.OmemoKeyUtil;
+
+import org.jxmpp.jid.BareJid;
+
+/**
+ * This class implements the Proxy Pattern in order to wrap an OmemoStore with a caching layer.
+ * This reduces access to the underlying storage layer (eg. database, filesystem) by only accessing it for
+ * missing/updated values.
+ *
+ * Alternatively this implementation can be used as an ephemeral keystore without a persisting backend.
+ *
+ * @param <T_IdKeyPair>
+ * @param <T_IdKey>
+ * @param <T_PreKey>
+ * @param <T_SigPreKey>
+ * @param <T_Sess>
+ * @param <T_Addr>
+ * @param <T_ECPub>
+ * @param <T_Bundle>
+ * @param <T_Ciph>
+ */
+public class CachingOmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph>
+        extends OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
+
+    private final HashMap<OmemoDevice, KeyCache<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess>> caches = new HashMap<>();
+    private final OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> persistent;
+    private final OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_ECPub, T_Bundle> keyUtil;
+
+    public CachingOmemoStore(OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_ECPub, T_Bundle> keyUtil) {
+        if (keyUtil == null) {
+            throw new IllegalArgumentException("KeyUtil MUST NOT be null!");
+        }
+        this.keyUtil = keyUtil;
+        persistent = null;
+    }
+
+    public CachingOmemoStore(OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> wrappedStore) {
+        if (wrappedStore == null) {
+            throw new NullPointerException("Wrapped OmemoStore MUST NOT be null!");
+        }
+        this.keyUtil = null;
+        persistent = wrappedStore;
+    }
+
+    @Override
+    public SortedSet<Integer> localDeviceIdsOf(BareJid localUser) {
+        if (persistent != null) {
+            return persistent.localDeviceIdsOf(localUser);
+        } else {
+            return new TreeSet<>(); //TODO: ?
+        }
+    }
+
+    @Override
+    public T_IdKeyPair loadOmemoIdentityKeyPair(OmemoDevice userDevice)
+            throws CorruptedOmemoKeyException {
+        T_IdKeyPair pair = getCache(userDevice).identityKeyPair;
+
+        if (pair == null && persistent != null) {
+            pair = persistent.loadOmemoIdentityKeyPair(userDevice);
+            if (pair != null) {
+                getCache(userDevice).identityKeyPair = pair;
+            }
+        }
+
+        return pair;
+    }
+
+    @Override
+    public void storeOmemoIdentityKeyPair(OmemoDevice userDevice, T_IdKeyPair identityKeyPair) {
+        getCache(userDevice).identityKeyPair = identityKeyPair;
+        if (persistent != null) {
+            persistent.storeOmemoIdentityKeyPair(userDevice, identityKeyPair);
+        }
+    }
+
+    @Override
+    public void removeOmemoIdentityKeyPair(OmemoDevice userDevice) {
+        getCache(userDevice).identityKeyPair = null;
+        if (persistent != null) {
+            persistent.removeOmemoIdentityKeyPair(userDevice);
+        }
+    }
+
+    @Override
+    public T_IdKey loadOmemoIdentityKey(OmemoDevice userDevice, OmemoDevice contactsDevice)
+            throws CorruptedOmemoKeyException {
+        T_IdKey idKey = getCache(userDevice).identityKeys.get(contactsDevice);
+
+        if (idKey == null && persistent != null) {
+            idKey = persistent.loadOmemoIdentityKey(userDevice, contactsDevice);
+            if (idKey != null) {
+                getCache(userDevice).identityKeys.put(contactsDevice, idKey);
+            }
+        }
+
+        return idKey;
+    }
+
+    @Override
+    public void storeOmemoIdentityKey(OmemoDevice userDevice, OmemoDevice device, T_IdKey t_idKey) {
+        getCache(userDevice).identityKeys.put(device, t_idKey);
+        if (persistent != null) {
+            persistent.storeOmemoIdentityKey(userDevice, device, t_idKey);
+        }
+    }
+
+    @Override
+    public void removeOmemoIdentityKey(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        getCache(userDevice).identityKeys.remove(contactsDevice);
+        if (persistent != null) {
+            persistent.removeOmemoIdentityKey(userDevice, contactsDevice);
+        }
+    }
+
+    @Override
+    public void setDateOfLastReceivedMessage(OmemoDevice userDevice, OmemoDevice from, Date date) {
+        getCache(userDevice).lastMessagesDates.put(from, date);
+        if (persistent != null) {
+            persistent.setDateOfLastReceivedMessage(userDevice, from, date);
+        }
+    }
+
+    @Override
+    public Date getDateOfLastReceivedMessage(OmemoDevice userDevice, OmemoDevice from) {
+        Date last = getCache(userDevice).lastMessagesDates.get(from);
+
+        if (last == null && persistent != null) {
+            last = persistent.getDateOfLastReceivedMessage(userDevice, from);
+            if (last != null) {
+                getCache(userDevice).lastMessagesDates.put(from, last);
+            }
+        }
+
+        return last;
+    }
+
+    @Override
+    public void setDateOfLastDeviceIdPublication(OmemoDevice userDevice, OmemoDevice contactsDevice, Date date) {
+        getCache(userDevice).lastDeviceIdPublicationDates.put(contactsDevice, date);
+        if (persistent != null) {
+            persistent.setDateOfLastReceivedMessage(userDevice, contactsDevice, date);
+        }
+    }
+
+    @Override
+    public Date getDateOfLastDeviceIdPublication(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        Date last = getCache(userDevice).lastDeviceIdPublicationDates.get(contactsDevice);
+
+        if (last == null && persistent != null) {
+            last = persistent.getDateOfLastDeviceIdPublication(userDevice, contactsDevice);
+            if (last != null) {
+                getCache(userDevice).lastDeviceIdPublicationDates.put(contactsDevice, last);
+            }
+        }
+
+        return last;
+    }
+
+    @Override
+    public void setDateOfLastSignedPreKeyRenewal(OmemoDevice userDevice, Date date) {
+        getCache(userDevice).lastRenewalDate = date;
+        if (persistent != null) {
+            persistent.setDateOfLastSignedPreKeyRenewal(userDevice, date);
+        }
+    }
+
+    @Override
+    public Date getDateOfLastSignedPreKeyRenewal(OmemoDevice userDevice) {
+        Date lastRenewal = getCache(userDevice).lastRenewalDate;
+
+        if (lastRenewal == null && persistent != null) {
+            lastRenewal = persistent.getDateOfLastSignedPreKeyRenewal(userDevice);
+            if (lastRenewal != null) {
+                getCache(userDevice).lastRenewalDate = lastRenewal;
+            }
+        }
+
+        return lastRenewal;
+    }
+
+    @Override
+    public T_PreKey loadOmemoPreKey(OmemoDevice userDevice, int preKeyId) {
+        T_PreKey preKey = getCache(userDevice).preKeys.get(preKeyId);
+
+        if (preKey == null && persistent != null) {
+            preKey = persistent.loadOmemoPreKey(userDevice, preKeyId);
+            if (preKey != null) {
+                getCache(userDevice).preKeys.put(preKeyId, preKey);
+            }
+        }
+
+        return preKey;
+    }
+
+    @Override
+    public void storeOmemoPreKey(OmemoDevice userDevice, int preKeyId, T_PreKey t_preKey) {
+        getCache(userDevice).preKeys.put(preKeyId, t_preKey);
+        if (persistent != null) {
+            persistent.storeOmemoPreKey(userDevice, preKeyId, t_preKey);
+        }
+    }
+
+    @Override
+    public void removeOmemoPreKey(OmemoDevice userDevice, int preKeyId) {
+        getCache(userDevice).preKeys.remove(preKeyId);
+        if (persistent != null) {
+            persistent.removeOmemoPreKey(userDevice, preKeyId);
+        }
+    }
+
+    @Override
+    public TreeMap<Integer, T_PreKey> loadOmemoPreKeys(OmemoDevice userDevice) {
+        TreeMap<Integer, T_PreKey> preKeys = getCache(userDevice).preKeys;
+
+        if (preKeys.isEmpty() && persistent != null) {
+            preKeys.putAll(persistent.loadOmemoPreKeys(userDevice));
+        }
+
+        return new TreeMap<>(preKeys);
+    }
+
+    @Override
+    public T_SigPreKey loadOmemoSignedPreKey(OmemoDevice userDevice, int signedPreKeyId) {
+        T_SigPreKey sigPreKey = getCache(userDevice).signedPreKeys.get(signedPreKeyId);
+
+        if (sigPreKey == null && persistent != null) {
+            sigPreKey = persistent.loadOmemoSignedPreKey(userDevice, signedPreKeyId);
+            if (sigPreKey != null) {
+                getCache(userDevice).signedPreKeys.put(signedPreKeyId, sigPreKey);
+            }
+        }
+
+        return sigPreKey;
+    }
+
+    @Override
+    public TreeMap<Integer, T_SigPreKey> loadOmemoSignedPreKeys(OmemoDevice userDevice) {
+        TreeMap<Integer, T_SigPreKey> sigPreKeys = getCache(userDevice).signedPreKeys;
+
+        if (sigPreKeys.isEmpty() && persistent != null) {
+            sigPreKeys.putAll(persistent.loadOmemoSignedPreKeys(userDevice));
+        }
+
+        return new TreeMap<>(sigPreKeys);
+    }
+
+    @Override
+    public void storeOmemoSignedPreKey(OmemoDevice userDevice,
+                                       int signedPreKeyId,
+                                       T_SigPreKey signedPreKey) {
+        getCache(userDevice).signedPreKeys.put(signedPreKeyId, signedPreKey);
+        if (persistent != null) {
+            persistent.storeOmemoSignedPreKey(userDevice, signedPreKeyId, signedPreKey);
+        }
+    }
+
+    @Override
+    public void removeOmemoSignedPreKey(OmemoDevice userDevice, int signedPreKeyId) {
+        getCache(userDevice).signedPreKeys.remove(signedPreKeyId);
+        if (persistent != null) {
+            persistent.removeOmemoSignedPreKey(userDevice, signedPreKeyId);
+        }
+    }
+
+    @Override
+    public T_Sess loadRawSession(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        HashMap<Integer, T_Sess> contactSessions = getCache(userDevice).sessions.get(contactsDevice.getJid());
+        if (contactSessions == null) {
+            contactSessions = new HashMap<>();
+            getCache(userDevice).sessions.put(contactsDevice.getJid(), contactSessions);
+        }
+
+        T_Sess session = contactSessions.get(contactsDevice.getDeviceId());
+        if (session == null && persistent != null) {
+            session = persistent.loadRawSession(userDevice, contactsDevice);
+            if (session != null) {
+                contactSessions.put(contactsDevice.getDeviceId(), session);
+            }
+        }
+
+        return session;
+    }
+
+    @Override
+    public HashMap<Integer, T_Sess> loadAllRawSessionsOf(OmemoDevice userDevice, BareJid contact) {
+        HashMap<Integer, T_Sess> sessions = getCache(userDevice).sessions.get(contact);
+        if (sessions == null) {
+            sessions = new HashMap<>();
+            getCache(userDevice).sessions.put(contact, sessions);
+        }
+
+        if (sessions.isEmpty() && persistent != null) {
+            sessions.putAll(persistent.loadAllRawSessionsOf(userDevice, contact));
+        }
+
+        return new HashMap<>(sessions);
+    }
+
+    @Override
+    public void storeRawSession(OmemoDevice userDevice, OmemoDevice contactsDevicece, T_Sess session) {
+        HashMap<Integer, T_Sess> sessions = getCache(userDevice).sessions.get(contactsDevicece.getJid());
+        if (sessions == null) {
+            sessions = new HashMap<>();
+            getCache(userDevice).sessions.put(contactsDevicece.getJid(), sessions);
+        }
+
+        sessions.put(contactsDevicece.getDeviceId(), session);
+        if (persistent != null) {
+            persistent.storeRawSession(userDevice, contactsDevicece, session);
+        }
+    }
+
+    @Override
+    public void removeRawSession(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        HashMap<Integer, T_Sess> sessions = getCache(userDevice).sessions.get(contactsDevice.getJid());
+        if (sessions != null) {
+            sessions.remove(contactsDevice.getDeviceId());
+        }
+
+        if (persistent != null) {
+            persistent.removeRawSession(userDevice, contactsDevice);
+        }
+    }
+
+    @Override
+    public void removeAllRawSessionsOf(OmemoDevice userDevice, BareJid contact) {
+        getCache(userDevice).sessions.remove(contact);
+        if (persistent != null) {
+            persistent.removeAllRawSessionsOf(userDevice, contact);
+        }
+    }
+
+    @Override
+    public boolean containsRawSession(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        HashMap<Integer, T_Sess> sessions = getCache(userDevice).sessions.get(contactsDevice.getJid());
+
+        return (sessions != null && sessions.get(contactsDevice.getDeviceId()) != null) ||
+                (persistent != null && persistent.containsRawSession(userDevice, contactsDevice));
+    }
+
+    @Override
+    public OmemoCachedDeviceList loadCachedDeviceList(OmemoDevice userDevice, BareJid contact) {
+        OmemoCachedDeviceList list = getCache(userDevice).deviceLists.get(contact);
+
+        if (list == null && persistent != null) {
+            list = persistent.loadCachedDeviceList(userDevice, contact);
+            if (list != null) {
+                getCache(userDevice).deviceLists.put(contact, list);
+            }
+        }
+
+        return list == null ? new OmemoCachedDeviceList() : new OmemoCachedDeviceList(list);
+    }
+
+    @Override
+    public void storeCachedDeviceList(OmemoDevice userDevice,
+                                      BareJid contact,
+                                      OmemoCachedDeviceList deviceList) {
+        getCache(userDevice).deviceLists.put(contact, new OmemoCachedDeviceList(deviceList));
+
+        if (persistent != null) {
+            persistent.storeCachedDeviceList(userDevice, contact, deviceList);
+        }
+    }
+
+    @Override
+    public void purgeOwnDeviceKeys(OmemoDevice userDevice) {
+        caches.remove(userDevice);
+
+        if (persistent != null) {
+            persistent.purgeOwnDeviceKeys(userDevice);
+        }
+    }
+
+    @Override
+    public OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_ECPub, T_Bundle>
+    keyUtil() {
+        if (persistent != null) {
+            return persistent.keyUtil();
+        } else {
+            return keyUtil;
+        }
+    }
+
+    /**
+     * Return the {@link KeyCache} object of an {@link OmemoManager}.
+     * @param device
+     * @return
+     */
+    private KeyCache<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess> getCache(OmemoDevice device) {
+        KeyCache<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess> cache = caches.get(device);
+        if (cache == null) {
+            cache = new KeyCache<>();
+            caches.put(device, cache);
+        }
+        return cache;
+    }
+
+    /**
+     * Cache that stores values for an {@link OmemoManager}.
+     * @param <T_IdKeyPair>
+     * @param <T_IdKey>
+     * @param <T_PreKey>
+     * @param <T_SigPreKey>
+     * @param <T_Sess>
+     */
+    private static class KeyCache<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess> {
+        private T_IdKeyPair identityKeyPair;
+        private final TreeMap<Integer, T_PreKey> preKeys = new TreeMap<>();
+        private final TreeMap<Integer, T_SigPreKey> signedPreKeys = new TreeMap<>();
+        private final HashMap<BareJid, HashMap<Integer, T_Sess>> sessions = new HashMap<>();
+        private final HashMap<OmemoDevice, T_IdKey> identityKeys = new HashMap<>();
+        private final HashMap<OmemoDevice, Date> lastMessagesDates = new HashMap<>();
+        private final HashMap<OmemoDevice, Date> lastDeviceIdPublicationDates = new HashMap<>();
+        private final HashMap<BareJid, OmemoCachedDeviceList> deviceLists = new HashMap<>();
+        private Date lastRenewalDate = null;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoConfiguration.java

@@ -16,8 +16,6 @@
  */
 package org.jivesoftware.smackx.omemo;
 
-import java.io.File;
-
 /**
  * Contains OMEMO related configuration options.
  *
@@ -33,37 +31,6 @@ public final class OmemoConfiguration {
     private static boolean IGNORE_STALE_DEVICES = true;
     private static int IGNORE_STALE_DEVICE_AFTER_HOURS = 24 * 7;         //One week
 
-    /**
-     * Delete stale devices from the device list after a period of time.
-     */
-    private static boolean DELETE_STALE_DEVICES = true;
-    private static int DELETE_STALE_DEVICE_AFTER_HOURS = 24 * 7 * 4;     //4 weeks
-
-    /**
-     * Upload a new signed prekey in intervals. This improves forward secrecy. Old keys are kept for some more time and
-     * then deleted.
-     */
-    private static boolean RENEW_OLD_SIGNED_PREKEYS = false;
-    private static int RENEW_OLD_SIGNED_PREKEYS_AFTER_HOURS = 24 * 7;    //One week
-    private static int MAX_NUMBER_OF_STORED_SIGNED_PREKEYS = 4;
-
-    /**
-     * Add a plaintext body hint about omemo encryption to the message.
-     */
-    private static boolean ADD_OMEMO_HINT_BODY = true;
-
-    /**
-     * Add Explicit Message Encryption hint (XEP-0380) to the message.
-     */
-    private static boolean ADD_EME_ENCRYPTION_HINT = true;
-
-    /**
-     * Add MAM storage hint to allow the server to store messages that do not contain a body.
-     */
-    private static boolean ADD_MAM_STORAGE_HINT = true;
-
-    private static File FILE_BASED_OMEMO_STORE_DEFAULT_PATH = null;
-
     public static void setIgnoreStaleDevices(boolean ignore) {
         IGNORE_STALE_DEVICES = ignore;
     }
@@ -83,6 +50,12 @@ public final class OmemoConfiguration {
         return IGNORE_STALE_DEVICE_AFTER_HOURS;
     }
 
+    /**
+     * Delete stale devices from the device list after a period of time.
+     */
+    private static boolean DELETE_STALE_DEVICES = true;
+    private static int DELETE_STALE_DEVICE_AFTER_HOURS = 24 * 7 * 4;     //4 weeks
+
     public static void setDeleteStaleDevices(boolean delete) {
         DELETE_STALE_DEVICES = delete;
     }
@@ -102,14 +75,39 @@ public final class OmemoConfiguration {
         return DELETE_STALE_DEVICE_AFTER_HOURS;
     }
 
+    /**
+     * Upload a new signed prekey in intervals. This improves forward secrecy. Old keys are kept for some more time and
+     * then deleted.
+     */
+    private static boolean RENEW_OLD_SIGNED_PREKEYS = false;
+    private static int RENEW_OLD_SIGNED_PREKEYS_AFTER_HOURS = 24 * 7;    //One week
+    private static int MAX_NUMBER_OF_STORED_SIGNED_PREKEYS = 4;
+
+    /**
+     * Decide, whether signed preKeys are automatically rotated or not.
+     * It is highly recommended to rotate signed preKeys to preserve forward secrecy.
+     *
+     * @param renew automatically rotate signed preKeys?
+     */
     public static void setRenewOldSignedPreKeys(boolean renew) {
         RENEW_OLD_SIGNED_PREKEYS = renew;
     }
 
+    /**
+     * Determine, whether signed preKeys are automatically rotated or not.
+     *
+     * @return auto-rotate signed preKeys?
+     */
     public static boolean getRenewOldSignedPreKeys() {
         return RENEW_OLD_SIGNED_PREKEYS;
     }
 
+    /**
+     * Set the interval in hours, after which the published signed preKey should be renewed.
+     * This value should be between one or two weeks.
+     *
+     * @param hours hours after which signed preKeys should be rotated.
+     */
     public static void setRenewOldSignedPreKeysAfterHours(int hours) {
         if (hours <= 0) {
             throw new IllegalArgumentException("Hours must be greater than 0.");
@@ -117,10 +115,23 @@ public final class OmemoConfiguration {
         RENEW_OLD_SIGNED_PREKEYS_AFTER_HOURS = hours;
     }
 
+    /**
+     * Get the interval in hours, after which the published signed preKey should be renewed.
+     * This value should be between one or two weeks.
+     *
+     * @return hours after which signed preKeys should be rotated.
+     */
     public static int getRenewOldSignedPreKeysAfterHours() {
         return RENEW_OLD_SIGNED_PREKEYS_AFTER_HOURS;
     }
 
+    /**
+     * Set the maximum number of signed preKeys that are cached until the oldest one gets deleted.
+     * This number should not be too small in order to prevent message loss, but also not too big
+     * to preserve forward secrecy.
+     *
+     * @param number number of cached signed preKeys.
+     */
     public static void setMaxNumberOfStoredSignedPreKeys(int number) {
         if (number <= 0) {
             throw new IllegalArgumentException("Number must be greater than 0.");
@@ -128,39 +139,79 @@ public final class OmemoConfiguration {
         MAX_NUMBER_OF_STORED_SIGNED_PREKEYS = number;
     }
 
+    /**
+     * Return the maximum number of signed preKeys that are cached until the oldest one gets deleted.
+     * @return max number of cached signed preKeys.
+     */
     public static int getMaxNumberOfStoredSignedPreKeys() {
         return MAX_NUMBER_OF_STORED_SIGNED_PREKEYS;
     }
 
+    /**
+     * Add a plaintext body hint about omemo encryption to the message.
+     */
+    private static boolean ADD_OMEMO_HINT_BODY = true;
+
+    /**
+     * Decide, whether an OMEMO message should carry a plaintext hint about OMEMO encryption.
+     * Eg. "I sent you an OMEMO encrypted message..."
+     *
+     * @param addHint shall we add a hint?
+     */
     public static void setAddOmemoHintBody(boolean addHint) {
         ADD_OMEMO_HINT_BODY = addHint;
     }
 
+    /**
+     * Determine, whether an OMEMO message should carry a plaintext hint about OMEMO encryption.
+     *
+     * @return true, if a hint is added to the message.
+     */
     public static boolean getAddOmemoHintBody() {
         return ADD_OMEMO_HINT_BODY;
     }
 
-    public static void setAddEmeEncryptionHint(boolean addHint) {
-        ADD_EME_ENCRYPTION_HINT = addHint;
+    private static boolean REPAIR_BROKEN_SESSIONS_WITH_PREKEY_MESSAGES = true;
+
+    /**
+     * Determine, whether incoming messages, which have broken sessions should automatically be answered by an empty
+     * preKeyMessage in order to establish a new session.
+     *
+     * @return true if session should be repaired automatically.
+     */
+    public static boolean getRepairBrokenSessionsWithPreKeyMessages() {
+        return REPAIR_BROKEN_SESSIONS_WITH_PREKEY_MESSAGES;
     }
 
-    public static boolean getAddEmeEncryptionHint() {
-        return ADD_EME_ENCRYPTION_HINT;
+    /**
+     * Decide, whether incoming messages, which have broken sessions should automatically be answered by an empty
+     * preKeyMessage in order to establish a new session.
+     *
+     * @param repair repair sessions?
+     */
+    public static void setRepairBrokenSessionsWithPrekeyMessages(boolean repair) {
+        REPAIR_BROKEN_SESSIONS_WITH_PREKEY_MESSAGES = repair;
     }
 
-    public static void setAddMAMStorageProcessingHint(boolean addStorageHint) {
-        ADD_MAM_STORAGE_HINT = addStorageHint;
+    private static boolean COMPLETE_SESSION_WITH_EMPTY_MESSAGE = true;
+
+    /**
+     * Determine, whether incoming preKeyMessages should automatically be answered by an empty message in order to
+     * complete the session.
+     *
+     * @return true if sessions should be completed.
+     */
+    public static boolean getCompleteSessionWithEmptyMessage() {
+        return COMPLETE_SESSION_WITH_EMPTY_MESSAGE;
     }
 
-    public static boolean getAddMAMStorageProcessingHint() {
-        return ADD_MAM_STORAGE_HINT;
-    }
-
-    public static void setFileBasedOmemoStoreDefaultPath(File path) {
-        FILE_BASED_OMEMO_STORE_DEFAULT_PATH = path;
-    }
-
-    public static File getFileBasedOmemoStoreDefaultPath() {
-        return FILE_BASED_OMEMO_STORE_DEFAULT_PATH;
+    /**
+     * Decide, whether incoming preKeyMessages should automatically be answered by an empty message in order to
+     * complete the session.
+     *
+     * @param complete complete the session or not
+     */
+    public static void setCompleteSessionWithEmptyMessage(boolean complete) {
+        COMPLETE_SESSION_WITH_EMPTY_MESSAGE = complete;
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoMessage.java

@@ -0,0 +1,213 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo;
+
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.BODY_OMEMO_HINT;
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO;
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO_NAMESPACE_V_AXOLOTL;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+
+import org.jivesoftware.smack.packet.Message;
+import org.jivesoftware.smackx.eme.element.ExplicitMessageEncryptionElement;
+import org.jivesoftware.smackx.hints.element.StoreHint;
+import org.jivesoftware.smackx.omemo.element.OmemoElement;
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
+
+import org.jxmpp.jid.Jid;
+
+public class OmemoMessage {
+
+    private final OmemoElement element;
+    private final byte[] messageKey, iv;
+
+    OmemoMessage(OmemoElement element, byte[] key, byte[] iv) {
+        this.element = element;
+        this.messageKey = key;
+        this.iv = iv;
+    }
+
+    /**
+     * Return the original OmemoElement (<encrypted/>).
+     *
+     * @return omemoElement
+     */
+    public OmemoElement getElement() {
+        return element;
+    }
+
+    /**
+     * Return the messageKey (or transported key in case of a KeyTransportMessage).
+     *
+     * @return key
+     */
+    public byte[] getKey() {
+        return messageKey.clone();
+    }
+
+    /**
+     * Return the initialization vector belonging to the key.
+     * @return initialization vector
+     */
+    public byte[] getIv() {
+        return iv.clone();
+    }
+
+    /**
+     * Outgoing OMEMO message.
+     */
+    public static class Sent extends OmemoMessage {
+        private final Set<OmemoDevice> intendedDevices = new HashSet<>();
+        private final HashMap<OmemoDevice, Throwable> skippedDevices = new HashMap<>();
+
+        /**
+         * Create a new outgoing OMEMO message.
+         * @param element OmemoElement
+         * @param key messageKey (or transported key)
+         * @param iv initialization vector belonging to key
+         * @param intendedDevices devices the client intended to encrypt the message for
+         * @param skippedDevices devices which were skipped during encryption process because encryption
+         *                       failed for some reason
+         */
+        Sent(OmemoElement element, byte[] key, byte[] iv, Set<OmemoDevice> intendedDevices, HashMap<OmemoDevice, Throwable> skippedDevices) {
+            super(element, key, iv);
+            this.intendedDevices.addAll(intendedDevices);
+            this.skippedDevices.putAll(skippedDevices);
+        }
+
+        /**
+         * Return a list of all devices the sender originally intended to encrypt the message for.
+         * @return list of intended recipients.
+         */
+        public Set<OmemoDevice> getIntendedDevices() {
+            return intendedDevices;
+        }
+
+        /**
+         * Return a map of all skipped recipients and the reasons for skipping.
+         * @return map of skipped recipients and reasons for that.
+         */
+        public HashMap<OmemoDevice, Throwable> getSkippedDevices() {
+            return skippedDevices;
+        }
+
+        /**
+         * Determine, if some recipients were skipped during encryption.
+         * @return true if recipients were skipped.
+         */
+        public boolean isMissingRecipients() {
+            return !getSkippedDevices().isEmpty();
+        }
+
+        /**
+         * Return the OmemoElement wrapped in a Message ready to be sent.
+         * The message is addressed to recipient, contains the OmemoElement
+         * as well as an optional clear text hint as body, a MAM storage hint
+         * and an EME hint about OMEMO encryption.
+         *
+         * @param recipient recipient for the to-field of the message.
+         * @return Message
+         */
+        public Message asMessage(Jid recipient) {
+
+            Message messageStanza = new Message();
+            messageStanza.setTo(recipient);
+            messageStanza.addExtension(getElement());
+
+            if (OmemoConfiguration.getAddOmemoHintBody()) {
+                messageStanza.setBody(BODY_OMEMO_HINT);
+            }
+
+            StoreHint.set(messageStanza);
+            messageStanza.addExtension(new ExplicitMessageEncryptionElement(OMEMO_NAMESPACE_V_AXOLOTL, OMEMO));
+
+            return messageStanza;
+        }
+    }
+
+    /**
+     * Incoming OMEMO message.
+     */
+    public static class Received extends OmemoMessage {
+        private final String message;
+        private final OmemoFingerprint sendersFingerprint;
+        private final OmemoDevice senderDevice;
+        private final boolean preKeyMessage;
+
+        /**
+         * Create a new incoming OMEMO message.
+         * @param element original OmemoElement
+         * @param key message key (or transported key)
+         * @param iv respective initialization vector
+         * @param body decrypted body
+         * @param sendersFingerprint OmemoFingerprint of the senders identityKey
+         * @param senderDevice OmemoDevice of the sender
+         * @param preKeyMessage if this was a preKeyMessage or not
+         */
+        Received(OmemoElement element, byte[] key, byte[] iv, String body, OmemoFingerprint sendersFingerprint, OmemoDevice senderDevice, boolean preKeyMessage) {
+            super(element, key, iv);
+            this.message = body;
+            this.sendersFingerprint = sendersFingerprint;
+            this.senderDevice = senderDevice;
+            this.preKeyMessage = preKeyMessage;
+        }
+
+        /**
+         * Return the decrypted body of the message.
+         * @return decrypted body
+         */
+        public String getBody() {
+            return message;
+        }
+
+        /**
+         * Return the fingerprint of the messages sender device.
+         * @return fingerprint of sender
+         */
+        public OmemoFingerprint getSendersFingerprint() {
+            return sendersFingerprint;
+        }
+
+        /**
+         * Return the OmemoDevice which sent the message.
+         * @return senderDevice
+         */
+        public OmemoDevice getSenderDevice() {
+            return senderDevice;
+        }
+
+        /**
+         * Return true, if this message was sent as a preKeyMessage.
+         * @return preKeyMessage or not
+         */
+        boolean isPreKeyMessage() {
+            return preKeyMessage;
+        }
+
+        /**
+         * Return true, if the message was a KeyTransportMessage.
+         * A KeyTransportMessage is a OmemoMessage without a payload.
+         * @return keyTransportMessage?
+         */
+        public boolean isKeyTransportMessage() {
+            return message == null;
+        }
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoRatchet.java

@@ -0,0 +1,197 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo;
+
+import java.io.UnsupportedEncodingException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.crypto.BadPaddingException;
+import javax.crypto.IllegalBlockSizeException;
+
+import org.jivesoftware.smack.util.StringUtils;
+import org.jivesoftware.smackx.omemo.element.OmemoElement;
+import org.jivesoftware.smackx.omemo.element.OmemoKeyElement;
+import org.jivesoftware.smackx.omemo.exceptions.CorruptedOmemoKeyException;
+import org.jivesoftware.smackx.omemo.exceptions.CryptoFailedException;
+import org.jivesoftware.smackx.omemo.exceptions.MultipleCryptoFailedException;
+import org.jivesoftware.smackx.omemo.exceptions.NoRawSessionException;
+import org.jivesoftware.smackx.omemo.exceptions.UntrustedOmemoIdentityException;
+import org.jivesoftware.smackx.omemo.internal.CipherAndAuthTag;
+import org.jivesoftware.smackx.omemo.internal.CiphertextTuple;
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+
+public abstract class OmemoRatchet<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
+    private static final Logger LOGGER = Logger.getLogger(OmemoRatchet.class.getName());
+
+    protected final OmemoManager omemoManager;
+    protected final OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> store;
+
+    /**
+     * Constructor.
+     *
+     * @param omemoManager omemoManager
+     * @param store omemoStore
+     */
+    public OmemoRatchet(OmemoManager omemoManager,
+                        OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> store) {
+        this.omemoManager = omemoManager;
+        this.store = store;
+    }
+
+    /**
+     * Decrypt a double-ratchet-encrypted message key.
+     *
+     * @param sender sender of the message.
+     * @param encryptedKey key encrypted with the ratchet of the sender.
+     * @return decrypted message key.
+     *
+     * @throws CorruptedOmemoKeyException
+     * @throws NoRawSessionException when no double ratchet session was found.
+     * @throws CryptoFailedException
+     * @throws UntrustedOmemoIdentityException
+     */
+    public abstract byte[] doubleRatchetDecrypt(OmemoDevice sender, byte[] encryptedKey)
+            throws CorruptedOmemoKeyException, NoRawSessionException, CryptoFailedException,
+            UntrustedOmemoIdentityException;
+
+    /**
+     * Encrypt a messageKey with the double ratchet session of the recipient.
+     *
+     * @param recipient recipient of the message.
+     * @param messageKey key we want to encrypt.
+     * @return encrypted message key.
+     */
+    public abstract CiphertextTuple doubleRatchetEncrypt(OmemoDevice recipient, byte[] messageKey);
+
+    /**
+     * Try to decrypt the transported message key using the double ratchet session.
+     *
+     * @param element omemoElement
+     * @return tuple of cipher generated from the unpacked message key and the auth-tag
+     * @throws CryptoFailedException if decryption using the double ratchet fails
+     * @throws NoRawSessionException if we have no session, but the element was NOT a PreKeyMessage
+     */
+    CipherAndAuthTag retrieveMessageKeyAndAuthTag(OmemoDevice sender, OmemoElement element) throws CryptoFailedException,
+            NoRawSessionException {
+        int keyId = omemoManager.getDeviceId();
+        byte[] unpackedKey = null;
+        List<CryptoFailedException> decryptExceptions = new ArrayList<>();
+        List<OmemoKeyElement> keys = element.getHeader().getKeys();
+
+        boolean preKey = false;
+
+        // Find key with our ID.
+        for (OmemoKeyElement k : keys) {
+            if (k.getId() == keyId) {
+                try {
+                    unpackedKey = doubleRatchetDecrypt(sender, k.getData());
+                    preKey = k.isPreKey();
+                    break;
+                } catch (CryptoFailedException e) {
+                    // There might be multiple keys with our id, but we can only decrypt one.
+                    // So we can't throw the exception, when decrypting the first duplicate which is not for us.
+                    decryptExceptions.add(e);
+                } catch (CorruptedOmemoKeyException e) {
+                    decryptExceptions.add(new CryptoFailedException(e));
+                } catch (UntrustedOmemoIdentityException e) {
+                    LOGGER.log(Level.WARNING, "Received message from " + sender + " contained unknown identityKey. Ignore message.", e);
+                }
+            }
+        }
+
+        if (unpackedKey == null) {
+            if (!decryptExceptions.isEmpty()) {
+                throw MultipleCryptoFailedException.from(decryptExceptions);
+            }
+
+            throw new CryptoFailedException("Transported key could not be decrypted, since no suitable message key " +
+                    "was provided. Provides keys: " + keys);
+        }
+
+        // Split in AES auth-tag and key
+        byte[] messageKey = new byte[16];
+        byte[] authTag = null;
+
+        if (unpackedKey.length == 32) {
+            authTag = new byte[16];
+            // copy key part into messageKey
+            System.arraycopy(unpackedKey, 0, messageKey, 0, 16);
+            // copy tag part into authTag
+            System.arraycopy(unpackedKey, 16, authTag, 0,16);
+        } else if (element.isKeyTransportElement() && unpackedKey.length == 16) {
+            messageKey = unpackedKey;
+        } else {
+            throw new CryptoFailedException("MessageKey has wrong length: "
+                    + unpackedKey.length + ". Probably legacy auth tag format.");
+        }
+
+        return new CipherAndAuthTag(messageKey, element.getHeader().getIv(), authTag, preKey);
+    }
+
+    /**
+     * Use the symmetric key in cipherAndAuthTag to decrypt the payload of the omemoMessage.
+     * The decrypted payload will be the body of the returned Message.
+     *
+     * @param element omemoElement containing a payload.
+     * @param cipherAndAuthTag cipher and authentication tag.
+     * @return decrypted plain text.
+     * @throws CryptoFailedException if decryption using AES key fails.
+     */
+    static String decryptMessageElement(OmemoElement element, CipherAndAuthTag cipherAndAuthTag)
+            throws CryptoFailedException {
+        if (!element.isMessageElement()) {
+            throw new IllegalArgumentException("decryptMessageElement cannot decrypt OmemoElement which is no MessageElement!");
+        }
+
+        if (cipherAndAuthTag.getAuthTag() == null || cipherAndAuthTag.getAuthTag().length != 16) {
+            throw new CryptoFailedException("AuthenticationTag is null or has wrong length: "
+                    + (cipherAndAuthTag.getAuthTag() == null ? "null" : cipherAndAuthTag.getAuthTag().length));
+        }
+
+        byte[] encryptedBody = payloadAndAuthTag(element, cipherAndAuthTag.getAuthTag());
+
+        try {
+            String plaintext = new String(cipherAndAuthTag.getCipher().doFinal(encryptedBody), StringUtils.UTF8);
+            return plaintext;
+
+        } catch (UnsupportedEncodingException | IllegalBlockSizeException | BadPaddingException e) {
+            throw new CryptoFailedException("decryptMessageElement could not decipher message body: "
+                    + e.getMessage());
+        }
+    }
+
+    /**
+     * Return the concatenation of the payload of the OmemoElement and the given auth tag.
+     *
+     * @param element omemoElement (message element)
+     * @param authTag authTag
+     * @return payload + authTag
+     */
+    static byte[] payloadAndAuthTag(OmemoElement element, byte[] authTag) {
+        if (!element.isMessageElement()) {
+            throw new IllegalArgumentException("OmemoElement has no payload.");
+        }
+
+        byte[] payload = new byte[element.getPayload().length + authTag.length];
+        System.arraycopy(element.getPayload(), 0, payload, 0, element.getPayload().length);
+        System.arraycopy(authTag, 0, payload, element.getPayload().length, authTag.length);
+        return payload;
+    }
+
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoBundleElement.java

@@ -16,12 +16,15 @@
  */
 package org.jivesoftware.smackx.omemo.element;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import org.jivesoftware.smack.packet.ExtensionElement;
 import org.jivesoftware.smack.util.XmlStringBuilder;
+import org.jivesoftware.smack.util.stringencoder.Base64;
 
 /**
  * Class that represents an OMEMO Bundle element.
- * TODO: Move functionality to here.
  *
  * @author Paul Schaub
  */
@@ -36,8 +39,173 @@ public abstract class OmemoBundleElement implements ExtensionElement {
     public static final String PRE_KEY_PUB = "preKeyPublic";
     public static final String PRE_KEY_ID = "preKeyId";
 
+    private final int signedPreKeyId;
+    private final String signedPreKeyB64;
+    private byte[] signedPreKey;
+    private final String signedPreKeySignatureB64;
+    private byte[] signedPreKeySignature;
+    private final String identityKeyB64;
+    private byte[] identityKey;
+    private final HashMap<Integer, String> preKeysB64;
+    private HashMap<Integer, byte[]> preKeys;
+
+    /**
+     * Constructor to create a Bundle Element from base64 Strings.
+     *
+     * @param signedPreKeyId id
+     * @param signedPreKeyB64 base64 encoded signedPreKey
+     * @param signedPreKeySigB64 base64 encoded signedPreKeySignature
+     * @param identityKeyB64 base64 encoded identityKey
+     * @param preKeysB64 HashMap of base64 encoded preKeys
+     */
+    public OmemoBundleElement(int signedPreKeyId, String signedPreKeyB64, String signedPreKeySigB64, String identityKeyB64, HashMap<Integer, String> preKeysB64) {
+        this.signedPreKeyId = signedPreKeyId;
+        this.signedPreKeyB64 = signedPreKeyB64;
+        this.signedPreKeySignatureB64 = signedPreKeySigB64;
+        this.identityKeyB64 = identityKeyB64;
+        this.preKeysB64 = preKeysB64;
+    }
+
+    /**
+     * Constructor to create a Bundle Element from decoded byte arrays.
+     *
+     * @param signedPreKeyId id
+     * @param signedPreKey signedPreKey
+     * @param signedPreKeySig signedPreKeySignature
+     * @param identityKey identityKey
+     * @param preKeys HashMap of preKeys
+     */
+    public OmemoBundleElement(int signedPreKeyId, byte[] signedPreKey, byte[] signedPreKeySig, byte[] identityKey, HashMap<Integer, byte[]> preKeys) {
+        this.signedPreKeyId = signedPreKeyId;
+
+        this.signedPreKey = signedPreKey;
+        this.signedPreKeyB64 = Base64.encodeToString(signedPreKey);
+
+        this.signedPreKeySignature = signedPreKeySig;
+        this.signedPreKeySignatureB64 = Base64.encodeToString(signedPreKeySignature);
+
+        this.identityKey = identityKey;
+        this.identityKeyB64 = Base64.encodeToString(identityKey);
+
+        this.preKeys = preKeys;
+        this.preKeysB64 = new HashMap<>();
+        for (int id : preKeys.keySet()) {
+            preKeysB64.put(id, Base64.encodeToString(preKeys.get(id)));
+        }
+    }
+
+    /**
+     * Return the signedPreKey of the OmemoBundleElement.
+     *
+     * @return signedPreKey as byte array
+     */
+    public byte[] getSignedPreKey() {
+        if (signedPreKey == null) {
+            signedPreKey = Base64.decode(signedPreKeyB64);
+        }
+        return this.signedPreKey.clone();
+    }
+
+    /**
+     * Return the id of the signedPreKey in the bundle.
+     *
+     * @return id of signedPreKey
+     */
+    public int getSignedPreKeyId() {
+        return this.signedPreKeyId;
+    }
+
+    /**
+     * Get the signature of the signedPreKey.
+     *
+     * @return signature as byte array
+     */
+    public byte[] getSignedPreKeySignature() {
+        if (signedPreKeySignature == null) {
+            signedPreKeySignature = Base64.decode(signedPreKeySignatureB64);
+        }
+        return signedPreKeySignature.clone();
+    }
+
+    /**
+     * Return the public identityKey of the bundles owner.
+     * This can be used to check the signedPreKeys signature.
+     * The fingerprint of this key is, what the user has to verify.
+     *
+     * @return public identityKey as byte array
+     */
+    public byte[] getIdentityKey() {
+        if (identityKey == null) {
+            identityKey = Base64.decode(identityKeyB64);
+        }
+        return this.identityKey.clone();
+    }
+
+    /**
+     * Return the HashMap of preKeys in the bundle.
+     * The map uses the preKeys ids as key and the preKeys as value.
+     *
+     * @return preKeys
+     */
+    public HashMap<Integer, byte[]> getPreKeys() {
+        if (preKeys == null) {
+            preKeys = new HashMap<>();
+            for (int id : preKeysB64.keySet()) {
+                preKeys.put(id, Base64.decode(preKeysB64.get(id)));
+            }
+        }
+        return this.preKeys;
+    }
+
+    /**
+     * Return a single preKey from the map.
+     *
+     * @param id id of the preKey
+     * @return the preKey
+     */
+    public byte[] getPreKey(int id) {
+        return getPreKeys().get(id);
+    }
+
     @Override
-    public abstract XmlStringBuilder toXML(String enclosingNamespace);
+    public String getElementName() {
+        return BUNDLE;
+    }
+
+    @Override
+    public XmlStringBuilder toXML(String enclosingNamespace) {
+        XmlStringBuilder sb = new XmlStringBuilder(this, enclosingNamespace).rightAngleBracket();
+
+        sb.halfOpenElement(SIGNED_PRE_KEY_PUB).attribute(SIGNED_PRE_KEY_ID, signedPreKeyId).rightAngleBracket()
+                .append(signedPreKeyB64).closeElement(SIGNED_PRE_KEY_PUB);
+
+        sb.openElement(SIGNED_PRE_KEY_SIG).append(signedPreKeySignatureB64).closeElement(SIGNED_PRE_KEY_SIG);
+
+        sb.openElement(IDENTITY_KEY).append(identityKeyB64).closeElement(IDENTITY_KEY);
+
+        sb.openElement(PRE_KEYS);
+        for (Map.Entry<Integer, String> p : this.preKeysB64.entrySet()) {
+            sb.halfOpenElement(PRE_KEY_PUB).attribute(PRE_KEY_ID, p.getKey()).rightAngleBracket()
+                    .append(p.getValue()).closeElement(PRE_KEY_PUB);
+        }
+        sb.closeElement(PRE_KEYS);
+
+        sb.closeElement(this);
+        return sb;
+    }
+
+    @Override
+    public String toString() {
+        String out = "OmemoBundleElement[\n";
+        out += SIGNED_PRE_KEY_PUB + " " + SIGNED_PRE_KEY_ID + "=" + signedPreKeyId + ": " + signedPreKeyB64 + "\n";
+        out += SIGNED_PRE_KEY_SIG + ": " + signedPreKeySignatureB64 + "\n";
+        out += IDENTITY_KEY + ": " + identityKeyB64 + "\n";
+        out += PRE_KEYS + " (" + preKeysB64.size() + ")\n";
+        for (Map.Entry<Integer, String> e : preKeysB64.entrySet()) {
+            out += PRE_KEY_PUB + " " + PRE_KEY_ID + "=" + e.getKey() + ": " + e.getValue() + "\n";
+        }
+        return out;
+    }
 
     @Override
     public boolean equals(Object other) {

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoBundleElement_VAxolotl.java

@@ -0,0 +1,43 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.element;
+
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO_NAMESPACE_V_AXOLOTL;
+
+import java.util.HashMap;
+
+/**
+ * OMEMO device bundle as described here:
+ * https://xmpp.org/extensions/xep-0384.html#usecases-announcing (Example 3).
+ *
+ * @author Paul Schaub
+ */
+public class OmemoBundleElement_VAxolotl extends OmemoBundleElement {
+
+    public OmemoBundleElement_VAxolotl(int signedPreKeyId, String signedPreKeyB64, String signedPreKeySigB64, String identityKeyB64, HashMap<Integer, String> preKeysB64) {
+        super(signedPreKeyId, signedPreKeyB64, signedPreKeySigB64, identityKeyB64, preKeysB64);
+    }
+
+    public OmemoBundleElement_VAxolotl(int signedPreKeyId, byte[] signedPreKey, byte[] signedPreKeySig, byte[] identityKey, HashMap<Integer, byte[]> preKeys) {
+        super(signedPreKeyId, signedPreKey, signedPreKeySig, identityKey, preKeys);
+    }
+
+    @Override
+    public String getNamespace() {
+        return OMEMO_NAMESPACE_V_AXOLOTL;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoBundleVAxolotlElement.java

@@ -1,207 +0,0 @@
-/**
- *
- * Copyright 2017 Paul Schaub
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jivesoftware.smackx.omemo.element;
-
-import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO_NAMESPACE_V_AXOLOTL;
-
-import java.util.HashMap;
-import java.util.Map;
-
-import org.jivesoftware.smack.util.XmlStringBuilder;
-import org.jivesoftware.smack.util.stringencoder.Base64;
-
-/**
- * OMEMO device bundle as described here:
- * https://xmpp.org/extensions/xep-0384.html#usecases-announcing (Example 3).
- *
- * @author Paul Schaub
- */
-public class OmemoBundleVAxolotlElement extends OmemoBundleElement {
-
-    private final int signedPreKeyId;
-    private final String signedPreKeyB64;
-    private byte[] signedPreKey;
-    private final String signedPreKeySignatureB64;
-    private byte[] signedPreKeySignature;
-    private final String identityKeyB64;
-    private byte[] identityKey;
-    private final HashMap<Integer, String> preKeysB64;
-    private HashMap<Integer, byte[]> preKeys;
-
-    /**
-     * Constructor to create a Bundle Element from base64 Strings.
-     *
-     * @param signedPreKeyId id
-     * @param signedPreKeyB64 base64 encoded signedPreKey
-     * @param signedPreKeySigB64 base64 encoded signedPreKeySignature
-     * @param identityKeyB64 base64 encoded identityKey
-     * @param preKeysB64 HashMap of base64 encoded preKeys
-     */
-    public OmemoBundleVAxolotlElement(int signedPreKeyId, String signedPreKeyB64, String signedPreKeySigB64, String identityKeyB64, HashMap<Integer, String> preKeysB64) {
-        this.signedPreKeyId = signedPreKeyId;
-        this.signedPreKeyB64 = signedPreKeyB64;
-        this.signedPreKeySignatureB64 = signedPreKeySigB64;
-        this.identityKeyB64 = identityKeyB64;
-        this.preKeysB64 = preKeysB64;
-    }
-
-    /**
-     * Constructor to create a Bundle Element from decoded byte arrays.
-     *
-     * @param signedPreKeyId id
-     * @param signedPreKey signedPreKey
-     * @param signedPreKeySig signedPreKeySignature
-     * @param identityKey identityKey
-     * @param preKeys HashMap of preKeys
-     */
-    public OmemoBundleVAxolotlElement(int signedPreKeyId, byte[] signedPreKey, byte[] signedPreKeySig, byte[] identityKey, HashMap<Integer, byte[]> preKeys) {
-        this.signedPreKeyId = signedPreKeyId;
-
-        this.signedPreKey = signedPreKey;
-        this.signedPreKeyB64 = Base64.encodeToString(signedPreKey);
-
-        this.signedPreKeySignature = signedPreKeySig;
-        this.signedPreKeySignatureB64 = Base64.encodeToString(signedPreKeySignature);
-
-        this.identityKey = identityKey;
-        this.identityKeyB64 = Base64.encodeToString(identityKey);
-
-        this.preKeys = preKeys;
-        this.preKeysB64 = new HashMap<>();
-        for (int id : preKeys.keySet()) {
-            preKeysB64.put(id, Base64.encodeToString(preKeys.get(id)));
-        }
-    }
-
-    /**
-     * Return the signedPreKey of the OmemoBundleElement.
-     *
-     * @return signedPreKey as byte array
-     */
-    public byte[] getSignedPreKey() {
-        if (signedPreKey == null) {
-            signedPreKey = Base64.decode(signedPreKeyB64);
-        }
-        return this.signedPreKey.clone();
-    }
-
-    /**
-     * Return the id of the signedPreKey in the bundle.
-     *
-     * @return id of signedPreKey
-     */
-    public int getSignedPreKeyId() {
-        return this.signedPreKeyId;
-    }
-
-    /**
-     * Get the signature of the signedPreKey.
-     *
-     * @return signature as byte array
-     */
-    public byte[] getSignedPreKeySignature() {
-        if (signedPreKeySignature == null) {
-            signedPreKeySignature = Base64.decode(signedPreKeySignatureB64);
-        }
-        return signedPreKeySignature.clone();
-    }
-
-    /**
-     * Return the public identityKey of the bundles owner.
-     * This can be used to check the signedPreKeys signature.
-     * The fingerprint of this key is, what the user has to verify.
-     *
-     * @return public identityKey as byte array
-     */
-    public byte[] getIdentityKey() {
-        if (identityKey == null) {
-            identityKey = Base64.decode(identityKeyB64);
-        }
-        return this.identityKey.clone();
-    }
-
-    /**
-     * Return the HashMap of preKeys in the bundle.
-     * The map uses the preKeys ids as key and the preKeys as value.
-     *
-     * @return preKeys
-     */
-    public HashMap<Integer, byte[]> getPreKeys() {
-        if (preKeys == null) {
-            preKeys = new HashMap<>();
-            for (int id : preKeysB64.keySet()) {
-                preKeys.put(id, Base64.decode(preKeysB64.get(id)));
-            }
-        }
-        return this.preKeys;
-    }
-
-    /**
-     * Return a single preKey from the map.
-     *
-     * @param id id of the preKey
-     * @return the preKey
-     */
-    public byte[] getPreKey(int id) {
-        return getPreKeys().get(id);
-    }
-
-    @Override
-    public String getElementName() {
-        return BUNDLE;
-    }
-
-    @Override
-    public XmlStringBuilder toXML(String enclosingNamespace) {
-        XmlStringBuilder sb = new XmlStringBuilder(this).rightAngleBracket();
-
-        sb.halfOpenElement(SIGNED_PRE_KEY_PUB).attribute(SIGNED_PRE_KEY_ID, signedPreKeyId).rightAngleBracket()
-                .append(signedPreKeyB64).closeElement(SIGNED_PRE_KEY_PUB);
-
-        sb.openElement(SIGNED_PRE_KEY_SIG).append(signedPreKeySignatureB64).closeElement(SIGNED_PRE_KEY_SIG);
-
-        sb.openElement(IDENTITY_KEY).append(identityKeyB64).closeElement(IDENTITY_KEY);
-
-        sb.openElement(PRE_KEYS);
-        for (Map.Entry<Integer, String> p : this.preKeysB64.entrySet()) {
-            sb.halfOpenElement(PRE_KEY_PUB).attribute(PRE_KEY_ID, p.getKey()).rightAngleBracket()
-                    .append(p.getValue()).closeElement(PRE_KEY_PUB);
-        }
-        sb.closeElement(PRE_KEYS);
-
-        sb.closeElement(this);
-        return sb;
-    }
-
-    @Override
-    public String getNamespace() {
-        return OMEMO_NAMESPACE_V_AXOLOTL;
-    }
-
-    @Override
-    public String toString() {
-        String out = "OmemoBundleElement[\n";
-        out += SIGNED_PRE_KEY_PUB + " " + SIGNED_PRE_KEY_ID + "=" + signedPreKeyId + ": " + signedPreKeyB64 + "\n";
-        out += SIGNED_PRE_KEY_SIG + ": " + signedPreKeySignatureB64 + "\n";
-        out += IDENTITY_KEY + ": " + identityKeyB64 + "\n";
-        out += PRE_KEYS + " (" + preKeysB64.size() + ")\n";
-        for (Map.Entry<Integer, String> e : preKeysB64.entrySet()) {
-            out += PRE_KEY_PUB + " " + PRE_KEY_ID + "=" + e.getKey() + ": " + e.getValue() + "\n";
-        }
-        return out;
-    }
-}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoDeviceListElement.java

@@ -23,6 +23,7 @@ import java.util.Set;
 import org.jivesoftware.smack.packet.ExtensionElement;
 import org.jivesoftware.smack.util.Objects;
 import org.jivesoftware.smack.util.XmlStringBuilder;
+import org.jivesoftware.smackx.omemo.internal.OmemoCachedDeviceList;
 
 /**
  * A OMEMO device list update containing the IDs of all active devices of a contact.
@@ -45,6 +46,10 @@ public abstract class OmemoDeviceListElement implements ExtensionElement {
         this.deviceIds = Collections.unmodifiableSet(deviceIds);
     }
 
+    public OmemoDeviceListElement(OmemoCachedDeviceList cachedList) {
+        this.deviceIds = Collections.unmodifiableSet(cachedList.getActiveDevices());
+    }
+
     public Set<Integer> getDeviceIds() {
         return deviceIds;
     }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoDeviceListElement_VAxolotl.java

@@ -20,17 +20,23 @@ import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO_NAMESPACE_
 
 import java.util.Set;
 
+import org.jivesoftware.smackx.omemo.internal.OmemoCachedDeviceList;
+
 /**
  * The OMEMO device list element with the legacy Axolotl namespace.
  *
  * @author Paul Schaub
  */
-public class OmemoDeviceListVAxolotlElement extends OmemoDeviceListElement {
+public class OmemoDeviceListElement_VAxolotl extends OmemoDeviceListElement {
 
-    public OmemoDeviceListVAxolotlElement(Set<Integer> deviceIds) {
+    public OmemoDeviceListElement_VAxolotl(Set<Integer> deviceIds) {
         super(deviceIds);
     }
 
+    public OmemoDeviceListElement_VAxolotl(OmemoCachedDeviceList cachedList) {
+        super(cachedList);
+    }
+
     @Override
     public String getNamespace() {
         return OMEMO_NAMESPACE_V_AXOLOTL;

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoElement.java

@@ -16,17 +16,13 @@
  */
 package org.jivesoftware.smackx.omemo.element;
 
-import java.util.ArrayList;
-
 import org.jivesoftware.smack.packet.ExtensionElement;
-import org.jivesoftware.smack.packet.NamedElement;
 import org.jivesoftware.smack.util.Objects;
 import org.jivesoftware.smack.util.XmlStringBuilder;
 import org.jivesoftware.smack.util.stringencoder.Base64;
 
 /**
- * Class that represents a OmemoElement.
- * TODO: Move functionality here.
+ * Class that represents an OmemoElement.
  *
  * @author Paul Schaub
  */
@@ -35,17 +31,11 @@ public abstract class OmemoElement implements ExtensionElement {
     public static final int TYPE_OMEMO_PREKEY_MESSAGE = 1;
     public static final int TYPE_OMEMO_MESSAGE = 0;
 
-    public static final String ENCRYPTED = "encrypted";
-    public static final String HEADER = "header";
-    public static final String SID = "sid";
-    public static final String KEY = "key";
-    public static final String RID = "rid";
-    public static final String PREKEY = "prekey";
-    public static final String IV = "iv";
-    public static final String PAYLOAD = "payload";
+    public static final String NAME_ENCRYPTED = "encrypted";
+    public static final String ATTR_PAYLOAD = "payload";
 
-    protected final OmemoElement.OmemoHeader header;
-    protected final byte[] payload;
+    private final OmemoHeaderElement header;
+    private final byte[] payload;
 
     /**
      * Create a new OmemoMessageElement from a header and a payload.
@@ -53,12 +43,12 @@ public abstract class OmemoElement implements ExtensionElement {
      * @param header  header of the message
      * @param payload payload
      */
-    public OmemoElement(OmemoElement.OmemoHeader header, byte[] payload) {
+    public OmemoElement(OmemoHeaderElement header, byte[] payload) {
         this.header = Objects.requireNonNull(header);
         this.payload = payload;
     }
 
-    public OmemoElement.OmemoHeader getHeader() {
+    public OmemoHeaderElement getHeader() {
         return header;
     }
 
@@ -82,113 +72,22 @@ public abstract class OmemoElement implements ExtensionElement {
         return payload != null;
     }
 
-    /**
-     * Header element of the message. The header contains information about the sender and the encrypted keys for
-     * the recipients, as well as the iv element for AES.
-     */
-    public static class OmemoHeader implements NamedElement {
-        private final int sid;
-        private final ArrayList<Key> keys;
-        private final byte[] iv;
+    @Override
+    public XmlStringBuilder toXML(String enclosingNamespace) {
+        XmlStringBuilder sb = new XmlStringBuilder(this, enclosingNamespace).rightAngleBracket();
 
-        public OmemoHeader(int sid, ArrayList<OmemoHeader.Key> keys, byte[] iv) {
-            this.sid = sid;
-            this.keys = keys;
-            this.iv = iv;
+        sb.element(header);
+
+        if (payload != null) {
+            sb.openElement(ATTR_PAYLOAD).append(Base64.encodeToString(payload)).closeElement(ATTR_PAYLOAD);
         }
 
-        /**
-         * Return the deviceId of the sender of the message.
-         *
-         * @return senders id
-         */
-        public int getSid() {
-            return sid;
-        }
+        sb.closeElement(this);
+        return sb;
+    }
 
-        public ArrayList<OmemoHeader.Key> getKeys() {
-            return new ArrayList<>(keys);
-        }
-
-        public byte[] getIv() {
-            return iv != null ? iv.clone() : null;
-        }
-
-        @Override
-        public String getElementName() {
-            return HEADER;
-        }
-
-        @Override
-        public CharSequence toXML(String enclosingNamespace) {
-            XmlStringBuilder sb = new XmlStringBuilder(this);
-            sb.attribute(SID, getSid()).rightAngleBracket();
-
-            for (OmemoHeader.Key k : getKeys()) {
-                sb.element(k);
-            }
-
-            sb.openElement(IV).append(Base64.encodeToString(getIv())).closeElement(IV);
-
-            return sb.closeElement(this);
-        }
-
-        /**
-         * Small class to collect key (byte[]), its id and whether its a prekey or not.
-         */
-        public static class Key implements NamedElement {
-            final byte[] data;
-            final int id;
-            final boolean preKey;
-
-            public Key(byte[] data, int id) {
-                this.data = data;
-                this.id = id;
-                this.preKey = false;
-            }
-
-            public Key(byte[] data, int id, boolean preKey) {
-                this.data = data;
-                this.id = id;
-                this.preKey = preKey;
-            }
-
-            public int getId() {
-                return this.id;
-            }
-
-            public byte[] getData() {
-                return this.data;
-            }
-
-            public boolean isPreKey() {
-                return this.preKey;
-            }
-
-            @Override
-            public String toString() {
-                return Integer.toString(id);
-            }
-
-            @Override
-            public String getElementName() {
-                return KEY;
-            }
-
-            @Override
-            public CharSequence toXML(String enclosingNamespace) {
-                XmlStringBuilder sb = new XmlStringBuilder(this);
-
-                if (isPreKey()) {
-                    sb.attribute(PREKEY, true);
-                }
-
-                sb.attribute(RID, getId());
-                sb.rightAngleBracket();
-                sb.append(Base64.encodeToString(getData()));
-                sb.closeElement(this);
-                return sb;
-            }
-        }
+    @Override
+    public String getElementName() {
+        return NAME_ENCRYPTED;
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoElement_VAxolotl.java

@@ -0,0 +1,44 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.element;
+
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO_NAMESPACE_V_AXOLOTL;
+
+/**
+ * An OMEMO (PreKey)WhisperMessage element.
+ *
+ * @author Paul Schaub
+ */
+public class OmemoElement_VAxolotl extends OmemoElement {
+
+    public static final String NAMESPACE = OMEMO_NAMESPACE_V_AXOLOTL;
+
+    /**
+     * Create a new OmemoMessageElement from a header and a payload.
+     *
+     * @param header  header of the message
+     * @param payload payload
+     */
+    public OmemoElement_VAxolotl(OmemoHeaderElement_VAxolotl header, byte[] payload) {
+        super(header, payload);
+    }
+
+    @Override
+    public String getNamespace() {
+        return NAMESPACE;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoHeaderElement.java

@@ -0,0 +1,83 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.element;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.jivesoftware.smack.packet.NamedElement;
+import org.jivesoftware.smack.util.XmlStringBuilder;
+import org.jivesoftware.smack.util.stringencoder.Base64;
+
+/**
+ * Header element of the message. The header contains information about the sender and the encrypted keys for
+ * the recipients, as well as the iv element for AES.
+ */
+public abstract class OmemoHeaderElement implements NamedElement {
+
+    public static final String NAME_HEADER = "header";
+    public static final String ATTR_SID = "sid";
+    public static final String ATTR_IV = "iv";
+
+    private final int sid;
+    private final List<OmemoKeyElement> keys;
+    private final byte[] iv;
+
+    public OmemoHeaderElement(int sid, List<OmemoKeyElement> keys, byte[] iv) {
+        this.sid = sid;
+        this.keys = keys;
+        this.iv = iv;
+    }
+
+    /**
+     * Return the deviceId of the sender of the message.
+     *
+     * @return senders id
+     */
+    public int getSid() {
+        return sid;
+    }
+
+    public ArrayList<OmemoKeyElement> getKeys() {
+        return new ArrayList<>(keys);
+    }
+
+    public byte[] getIv() {
+        return iv != null ? iv.clone() : null;
+    }
+
+    @Override
+    public String getElementName() {
+        return NAME_HEADER;
+    }
+
+    @Override
+    public CharSequence toXML(String enclosingNamespace) {
+        XmlStringBuilder sb = new XmlStringBuilder(this);
+        sb.attribute(ATTR_SID, getSid()).rightAngleBracket();
+
+        for (OmemoKeyElement k : getKeys()) {
+            sb.element(k);
+        }
+
+        sb.openElement(ATTR_IV).append(Base64.encodeToString(getIv())).closeElement(ATTR_IV);
+
+        return sb.closeElement(this);
+    }
+
+
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoHeaderElement_VAxolotl.java

@@ -0,0 +1,27 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.element;
+
+import java.util.List;
+
+public class OmemoHeaderElement_VAxolotl extends OmemoHeaderElement {
+
+    public OmemoHeaderElement_VAxolotl(int sid, List<OmemoKeyElement> keys, byte[] iv) {
+        super(sid, keys, iv);
+    }
+
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoKeyElement.java

@@ -0,0 +1,84 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.element;
+
+import org.jivesoftware.smack.packet.NamedElement;
+import org.jivesoftware.smack.util.XmlStringBuilder;
+import org.jivesoftware.smack.util.stringencoder.Base64;
+
+/**
+ * Small class to collect key (byte[]), its id and whether its a preKey or not.
+ */
+public class OmemoKeyElement implements NamedElement {
+
+    public static final String NAME_KEY = "key";
+    public static final String ATTR_RID = "rid";
+    public static final String ATTR_PREKEY = "prekey";
+
+    private final byte[] data;
+    private final int id;
+    private final boolean preKey;
+
+    public OmemoKeyElement(byte[] data, int id) {
+        this.data = data;
+        this.id = id;
+        this.preKey = false;
+    }
+
+    public OmemoKeyElement(byte[] data, int id, boolean preKey) {
+        this.data = data;
+        this.id = id;
+        this.preKey = preKey;
+    }
+
+    public int getId() {
+        return this.id;
+    }
+
+    public byte[] getData() {
+        return this.data;
+    }
+
+    public boolean isPreKey() {
+        return this.preKey;
+    }
+
+    @Override
+    public String toString() {
+        return Integer.toString(id);
+    }
+
+    @Override
+    public String getElementName() {
+        return NAME_KEY;
+    }
+
+    @Override
+    public CharSequence toXML(String enclosingNamespace) {
+        XmlStringBuilder sb = new XmlStringBuilder(this);
+
+        if (isPreKey()) {
+            sb.attribute(ATTR_PREKEY, true);
+        }
+
+        sb.attribute(ATTR_RID, getId());
+        sb.rightAngleBracket();
+        sb.append(Base64.encodeToString(getData()));
+        sb.closeElement(this);
+        return sb;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/element/OmemoVAxolotlElement.java

@@ -1,86 +0,0 @@
-/**
- *
- * Copyright 2017 Paul Schaub
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jivesoftware.smackx.omemo.element;
-
-import static org.jivesoftware.smackx.omemo.util.OmemoConstants.OMEMO_NAMESPACE_V_AXOLOTL;
-
-import java.io.UnsupportedEncodingException;
-
-import org.jivesoftware.smack.util.StringUtils;
-import org.jivesoftware.smack.util.XmlStringBuilder;
-import org.jivesoftware.smack.util.stringencoder.Base64;
-
-/**
- * An OMEMO (PreKey)WhisperMessage element.
- *
- * @author Paul Schaub
- */
-public class OmemoVAxolotlElement extends OmemoElement {
-
-    /**
-     * Create a new OmemoMessageElement from a header and a payload.
-     *
-     * @param header  header of the message
-     * @param payload payload
-     */
-    public OmemoVAxolotlElement(OmemoHeader header, byte[] payload) {
-        super(header, payload);
-    }
-
-    @Override
-    public String getElementName() {
-        return ENCRYPTED;
-    }
-
-    @Override
-    public XmlStringBuilder toXML(String enclosingNamespace) {
-        XmlStringBuilder sb = new XmlStringBuilder(this).rightAngleBracket();
-
-        sb.element(header);
-
-        if (payload != null) {
-            sb.openElement(PAYLOAD).append(Base64.encodeToString(payload)).closeElement(PAYLOAD);
-        }
-
-        sb.closeElement(this);
-        return sb;
-    }
-
-    @Override
-    public String getNamespace() {
-        return OMEMO_NAMESPACE_V_AXOLOTL;
-    }
-
-    @Override
-    public String toString() {
-        try {
-            StringBuilder s = new StringBuilder("Encrypted:\n")
-                    .append("   header: sid: ").append(getHeader().getSid()).append('\n');
-            for (OmemoHeader.Key k : getHeader().getKeys()) {
-                s.append("      key: prekey: ").append(k.isPreKey()).append(" rid: ")
-                        .append(k.getId()).append(' ')
-                        .append(new String(k.getData(), StringUtils.UTF8)).append('\n');
-            }
-            s.append("      iv: ").append(new String(getHeader().getIv(), StringUtils.UTF8)).append('\n');
-            s.append("  payload: ").append(new String(getPayload(), StringUtils.UTF8));
-            return s.toString();
-        } catch (UnsupportedEncodingException e) {
-            // UTF-8 must be supported on all platforms claiming to be java compatible.
-            throw new AssertionError(e);
-        }
-    }
-}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/exceptions/CryptoFailedException.java

@@ -16,6 +16,10 @@
  */
 package org.jivesoftware.smackx.omemo.exceptions;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
 /**
  * Exception gets thrown when some cryptographic function failed.
  *
@@ -25,11 +29,18 @@ public class CryptoFailedException extends Exception {
 
     private static final long serialVersionUID = 3466888654338119924L;
 
+    private final ArrayList<Exception> exceptions = new ArrayList<>();
+
     public CryptoFailedException(String message) {
         super(message);
     }
 
     public CryptoFailedException(Exception e) {
         super(e);
+        exceptions.add(e);
+    }
+
+    public List<Exception> getExceptions() {
+        return Collections.unmodifiableList(exceptions);
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/exceptions/NoIdentityKeyException.java

@@ -0,0 +1,33 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.exceptions;
+
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+
+public class NoIdentityKeyException extends Exception {
+    private static final long serialVersionUID = 1L;
+
+    private final OmemoDevice device;
+
+    public NoIdentityKeyException(OmemoDevice device) {
+        this.device = device;
+    }
+
+    public OmemoDevice getDevice() {
+        return device;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/exceptions/NoRawSessionException.java

@@ -16,6 +16,8 @@
  */
 package org.jivesoftware.smackx.omemo.exceptions;
 
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+
 /**
  * Exception that gets thrown whenever a OmemoMessage arrives that no OmemoSession was found for to decrypt it.
  *
@@ -25,7 +27,14 @@ public class NoRawSessionException extends Exception {
 
     private static final long serialVersionUID = 3466888654338119954L;
 
-    public NoRawSessionException(Exception e) {
+    private final OmemoDevice device;
+
+    public NoRawSessionException(OmemoDevice device, Exception e) {
         super(e);
+        this.device = device;
+    }
+
+    public OmemoDevice getDeviceWithoutSession() {
+        return device;
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/exceptions/StaleDeviceException.java

@@ -0,0 +1,66 @@
+/**
+ *
+ * Copyright 2018 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.exceptions;
+
+import java.util.Date;
+
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+
+public class StaleDeviceException extends Exception {
+    private static final long serialVersionUID = 1L;
+
+    private final OmemoDevice device;
+    private final Date lastMessageDate;
+    private final Date lastDeviceIdPublication;
+
+    /**
+     * This exception gets thrown if a message cannot be encrypted for a device due to the device being inactive for too long (stale).
+     *
+     * @param device OmemoDevice.
+     * @param lastMessageDate
+     * @param lastDeviceIdPublicationDate
+     */
+    public StaleDeviceException(OmemoDevice device, Date lastMessageDate, Date lastDeviceIdPublicationDate) {
+        this.device = device;
+        this.lastMessageDate = lastMessageDate;
+        this.lastDeviceIdPublication = lastDeviceIdPublicationDate;
+    }
+
+    /**
+     * Return the date on which the last OMEMO message sent from the device was received.
+     * @return last messages date
+     */
+    public Date getLastMessageDate() {
+        return lastMessageDate;
+    }
+
+    /**
+     * Return the date of the last time the deviceId was republished after being inactive/non-existent before.
+     * @return date of last deviceId (re)publication.
+     */
+    public Date getLastDeviceIdPublicationDate() {
+        return lastDeviceIdPublication;
+    }
+
+    /**
+     * Return the stale OMEMO device.
+     * @return stale device
+     */
+    public OmemoDevice getDevice() {
+        return device;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/exceptions/UndecidedOmemoIdentityException.java

@@ -16,6 +16,7 @@
  */
 package org.jivesoftware.smackx.omemo.exceptions;
 
+import java.util.Collection;
 import java.util.HashSet;
 
 import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
@@ -34,6 +35,11 @@ public class UndecidedOmemoIdentityException extends Exception {
         this.devices.add(contact);
     }
 
+    public UndecidedOmemoIdentityException(Collection<OmemoDevice> devices) {
+        super();
+        this.devices.addAll(devices);
+    }
+
     /**
      * Return the HashSet of undecided devices.
      *

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/exceptions/UntrustedOmemoIdentityException.java

@@ -0,0 +1,93 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.exceptions;
+
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
+
+/**
+ * Exception that gets thrown when we try to en-/decrypt a message for an untrusted contact.
+ * This might either be because the user actively untrusted a device, or we receive a message from a contact
+ * which contains an identityKey that differs from the one the user trusted.
+ */
+public class UntrustedOmemoIdentityException extends Exception {
+
+    private static final long serialVersionUID = 1L;
+    private final OmemoDevice device;
+    private final OmemoFingerprint trustedKey, untrustedKey;
+
+    /**
+     * Constructor for when we receive a message with an identityKey different from the one we trusted.
+     *
+     * @param device device which sent the message.
+     * @param fpTrusted fingerprint of the identityKey we previously had and trusted.
+     * @param fpUntrusted fingerprint of the new key which is untrusted.
+     */
+    public UntrustedOmemoIdentityException(OmemoDevice device, OmemoFingerprint fpTrusted, OmemoFingerprint fpUntrusted) {
+        super();
+        this.device = device;
+        this.trustedKey = fpTrusted;
+        this.untrustedKey = fpUntrusted;
+    }
+
+    /**
+     * Constructor for when encryption fails because the user untrusted a recipients device.
+     *
+     * @param device device the user wants to encrypt for, but which has been marked as untrusted.
+     * @param untrustedKey fingerprint of that device.
+     */
+    public UntrustedOmemoIdentityException(OmemoDevice device, OmemoFingerprint untrustedKey) {
+        this(device, null, untrustedKey);
+    }
+
+    /**
+     * Return the device which sent the message.
+     * @return omemoDevice.
+     */
+    public OmemoDevice getDevice() {
+        return device;
+    }
+
+    /**
+     * Return the fingerprint of the key we expected.
+     * This might return null in case this exception got thrown during encryption process.
+     * @return
+     */
+    public OmemoFingerprint getTrustedFingerprint() {
+        return trustedKey;
+    }
+
+    /**
+     * Return the fingerprint of the unexpected untrusted key.
+     * @return
+     */
+    public OmemoFingerprint getUntrustedFingerprint() {
+        return untrustedKey;
+    }
+
+    @Override
+    public String toString() {
+        if (trustedKey != null) {
+            return "Untrusted OMEMO Identity encountered:\n" +
+                    "Fingerprint of trusted key:\n" + trustedKey.blocksOf8Chars() + "\n" +
+                    "Fingerprint of untrusted key:\n" + untrustedKey.blocksOf8Chars();
+        } else {
+            return "Untrusted OMEMO Identity encountered:\n" +
+                    "Fingerprint of untrusted key:\n" + untrustedKey.blocksOf8Chars();
+        }
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/CipherAndAuthTag.java

@@ -23,7 +23,6 @@ import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.PROVIDER;
 import java.security.InvalidAlgorithmParameterException;
 import java.security.NoSuchAlgorithmException;
 import java.security.NoSuchProviderException;
-
 import javax.crypto.Cipher;
 import javax.crypto.NoSuchPaddingException;
 import javax.crypto.spec.IvParameterSpec;
@@ -38,11 +37,13 @@ import org.jivesoftware.smackx.omemo.exceptions.CryptoFailedException;
  */
 public class CipherAndAuthTag {
     private final byte[] key, iv, authTag;
+    private final boolean wasPreKey;
 
-    public CipherAndAuthTag(byte[] key, byte[] iv, byte[] authTag) throws CryptoFailedException {
+    public CipherAndAuthTag(byte[] key, byte[] iv, byte[] authTag, boolean wasPreKey) {
         this.authTag = authTag;
         this.key = key;
         this.iv = iv;
+        this.wasPreKey = wasPreKey;
     }
 
     public Cipher getCipher() throws CryptoFailedException {
@@ -82,4 +83,8 @@ public class CipherAndAuthTag {
         }
         return null;
     }
+
+    public boolean wasPreKeyEncrypted() {
+        return wasPreKey;
+    }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/ClearTextMessage.java

@@ -1,63 +0,0 @@
-/**
- *
- * Copyright 2017 Paul Schaub
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jivesoftware.smackx.omemo.internal;
-
-import org.jivesoftware.smack.packet.Message;
-
-/**
- * Class that bundles a decrypted message together with the original message and some information about the encryption.
- *
- * @author Paul Schaub
- */
-public class ClearTextMessage {
-    private final String body;
-    private final Message encryptedMessage;
-    private final OmemoMessageInformation messageInformation;
-
-    public ClearTextMessage(String message, Message original, OmemoMessageInformation messageInfo) {
-        this.body = message;
-        this.encryptedMessage = original;
-        this.messageInformation = messageInfo;
-    }
-
-    /**
-     * Return the body of the decrypted message.
-     *
-     * @return plaintext body
-     */
-    public String getBody() {
-        return body;
-    }
-
-    /**
-     * Return the original Message object.
-     *
-     * @return original message
-     */
-    public Message getOriginalMessage() {
-        return encryptedMessage;
-    }
-
-    /**
-     * Return the OmemoMessageInformation.
-     *
-     * @return omemoMessageInformation
-     */
-    public OmemoMessageInformation getMessageInformation() {
-        return messageInformation;
-    }
-}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/OmemoCachedDeviceList.java

@@ -32,17 +32,27 @@ import java.util.Set;
  *
  * @author Paul Schaub
  */
-public class CachedDeviceList implements Serializable {
+public class OmemoCachedDeviceList implements Serializable {
     private static final long serialVersionUID = 3153579238321261203L;
 
     private final Set<Integer> activeDevices;
     private final Set<Integer> inactiveDevices;
 
-    public CachedDeviceList() {
+    public OmemoCachedDeviceList() {
         this.activeDevices = new HashSet<>();
         this.inactiveDevices = new HashSet<>();
     }
 
+    public OmemoCachedDeviceList(Set<Integer> activeDevices, Set<Integer> inactiveDevices) {
+        this();
+        this.activeDevices.addAll(activeDevices);
+        this.inactiveDevices.addAll(inactiveDevices);
+    }
+
+    public OmemoCachedDeviceList(OmemoCachedDeviceList original) {
+        this(original.getActiveDevices(), original.getInactiveDevices());
+    }
+
     /**
      * Returns all active devices.
      * Active devices are all devices that were in the latest DeviceList update.
@@ -90,12 +100,18 @@ public class CachedDeviceList implements Serializable {
     }
 
     /**
-     * Add a device to the list of active devices.
+     * Add a device to the list of active devices and remove it from inactive.
      *
      * @param deviceId deviceId that will be added
      */
     public void addDevice(int deviceId) {
         activeDevices.add(deviceId);
+        inactiveDevices.remove(deviceId);
+    }
+
+    public void addInactiveDevice(int deviceId) {
+        activeDevices.remove(deviceId);
+        inactiveDevices.add(deviceId);
     }
 
     /**
@@ -108,6 +124,10 @@ public class CachedDeviceList implements Serializable {
         return activeDevices.contains(deviceId) || inactiveDevices.contains(deviceId);
     }
 
+    public boolean isActive(int deviceId) {
+        return getActiveDevices().contains(deviceId);
+    }
+
     @Override
     public String toString() {
         String out = "active: [";

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/OmemoDevice.java

@@ -16,6 +16,8 @@
  */
 package org.jivesoftware.smackx.omemo.internal;
 
+import org.jivesoftware.smackx.omemo.util.OmemoConstants;
+
 import org.jxmpp.jid.BareJid;
 
 /**
@@ -74,4 +76,12 @@ public class OmemoDevice {
         i = jid.hashCode() + deviceId;
         return i.hashCode();
     }
+
+    /**
+     * Return the name of the PubSub {@link org.jivesoftware.smackx.pubsub.LeafNode} of this device.
+     * @return node name.
+     */
+    public String getBundleNodeName() {
+        return OmemoConstants.PEP_NODE_BUNDLE_FROM_DEVICE_ID(getDeviceId());
+    }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/OmemoMessageInformation.java

@@ -1,141 +0,0 @@
-/**
- *
- * Copyright 2017 Paul Schaub
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jivesoftware.smackx.omemo.internal;
-
-/**
- * Class that contains information about a decrypted message (eg. which key was used, if it was a carbon...).
- *
- * @author Paul Schaub
- */
-public class OmemoMessageInformation {
-    private boolean isOmemoMessage;
-    private IdentityKeyWrapper senderIdentityKey;
-    private OmemoDevice senderDevice;
-    private CARBON carbon = CARBON.NONE;
-
-    /**
-     * Empty constructor.
-     */
-    // TOOD Move this class into smackx.omemo and make this constructor package protected. -Flow
-    public OmemoMessageInformation() {
-    }
-
-    /**
-     * Creates a new OmemoMessageInformation object. Its assumed, that this is about an OMEMO message.
-     *
-     * @param senderIdentityKey identityKey of the sender device
-     * @param senderDevice      device that sent the message
-     * @param carbon            Carbon type
-     */
-    public OmemoMessageInformation(IdentityKeyWrapper senderIdentityKey, OmemoDevice senderDevice, CARBON carbon) {
-        this.senderIdentityKey = senderIdentityKey;
-        this.senderDevice = senderDevice;
-        this.carbon = carbon;
-        this.isOmemoMessage = true;
-    }
-
-    /**
-     * Create a new OmemoMessageInformation.
-     *
-     * @param senderIdentityKey identityKey of the sender device
-     * @param senderDevice      device that sent the message
-     * @param carbon            Carbon type
-     * @param omemo             is this an omemo message?
-     */
-    public OmemoMessageInformation(IdentityKeyWrapper senderIdentityKey, OmemoDevice senderDevice, CARBON carbon, boolean omemo) {
-        this(senderIdentityKey, senderDevice, carbon);
-        this.isOmemoMessage = omemo;
-    }
-
-    /**
-     * Return the sender devices identityKey.
-     *
-     * @return identityKey
-     */
-    public IdentityKeyWrapper getSenderIdentityKey() {
-        return senderIdentityKey;
-    }
-
-    /**
-     * Set the sender devices identityKey.
-     *
-     * @param senderIdentityKey identityKey
-     */
-    public void setSenderIdentityKey(IdentityKeyWrapper senderIdentityKey) {
-        this.senderIdentityKey = senderIdentityKey;
-    }
-
-    /**
-     * Return the sender device.
-     *
-     * @return sender device
-     */
-    public OmemoDevice getSenderDevice() {
-        return senderDevice;
-    }
-
-    /**
-     * Return true, if this is (was) an OMEMO message.
-     * @return true if omemo
-     */
-    public boolean isOmemoMessage() {
-        return this.isOmemoMessage;
-    }
-
-    /**
-     * Set the sender device.
-     *
-     * @param senderDevice sender device
-     */
-    public void setSenderDevice(OmemoDevice senderDevice) {
-        this.senderDevice = senderDevice;
-    }
-
-    /**
-     * Return the carbon type.
-     *
-     * @return carbon type
-     */
-    public CARBON getCarbon() {
-        return carbon;
-    }
-
-    /**
-     * Set the carbon type.
-     *
-     * @param carbon carbon type
-     */
-    public void setCarbon(CARBON carbon) {
-        this.carbon = carbon;
-    }
-
-    /**
-     * Types of Carbon Messages.
-     */
-    public enum CARBON {
-        NONE,   //No carbon
-        SENT,   //Sent carbon
-        RECV    //Received Carbon
-    }
-
-    @Override
-    public String toString() {
-        return (senderDevice != null ? senderDevice.toString() : "") + " " + carbon;
-    }
-}
-
-

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/OmemoSession.java

@@ -1,266 +0,0 @@
-/**
- *
- * Copyright 2017 Paul Schaub
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jivesoftware.smackx.omemo.internal;
-
-import java.io.UnsupportedEncodingException;
-import java.util.ArrayList;
-import java.util.List;
-
-import javax.crypto.BadPaddingException;
-import javax.crypto.IllegalBlockSizeException;
-
-import org.jivesoftware.smack.packet.Message;
-import org.jivesoftware.smack.util.StringUtils;
-
-import org.jivesoftware.smackx.omemo.OmemoFingerprint;
-import org.jivesoftware.smackx.omemo.OmemoManager;
-import org.jivesoftware.smackx.omemo.OmemoStore;
-import org.jivesoftware.smackx.omemo.element.OmemoElement;
-import org.jivesoftware.smackx.omemo.element.OmemoElement.OmemoHeader.Key;
-import org.jivesoftware.smackx.omemo.exceptions.CryptoFailedException;
-import org.jivesoftware.smackx.omemo.exceptions.MultipleCryptoFailedException;
-import org.jivesoftware.smackx.omemo.exceptions.NoRawSessionException;
-
-/**
- * This class represents a OMEMO session between us and another device.
- *
- * @param <T_IdKeyPair> IdentityKeyPair class
- * @param <T_IdKey>     IdentityKey class
- * @param <T_PreKey>    PreKey class
- * @param <T_SigPreKey> SignedPreKey class
- * @param <T_Sess>      Session class
- * @param <T_Addr>      Address class
- * @param <T_ECPub>     Elliptic Curve PublicKey class
- * @param <T_Bundle>    Bundle class
- * @param <T_Ciph>      Cipher class
- * @author Paul Schaub
- */
-public abstract class OmemoSession<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
-
-    protected final T_Ciph cipher;
-    protected final OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore;
-    protected final OmemoDevice remoteDevice;
-    protected final OmemoManager omemoManager;
-    protected T_IdKey identityKey;
-    protected int preKeyId = -1;
-
-    /**
-     * Constructor used when we establish the session.
-     *
-     * @param omemoManager OmemoManager of our device
-     * @param omemoStore   OmemoStore where we want to store the session and get key information from
-     * @param remoteDevice the OmemoDevice we want to establish the session with
-     * @param identityKey  identityKey of the recipient
-     */
-    public OmemoSession(OmemoManager omemoManager,
-                        OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore,
-                        OmemoDevice remoteDevice, T_IdKey identityKey) {
-        this(omemoManager, omemoStore, remoteDevice);
-        this.identityKey = identityKey;
-    }
-
-    /**
-     * Another constructor used when they establish the session with us.
-     *
-     * @param omemoManager OmemoManager of our device
-     * @param omemoStore   OmemoStore we want to store the session and their key in
-     * @param remoteDevice identityKey of the partner
-     */
-    public OmemoSession(OmemoManager omemoManager, OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore,
-                        OmemoDevice remoteDevice) {
-        this.omemoManager = omemoManager;
-        this.omemoStore = omemoStore;
-        this.remoteDevice = remoteDevice;
-        this.cipher = createCipher(remoteDevice);
-    }
-
-    /**
-     * Try to decrypt the transported message key using the double ratchet session.
-     *
-     * @param element omemoElement
-     * @param keyId our keyId
-     * @return tuple of cipher generated from the unpacked message key and the authtag
-     * @throws CryptoFailedException if decryption using the double ratchet fails
-     * @throws NoRawSessionException if we have no session, but the element was NOT a PreKeyMessage
-     */
-    public CipherAndAuthTag decryptTransportedKey(OmemoElement element, int keyId) throws CryptoFailedException,
-            NoRawSessionException {
-        byte[] unpackedKey = null;
-        List<CryptoFailedException> decryptExceptions = new ArrayList<>();
-        List<Key> keys = element.getHeader().getKeys();
-        // Find key with our ID.
-        for (OmemoElement.OmemoHeader.Key k : keys) {
-            if (k.getId() == keyId) {
-                try {
-                    unpackedKey = decryptMessageKey(k.getData());
-                    break;
-                } catch (CryptoFailedException e) {
-                    // There might be multiple keys with our id, but we can only decrypt one.
-                    // So we can't throw the exception, when decrypting the first duplicate which is not for us.
-                    decryptExceptions.add(e);
-                }
-            }
-        }
-
-        if (unpackedKey == null) {
-            if (!decryptExceptions.isEmpty()) {
-                throw MultipleCryptoFailedException.from(decryptExceptions);
-            }
-
-            throw new CryptoFailedException("Transported key could not be decrypted, since no provided message key. Provides keys: " + keys);
-        }
-
-        byte[] messageKey = new byte[16];
-        byte[] authTag = null;
-
-        if (unpackedKey.length == 32) {
-            authTag = new byte[16];
-            // copy key part into messageKey
-            System.arraycopy(unpackedKey, 0, messageKey, 0, 16);
-            // copy tag part into authTag
-            System.arraycopy(unpackedKey, 16, authTag, 0,16);
-        } else if (element.isKeyTransportElement() && unpackedKey.length == 16) {
-            messageKey = unpackedKey;
-        } else {
-            throw new CryptoFailedException("MessageKey has wrong length: "
-                    + unpackedKey.length + ". Probably legacy auth tag format.");
-        }
-
-        return new CipherAndAuthTag(messageKey, element.getHeader().getIv(), authTag);
-    }
-
-    /**
-     * Use the symmetric key in cipherAndAuthTag to decrypt the payload of the omemoMessage.
-     * The decrypted payload will be the body of the returned Message.
-     *
-     * @param element omemoElement containing a payload.
-     * @param cipherAndAuthTag cipher and authentication tag.
-     * @return Message containing the decrypted payload in its body.
-     * @throws CryptoFailedException
-     */
-    public static Message decryptMessageElement(OmemoElement element, CipherAndAuthTag cipherAndAuthTag) throws CryptoFailedException {
-        if (!element.isMessageElement()) {
-            throw new IllegalArgumentException("decryptMessageElement cannot decrypt OmemoElement which is no MessageElement!");
-        }
-
-        if (cipherAndAuthTag.getAuthTag() == null || cipherAndAuthTag.getAuthTag().length != 16) {
-            throw new CryptoFailedException("AuthenticationTag is null or has wrong length: "
-                    + (cipherAndAuthTag.getAuthTag() == null ? "null" : cipherAndAuthTag.getAuthTag().length));
-        }
-        byte[] encryptedBody = new byte[element.getPayload().length + 16];
-        byte[] payload = element.getPayload();
-        System.arraycopy(payload, 0, encryptedBody, 0, payload.length);
-        System.arraycopy(cipherAndAuthTag.getAuthTag(), 0, encryptedBody, payload.length, 16);
-
-        try {
-            String plaintext = new String(cipherAndAuthTag.getCipher().doFinal(encryptedBody), StringUtils.UTF8);
-            Message decrypted = new Message();
-            decrypted.setBody(plaintext);
-            return decrypted;
-
-        } catch (UnsupportedEncodingException | IllegalBlockSizeException | BadPaddingException e) {
-            throw new CryptoFailedException("decryptMessageElement could not decipher message body: "
-                    + e.getMessage());
-        }
-    }
-
-    /**
-     * Try to decrypt the message.
-     * First decrypt the message key using our session with the sender.
-     * Second use the decrypted key to decrypt the message.
-     * The decrypted content of the 'encrypted'-element becomes the body of the clear text message.
-     *
-     * @param element OmemoElement
-     * @param keyId   the key we want to decrypt (usually our own device id)
-     * @return message as plaintext
-     * @throws CryptoFailedException
-     * @throws NoRawSessionException
-     */
-    // TODO find solution for what we actually want to decrypt (String, Message, List<ExtensionElements>...)
-    public Message decryptMessageElement(OmemoElement element, int keyId) throws CryptoFailedException, NoRawSessionException {
-        if (!element.isMessageElement()) {
-            throw new IllegalArgumentException("OmemoElement is not a messageElement!");
-        }
-
-        CipherAndAuthTag cipherAndAuthTag = decryptTransportedKey(element, keyId);
-        return decryptMessageElement(element, cipherAndAuthTag);
-    }
-
-    /**
-     * Create a new SessionCipher used to encrypt/decrypt keys. The cipher typically implements the ratchet and KDF-chains.
-     *
-     * @param contact    OmemoDevice
-     * @return SessionCipher
-     */
-    public abstract T_Ciph createCipher(OmemoDevice contact);
-
-    /**
-     * Get the id of the preKey used to establish the session.
-     *
-     * @return id
-     */
-    public int getPreKeyId() {
-        return this.preKeyId;
-    }
-
-    /**
-     * Encrypt a message key for the recipient. This key can be deciphered by the recipient with its corresponding
-     * session cipher. The key is then used to decipher the message.
-     *
-     * @param messageKey serialized key to encrypt
-     * @return A CiphertextTuple containing the ciphertext and the messageType
-     * @throws CryptoFailedException
-     */
-    public abstract CiphertextTuple encryptMessageKey(byte[] messageKey) throws CryptoFailedException;
-
-    /**
-     * Decrypt a messageKey using our sessionCipher. We can use that key to decipher the actual message.
-     * Same as encryptMessageKey, just the other way round.
-     *
-     * @param encryptedKey encrypted key
-     * @return serialized decrypted key or null
-     * @throws CryptoFailedException when decryption fails.
-     * @throws NoRawSessionException when no session was found in the double ratchet library
-     */
-    public abstract byte[] decryptMessageKey(byte[] encryptedKey) throws CryptoFailedException, NoRawSessionException;
-
-    /**
-     * Return the identityKey of the session.
-     *
-     * @return identityKey
-     */
-    public T_IdKey getIdentityKey() {
-        return identityKey;
-    }
-
-    /**
-     * Set the identityKey of the remote device.
-     * @param identityKey identityKey
-     */
-    public void setIdentityKey(T_IdKey identityKey) {
-        this.identityKey = identityKey;
-    }
-
-    /**
-     * Return the fingerprint of the contacts identityKey.
-     *
-     * @return fingerprint or null
-     */
-    public OmemoFingerprint getFingerprint() {
-        return (this.identityKey != null ? omemoStore.keyUtil().getFingerprint(this.identityKey) : null);
-    }
-}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/listener/OmemoCarbonCopyStanzaReceivedListener.java

@@ -0,0 +1,29 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.internal.listener;
+
+import org.jivesoftware.smack.packet.Message;
+import org.jivesoftware.smackx.carbons.packet.CarbonExtension;
+import org.jivesoftware.smackx.omemo.OmemoManager;
+
+/**
+ * Internal listener for OMEMO encrypted carbon copies.
+ */
+public interface OmemoCarbonCopyStanzaReceivedListener {
+
+    void onOmemoCarbonCopyReceived(CarbonExtension.Direction direction, Message carbonCopy, Message wrappingMessage, OmemoManager.LoggedInOmemoManager omemoManager);
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/listener/OmemoMessageStanzaReceivedListener.java

@@ -0,0 +1,25 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.internal.listener;
+
+import org.jivesoftware.smack.packet.Stanza;
+import org.jivesoftware.smackx.omemo.OmemoManager;
+
+public interface OmemoMessageStanzaReceivedListener {
+
+    void onOmemoMessageStanzaReceived(Stanza stanza, OmemoManager.LoggedInOmemoManager omemoManager);
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/internal/listener/package-info.java

@@ -14,23 +14,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smackx.omemo.internal;
-
 /**
- * Wrapper for IdentityKey objects.
+ * StanzaListeners used for internal purposes.
  *
  * @author Paul Schaub
+ * @see <a href="https://conversations.im/xeps/multi-end.html">XEP-0384: OMEMO</a>
  */
-public class IdentityKeyWrapper {
-    private final Object identityKey;
-
-    public IdentityKeyWrapper(Object wrapped) {
-        identityKey = wrapped;
-    }
-
-    public Object getIdentityKey() {
-        return identityKey;
-    }
-
-
-}
+package org.jivesoftware.smackx.omemo.internal.listener;

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/listener/OmemoMessageListener.java

@@ -17,9 +17,9 @@
 package org.jivesoftware.smackx.omemo.listener;
 
 import org.jivesoftware.smack.packet.Message;
-
-import org.jivesoftware.smackx.omemo.internal.CipherAndAuthTag;
-import org.jivesoftware.smackx.omemo.internal.OmemoMessageInformation;
+import org.jivesoftware.smack.packet.Stanza;
+import org.jivesoftware.smackx.carbons.packet.CarbonExtension;
+import org.jivesoftware.smackx.omemo.OmemoMessage;
 
 /**
  * Listener interface that allows implementations to receive decrypted OMEMO messages.
@@ -30,20 +30,13 @@ public interface OmemoMessageListener {
     /**
      * Gets called, whenever an OmemoMessage has been received and was successfully decrypted.
      *
-     * @param decryptedBody    Decrypted body
-     * @param encryptedMessage Encrypted Message
-     * @param wrappingMessage  Wrapping carbon message, in case the message was a carbon copy, else null.
-     * @param omemoInformation Information about the messages encryption etc.
+     * @param stanza Received (encrypted) stanza.
+     * @param decryptedMessage decrypted OmemoMessage.
      */
-    void onOmemoMessageReceived(String decryptedBody, Message encryptedMessage, Message wrappingMessage, OmemoMessageInformation omemoInformation);
+    void onOmemoMessageReceived(Stanza stanza, OmemoMessage.Received decryptedMessage);
 
-    /**
-     * Gets called, whenever an OmemoElement without a body (an OmemoKeyTransportElement) is received.
-     *
-     * @param cipherAndAuthTag  transported Cipher along with an optional AuthTag
-     * @param message           Message that contained the KeyTransport
-     * @param wrappingMessage   Wrapping message (eg. carbon), or null
-     * @param omemoInformation  Information about the messages encryption etc.
-     */
-    void onOmemoKeyTransportReceived(CipherAndAuthTag cipherAndAuthTag, Message message, Message wrappingMessage, OmemoMessageInformation omemoInformation);
+    void onOmemoCarbonCopyReceived(CarbonExtension.Direction direction,
+                                   Message carbonCopy,
+                                   Message wrappingMessage,
+                                   OmemoMessage.Received decryptedCarbonCopy);
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/listener/OmemoMucMessageListener.java

@@ -16,13 +16,9 @@
  */
 package org.jivesoftware.smackx.omemo.listener;
 
-import org.jivesoftware.smack.packet.Message;
-
+import org.jivesoftware.smack.packet.Stanza;
 import org.jivesoftware.smackx.muc.MultiUserChat;
-import org.jivesoftware.smackx.omemo.internal.CipherAndAuthTag;
-import org.jivesoftware.smackx.omemo.internal.OmemoMessageInformation;
-
-import org.jxmpp.jid.BareJid;
+import org.jivesoftware.smackx.omemo.OmemoMessage;
 
 /**
  * Listener interface that allows implementations to receive decrypted OMEMO MUC messages.
@@ -33,24 +29,8 @@ public interface OmemoMucMessageListener {
     /**
      * Gets called whenever an OMEMO message has been received in a MultiUserChat and successfully decrypted.
      * @param muc MultiUserChat the message was sent in
-     * @param from the bareJid of the sender
-     * @param decryptedBody the decrypted Body of the message
-     * @param message the original message with encrypted element
-     * @param wrappingMessage in case of a carbon copy, this is the wrapping message
-     * @param omemoInformation information about the encryption of the message
+     * @param stanza Original Stanza
+     * @param decryptedOmemoMessage decrypted Omemo message
      */
-    void onOmemoMucMessageReceived(MultiUserChat muc, BareJid from, String decryptedBody, Message message,
-                                   Message wrappingMessage, OmemoMessageInformation omemoInformation);
-
-    /**
-     * Gets called, whenever an OmemoElement without a body (an OmemoKeyTransportElement) is received.
-     *
-     * @param muc               MultiUserChat the message was sent in
-     * @param from              bareJid of the sender
-     * @param cipherAndAuthTag  transportedKey along with an optional authTag
-     * @param message           Message that contained the KeyTransport
-     * @param wrappingMessage   Wrapping message (eg. carbon), or null
-     * @param omemoInformation  Information about the messages encryption etc.
-     */
-    void onOmemoKeyTransportReceived(MultiUserChat muc, BareJid from, CipherAndAuthTag cipherAndAuthTag, Message message, Message wrappingMessage, OmemoMessageInformation omemoInformation);
+    void onOmemoMucMessageReceived(MultiUserChat muc, Stanza stanza, OmemoMessage.Received decryptedOmemoMessage);
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/provider/OmemoBundleVAxolotlProvider.java

@@ -31,7 +31,7 @@ import java.util.HashMap;
 
 import org.jivesoftware.smack.provider.ExtensionElementProvider;
 
-import org.jivesoftware.smackx.omemo.element.OmemoBundleVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoBundleElement_VAxolotl;
 
 import org.xmlpull.v1.XmlPullParser;
 
@@ -40,9 +40,9 @@ import org.xmlpull.v1.XmlPullParser;
  *
  * @author Paul Schaub
  */
-public class OmemoBundleVAxolotlProvider extends ExtensionElementProvider<OmemoBundleVAxolotlElement> {
+public class OmemoBundleVAxolotlProvider extends ExtensionElementProvider<OmemoBundleElement_VAxolotl> {
     @Override
-    public OmemoBundleVAxolotlElement parse(XmlPullParser parser, int initialDepth) throws Exception {
+    public OmemoBundleElement_VAxolotl parse(XmlPullParser parser, int initialDepth) throws Exception {
         boolean stop = false;
         boolean inPreKeys = false;
 
@@ -96,6 +96,6 @@ public class OmemoBundleVAxolotlProvider extends ExtensionElementProvider<OmemoB
                     break;
             }
         }
-        return new OmemoBundleVAxolotlElement(signedPreKeyId, signedPreKey, signedPreKeySignature, identityKey, preKeys);
+        return new OmemoBundleElement_VAxolotl(signedPreKeyId, signedPreKey, signedPreKeySignature, identityKey, preKeys);
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/provider/OmemoDeviceListVAxolotlProvider.java

@@ -27,7 +27,7 @@ import java.util.Set;
 
 import org.jivesoftware.smack.provider.ExtensionElementProvider;
 
-import org.jivesoftware.smackx.omemo.element.OmemoDeviceListVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoDeviceListElement_VAxolotl;
 
 import org.xmlpull.v1.XmlPullParser;
 
@@ -36,10 +36,10 @@ import org.xmlpull.v1.XmlPullParser;
  *
  * @author Paul Schaub
  */
-public class OmemoDeviceListVAxolotlProvider extends ExtensionElementProvider<OmemoDeviceListVAxolotlElement> {
+public class OmemoDeviceListVAxolotlProvider extends ExtensionElementProvider<OmemoDeviceListElement_VAxolotl> {
 
     @Override
-    public OmemoDeviceListVAxolotlElement parse(XmlPullParser parser, int initialDepth) throws Exception {
+    public OmemoDeviceListElement_VAxolotl parse(XmlPullParser parser, int initialDepth) throws Exception {
         Set<Integer> deviceListIds = new HashSet<>();
         boolean stop = false;
         while (!stop) {
@@ -63,6 +63,6 @@ public class OmemoDeviceListVAxolotlProvider extends ExtensionElementProvider<Om
                     break;
             }
         }
-        return new OmemoDeviceListVAxolotlElement(deviceListIds);
+        return new OmemoDeviceListElement_VAxolotl(deviceListIds);
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/provider/OmemoVAxolotlProvider.java

@@ -16,14 +16,8 @@
  */
 package org.jivesoftware.smackx.omemo.provider;
 
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.ENCRYPTED;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.HEADER;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.IV;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.KEY;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.PAYLOAD;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.PREKEY;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.RID;
-import static org.jivesoftware.smackx.omemo.element.OmemoElement.SID;
+import static org.jivesoftware.smackx.omemo.element.OmemoElement.ATTR_PAYLOAD;
+import static org.jivesoftware.smackx.omemo.element.OmemoElement.NAME_ENCRYPTED;
 import static org.xmlpull.v1.XmlPullParser.END_TAG;
 import static org.xmlpull.v1.XmlPullParser.START_TAG;
 
@@ -31,8 +25,10 @@ import java.util.ArrayList;
 
 import org.jivesoftware.smack.provider.ExtensionElementProvider;
 import org.jivesoftware.smack.util.stringencoder.Base64;
-
-import org.jivesoftware.smackx.omemo.element.OmemoVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoElement_VAxolotl;
+import org.jivesoftware.smackx.omemo.element.OmemoHeaderElement;
+import org.jivesoftware.smackx.omemo.element.OmemoHeaderElement_VAxolotl;
+import org.jivesoftware.smackx.omemo.element.OmemoKeyElement;
 
 import org.xmlpull.v1.XmlPullParser;
 
@@ -41,13 +37,13 @@ import org.xmlpull.v1.XmlPullParser;
  *
  * @author Paul Schaub
  */
-public class OmemoVAxolotlProvider extends ExtensionElementProvider<OmemoVAxolotlElement> {
+public class OmemoVAxolotlProvider extends ExtensionElementProvider<OmemoElement_VAxolotl> {
 
     @Override
-    public OmemoVAxolotlElement parse(XmlPullParser parser, int initialDepth) throws Exception {
+    public OmemoElement_VAxolotl parse(XmlPullParser parser, int initialDepth) throws Exception {
         boolean inEncrypted = true;
         int sid = -1;
-        ArrayList<OmemoVAxolotlElement.OmemoHeader.Key> keys = new ArrayList<>();
+        ArrayList<OmemoKeyElement> keys = new ArrayList<>();
         byte[] iv = null;
         byte[] payload = null;
 
@@ -57,41 +53,41 @@ public class OmemoVAxolotlProvider extends ExtensionElementProvider<OmemoVAxolot
             switch (tag) {
                 case START_TAG:
                     switch (name) {
-                        case HEADER:
+                        case OmemoHeaderElement.NAME_HEADER:
                             for (int i = 0; i < parser.getAttributeCount(); i++) {
-                                if (parser.getAttributeName(i).equals(SID)) {
+                                if (parser.getAttributeName(i).equals(OmemoHeaderElement.ATTR_SID)) {
                                     sid = Integer.parseInt(parser.getAttributeValue(i));
                                 }
                             }
                             break;
-                        case KEY:
+                        case OmemoKeyElement.NAME_KEY:
                             boolean prekey = false;
                             int rid = -1;
                             for (int i = 0; i < parser.getAttributeCount(); i++) {
-                                if (parser.getAttributeName(i).equals(PREKEY)) {
+                                if (parser.getAttributeName(i).equals(OmemoKeyElement.ATTR_PREKEY)) {
                                     prekey = Boolean.parseBoolean(parser.getAttributeValue(i));
-                                } else if (parser.getAttributeName(i).equals(RID)) {
+                                } else if (parser.getAttributeName(i).equals(OmemoKeyElement.ATTR_RID)) {
                                     rid = Integer.parseInt(parser.getAttributeValue(i));
                                 }
                             }
-                            keys.add(new OmemoVAxolotlElement.OmemoHeader.Key(Base64.decode(parser.nextText()), rid, prekey));
+                            keys.add(new OmemoKeyElement(Base64.decode(parser.nextText()), rid, prekey));
                             break;
-                        case IV:
+                        case OmemoHeaderElement.ATTR_IV:
                             iv = Base64.decode(parser.nextText());
                             break;
-                        case PAYLOAD:
+                        case ATTR_PAYLOAD:
                             payload = Base64.decode(parser.nextText());
                             break;
                     }
                     break;
                 case END_TAG:
-                    if (name.equals(ENCRYPTED)) {
+                    if (name.equals(NAME_ENCRYPTED)) {
                         inEncrypted = false;
                     }
                     break;
             }
         }
-        OmemoVAxolotlElement.OmemoHeader header = new OmemoVAxolotlElement.OmemoHeader(sid, keys, iv);
-        return new OmemoVAxolotlElement(header, payload);
+        OmemoHeaderElement_VAxolotl header = new OmemoHeaderElement_VAxolotl(sid, keys, iv);
+        return new OmemoElement_VAxolotl(header, payload);
     }
 }

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/trust/OmemoFingerprint.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smackx.omemo;
+package org.jivesoftware.smackx.omemo.trust;
 
 public class OmemoFingerprint implements CharSequence {
 
@@ -39,10 +39,6 @@ public class OmemoFingerprint implements CharSequence {
         return fingerprintString.subSequence(start, end);
     }
 
-    public CharSequence subSequence(int start) {
-        return fingerprintString.subSequence(start, fingerprintString.length() - 1);
-    }
-
     @Override
     public String toString() {
         return fingerprintString;
@@ -57,6 +53,20 @@ public class OmemoFingerprint implements CharSequence {
         return this.toString().trim().equals(otherFingerprint.toString().trim());
     }
 
+    /**
+     * Split the fingerprint in blocks of 8 characters with spaces between.
+     *
+     * @return Block representation of the fingerprint.
+     */
+    public String blocksOf8Chars() {
+        StringBuilder pretty = new StringBuilder();
+        for (int i = 0; i < 8; i++) {
+            if (i != 0) pretty.append(' ');
+            pretty.append(this.fingerprintString.substring(8 * i, 8 * (i + 1)));
+        }
+        return pretty.toString();
+    }
+
     @Override
     public int hashCode() {
         return toString().hashCode();

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/trust/OmemoTrustCallback.java

@@ -0,0 +1,27 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.trust;
+
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+
+public interface OmemoTrustCallback {
+
+    TrustState getTrust(OmemoDevice device, OmemoFingerprint fingerprint);
+
+    void setTrust(OmemoDevice device, OmemoFingerprint fingerprint, TrustState state);
+
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/trust/TrustState.java

@@ -0,0 +1,23 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.trust;
+
+public enum TrustState {
+    undecided, // User has yet to decide, whether the identity is trusted or not.
+    untrusted, // User decided NOT to trust this device.
+    trusted    // User decided to trust this device.
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/trust/package-info.java

@@ -0,0 +1,23 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/**
+ * Callbacks used to pass trust decisions up to the client.
+ *
+ * @author Paul Schaub
+ * @see <a href="https://conversations.im/xeps/multi-end.html">XEP-0384: OMEMO</a>
+ */
+package org.jivesoftware.smackx.omemo.trust;

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/util/MessageOrOmemoMessage.java

@@ -0,0 +1,49 @@
+/**
+ *
+ * Copyright 2018 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.util;
+
+import org.jivesoftware.smack.packet.Message;
+import org.jivesoftware.smack.util.Objects;
+import org.jivesoftware.smackx.omemo.OmemoMessage;
+
+public class MessageOrOmemoMessage {
+
+    private final Message message;
+    private final OmemoMessage.Received omemoMessage;
+
+    public MessageOrOmemoMessage(Message message) {
+        this.message = Objects.requireNonNull(message);
+        this.omemoMessage = null;
+    }
+
+    public MessageOrOmemoMessage(OmemoMessage.Received omemoMessage) {
+        this.omemoMessage = Objects.requireNonNull(omemoMessage);
+        this.message = null;
+    }
+
+    public boolean isOmemoMessage() {
+        return omemoMessage != null;
+    }
+
+    public Message getMessage() {
+        return message;
+    }
+
+    public OmemoMessage.Received getOmemoMessage() {
+        return omemoMessage;
+    }
+}

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/util/OmemoConstants.java

@@ -37,7 +37,7 @@ public final class OmemoConstants {
     /**
      * How many preKeys do we want to publish?
      */
-    public static final int TARGET_PRE_KEY_COUNT = 100;
+    public static final int PRE_KEY_COUNT_PER_BUNDLE = 100;
 
     /**
      * Return the node name of the PEP node containing the device bundle of the device with device id deviceId.
@@ -49,7 +49,7 @@ public final class OmemoConstants {
         return PEP_NODE_BUNDLES + ":" + deviceId;
     }
 
-    public static final String BODY_OMEMO_HINT = "I sent you an OMEMO encrypted message but your client doesn’t seem to support that. Find more information on https://conversations.im/omemo";
+    public static final String BODY_OMEMO_HINT = "I sent you an OMEMO encrypted message but your client doesn't seem to support that. Find more information on https://conversations.im/omemo";
 
     /**
      * Information about the keys used for message encryption.

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/util/OmemoKeyUtil.java

@@ -19,18 +19,14 @@ package org.jivesoftware.smackx.omemo.util;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.TreeMap;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-import org.jivesoftware.smackx.omemo.OmemoFingerprint;
-import org.jivesoftware.smackx.omemo.OmemoManager;
-import org.jivesoftware.smackx.omemo.OmemoStore;
-import org.jivesoftware.smackx.omemo.element.OmemoBundleVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoBundleElement;
 import org.jivesoftware.smackx.omemo.exceptions.CorruptedOmemoKeyException;
 import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
-import org.jivesoftware.smackx.omemo.internal.OmemoSession;
-
-import org.jxmpp.stringprep.XmppStringprepException;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
 
 /**
  * Class that is used to convert bytes to keys and vice versa.
@@ -40,13 +36,11 @@ import org.jxmpp.stringprep.XmppStringprepException;
  * @param <T_PreKey>    PreKey class
  * @param <T_SigPreKey> SignedPreKey class
  * @param <T_Sess>      Session class
- * @param <T_Addr>      Address class
  * @param <T_ECPub>     Elliptic Curve PublicKey class
  * @param <T_Bundle>    Bundle class
- * @param <T_Ciph>      Cipher class
  * @author Paul Schaub
  */
-public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
+public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_ECPub, T_Bundle> {
     private static final Logger LOGGER = Logger.getLogger(OmemoKeyUtil.class.getName());
 
     public final Bundle BUNDLE = new Bundle();
@@ -63,7 +57,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
          * @return identityKey
          * @throws CorruptedOmemoKeyException if the key is damaged/malformed
          */
-        public T_IdKey identityKey(OmemoBundleVAxolotlElement bundle) throws CorruptedOmemoKeyException {
+        public T_IdKey identityKey(OmemoBundleElement bundle) throws CorruptedOmemoKeyException {
             return identityKeyFromBytes(bundle.getIdentityKey());
         }
 
@@ -74,7 +68,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
          * @return singedPreKey
          * @throws CorruptedOmemoKeyException if the key is damaged/malformed
          */
-        public T_ECPub signedPreKeyPublic(OmemoBundleVAxolotlElement bundle) throws CorruptedOmemoKeyException {
+        public T_ECPub signedPreKeyPublic(OmemoBundleElement bundle) throws CorruptedOmemoKeyException {
             return signedPreKeyPublicFromBytes(bundle.getSignedPreKey());
         }
 
@@ -84,7 +78,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
          * @param bundle OmemoBundleElement
          * @return signedPreKeyId
          */
-        public int signedPreKeyId(OmemoBundleVAxolotlElement bundle) {
+        public int signedPreKeyId(OmemoBundleElement bundle) {
             return bundle.getSignedPreKeyId();
         }
 
@@ -94,7 +88,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
          * @param bundle OmemoBundleElement
          * @return signature
          */
-        public byte[] signedPreKeySignature(OmemoBundleVAxolotlElement bundle) {
+        public byte[] signedPreKeySignature(OmemoBundleElement bundle) {
             return bundle.getSignedPreKeySignature();
         }
 
@@ -106,7 +100,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
          * @return the preKey
          * @throws CorruptedOmemoKeyException when the key cannot be parsed from bytes
          */
-        public T_ECPub preKeyPublic(OmemoBundleVAxolotlElement bundle, int keyId) throws CorruptedOmemoKeyException {
+        public T_ECPub preKeyPublic(OmemoBundleElement bundle, int keyId) throws CorruptedOmemoKeyException {
             return preKeyPublicFromBytes(bundle.getPreKey(keyId));
         }
 
@@ -120,7 +114,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
          * @return a HashMap with one T_Bundle per preKey and the preKeyId as key
          * @throws CorruptedOmemoKeyException when one of the keys cannot be parsed
          */
-        public HashMap<Integer, T_Bundle> bundles(OmemoBundleVAxolotlElement bundle, OmemoDevice contact) throws CorruptedOmemoKeyException {
+        public HashMap<Integer, T_Bundle> bundles(OmemoBundleElement bundle, OmemoDevice contact) throws CorruptedOmemoKeyException {
             HashMap<Integer, T_Bundle> bundles = new HashMap<>();
             for (int deviceId : bundle.getPreKeys().keySet()) {
                 try {
@@ -207,7 +201,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
      * @param count   how many keys do we want to generate
      * @return Map of new preKeys
      */
-    public abstract HashMap<Integer, T_PreKey> generateOmemoPreKeys(int startId, int count);
+    public abstract TreeMap<Integer, T_PreKey> generateOmemoPreKeys(int startId, int count);
 
     /**
      * Generate a new signed preKey.
@@ -259,7 +253,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
      * @return PreKeyBundle (T_PreKey)
      * @throws CorruptedOmemoKeyException if some key is damaged or malformed
      */
-    public abstract T_Bundle bundleFromOmemoBundle(OmemoBundleVAxolotlElement bundle, OmemoDevice contact, int keyId) throws CorruptedOmemoKeyException;
+    public abstract T_Bundle bundleFromOmemoBundle(OmemoBundleElement bundle, OmemoDevice contact, int keyId) throws CorruptedOmemoKeyException;
 
     /**
      * Extract the signature from a signedPreKey.
@@ -330,7 +324,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
      * @param preKeyHashMap HashMap of preKeys
      * @return HashMap of byte arrays but with the same keyIds as key
      */
-    public HashMap<Integer, byte[]> preKeyPublisKeysForBundle(HashMap<Integer, T_PreKey> preKeyHashMap) {
+    public HashMap<Integer, byte[]> preKeyPublicKeysForBundle(TreeMap<Integer, T_PreKey> preKeyHashMap) {
         HashMap<Integer, byte[]> out = new HashMap<>();
         for (Map.Entry<Integer, T_PreKey> e : preKeyHashMap.entrySet()) {
             out.put(e.getKey(), preKeyForBundle(e.getValue()));
@@ -341,7 +335,7 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
     /**
      * Prepare a public signedPreKey for transport in a bundle.
      *
-     * @param signedPreKey signedPrekey
+     * @param signedPreKey signedPreKey
      * @return signedPreKey as byte array
      */
     public abstract byte[] signedPreKeyPublicForBundle(T_SigPreKey signedPreKey);
@@ -352,32 +346,14 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
      * @param identityKey identityKey
      * @return fingerprint of the key
      */
-    public abstract OmemoFingerprint getFingerprint(T_IdKey identityKey);
+    public abstract OmemoFingerprint getFingerprintOfIdentityKey(T_IdKey identityKey);
 
     /**
-     * Create a new crypto-specific Session object.
-     *
-     * @param omemoManager  omemoManager of our device.
-     * @param omemoStore    omemoStore where we can save the session, get keys from etc.
-     * @param from          the device we want to create the session with.
-     * @return a new session
+     * Returns the fingerprint of the public key of an identityKeyPair.
+     * @param identityKeyPair IdentityKeyPair.
+     * @return fingerprint of the public key.
      */
-    public abstract OmemoSession<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph>
-    createOmemoSession(OmemoManager omemoManager, OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore,
-                       OmemoDevice from);
-
-    /**
-     * Create a new concrete OmemoSession with a contact.
-     *
-     * @param omemoManager  omemoManager of our device.
-     * @param omemoStore    omemoStore
-     * @param device        device to establish the session with
-     * @param identityKey   identityKey of the device
-     * @return concrete OmemoSession
-     */
-    public abstract OmemoSession<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph>
-    createOmemoSession(OmemoManager omemoManager, OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore,
-                       OmemoDevice device, T_IdKey identityKey);
+    public abstract OmemoFingerprint getFingerprintOfIdentityKeyPair(T_IdKeyPair identityKeyPair);
 
     /**
      * Deserialize a raw OMEMO Session from bytes.
@@ -396,43 +372,6 @@ public abstract class OmemoKeyUtil<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
      */
     public abstract byte[] rawSessionToBytes(T_Sess session);
 
-    /**
-     * Convert an OmemoDevice to a crypto-lib specific contact format.
-     *
-     * @param contact omemoContact
-     * @return crypto-lib specific contact object
-     */
-    public abstract T_Addr omemoDeviceAsAddress(OmemoDevice contact);
-
-    /**
-     * Convert a crypto-lib specific contact object into an OmemoDevice.
-     *
-     * @param address contact
-     * @return as OmemoDevice
-     * @throws XmppStringprepException if the address is not a valid BareJid
-     */
-    public abstract OmemoDevice addressAsOmemoDevice(T_Addr address) throws XmppStringprepException;
-
-    public static String prettyFingerprint(OmemoFingerprint fingerprint) {
-        return prettyFingerprint(fingerprint.toString());
-    }
-
-    /**
-     * Split the fingerprint in blocks of 8 characters with spaces between.
-     *
-     * @param ugly fingerprint as continuous string
-     * @return fingerprint with spaces for better readability
-     */
-    public static String prettyFingerprint(String ugly) {
-        if (ugly == null) return null;
-        String pretty = "";
-        for (int i = 0; i < 8; i++) {
-            if (i != 0) pretty += " ";
-            pretty += ugly.substring(8 * i, 8 * (i + 1));
-        }
-        return pretty;
-    }
-
     /**
      * Add integers modulo MAX_VALUE.
      *

smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/util/OmemoMessageBuilder.java

@@ -28,7 +28,6 @@ import java.security.NoSuchAlgorithmException;
 import java.security.NoSuchProviderException;
 import java.security.SecureRandom;
 import java.util.ArrayList;
-
 import javax.crypto.BadPaddingException;
 import javax.crypto.Cipher;
 import javax.crypto.IllegalBlockSizeException;
@@ -39,16 +38,21 @@ import javax.crypto.spec.IvParameterSpec;
 import javax.crypto.spec.SecretKeySpec;
 
 import org.jivesoftware.smack.util.StringUtils;
-
-import org.jivesoftware.smackx.omemo.OmemoManager;
-import org.jivesoftware.smackx.omemo.OmemoStore;
-import org.jivesoftware.smackx.omemo.element.OmemoVAxolotlElement;
+import org.jivesoftware.smackx.omemo.OmemoRatchet;
+import org.jivesoftware.smackx.omemo.OmemoService;
+import org.jivesoftware.smackx.omemo.element.OmemoElement;
+import org.jivesoftware.smackx.omemo.element.OmemoElement_VAxolotl;
+import org.jivesoftware.smackx.omemo.element.OmemoHeaderElement_VAxolotl;
+import org.jivesoftware.smackx.omemo.element.OmemoKeyElement;
 import org.jivesoftware.smackx.omemo.exceptions.CorruptedOmemoKeyException;
-import org.jivesoftware.smackx.omemo.exceptions.CryptoFailedException;
+import org.jivesoftware.smackx.omemo.exceptions.NoIdentityKeyException;
 import org.jivesoftware.smackx.omemo.exceptions.UndecidedOmemoIdentityException;
+import org.jivesoftware.smackx.omemo.exceptions.UntrustedOmemoIdentityException;
 import org.jivesoftware.smackx.omemo.internal.CiphertextTuple;
 import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
-import org.jivesoftware.smackx.omemo.internal.OmemoSession;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
+import org.jivesoftware.smackx.omemo.trust.OmemoTrustCallback;
+
 
 /**
  * Class used to build OMEMO messages.
@@ -65,22 +69,27 @@ import org.jivesoftware.smackx.omemo.internal.OmemoSession;
  * @author Paul Schaub
  */
 public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
-    private final OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore;
-    private final OmemoManager omemoManager;
 
-    private byte[] messageKey = generateKey();
-    private byte[] initializationVector = generateIv();
+    private final OmemoDevice userDevice;
+    private final OmemoRatchet<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> ratchet;
+    private final OmemoTrustCallback trustCallback;
+
+    private byte[] messageKey;
+    private final byte[] initializationVector;
 
     private byte[] ciphertextMessage;
-    private final ArrayList<OmemoVAxolotlElement.OmemoHeader.Key> keys = new ArrayList<>();
+    private final ArrayList<OmemoKeyElement> keys = new ArrayList<>();
 
     /**
-     * Create a OmemoMessageBuilder.
+     * Create an OmemoMessageBuilder.
+     *
+     * @param userDevice our OmemoDevice
+     * @param callback trustCallback for querying trust decisions
+     * @param ratchet our OmemoRatchet
+     * @param aesKey aes message key used for message encryption
+     * @param iv initialization vector used for message encryption
+     * @param message message we want to send
      *
-     * @param omemoManager      OmemoManager of our device.
-     * @param omemoStore        OmemoStore.
-     * @param aesKey            AES key that will be transported to the recipient. This is used eg. to encrypt the body.
-     * @param iv                IV
      * @throws NoSuchPaddingException
      * @throws BadPaddingException
      * @throws InvalidKeyException
@@ -90,23 +99,31 @@ public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_
      * @throws NoSuchProviderException
      * @throws InvalidAlgorithmParameterException
      */
-    public OmemoMessageBuilder(OmemoManager omemoManager,
-                               OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore,
-                               byte[] aesKey, byte[] iv)
-            throws NoSuchPaddingException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, IllegalBlockSizeException,
+    public OmemoMessageBuilder(OmemoDevice userDevice,
+                               OmemoTrustCallback callback,
+                               OmemoRatchet<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> ratchet,
+                               byte[] aesKey,
+                               byte[] iv,
+                               String message)
+            throws NoSuchPaddingException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException,
+            IllegalBlockSizeException,
             UnsupportedEncodingException, NoSuchProviderException, InvalidAlgorithmParameterException {
-        this.omemoStore = omemoStore;
-        this.omemoManager = omemoManager;
+        this.userDevice = userDevice;
+        this.trustCallback = callback;
+        this.ratchet = ratchet;
         this.messageKey = aesKey;
         this.initializationVector = iv;
+        setMessage(message);
     }
 
     /**
-     * Create a new OmemoMessageBuilder with random IV and AES key.
+     * Create an OmemoMessageBuilder.
+     *
+     * @param userDevice our OmemoDevice
+     * @param callback trustCallback for querying trust decisions
+     * @param ratchet our OmemoRatchet
+     * @param message message we want to send
      *
-     * @param omemoManager  omemoManager of our device.
-     * @param omemoStore    omemoStore.
-     * @param message       Messages body.
      * @throws NoSuchPaddingException
      * @throws BadPaddingException
      * @throws InvalidKeyException
@@ -116,30 +133,32 @@ public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_
      * @throws NoSuchProviderException
      * @throws InvalidAlgorithmParameterException
      */
-    public OmemoMessageBuilder(OmemoManager omemoManager,
-                               OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> omemoStore, String message)
+    public OmemoMessageBuilder(OmemoDevice userDevice,
+                               OmemoTrustCallback callback,
+                               OmemoRatchet<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> ratchet,
+                               String message)
             throws NoSuchPaddingException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, IllegalBlockSizeException,
             UnsupportedEncodingException, NoSuchProviderException, InvalidAlgorithmParameterException {
-        this.omemoManager = omemoManager;
-        this.omemoStore = omemoStore;
-        this.setMessage(message);
+        this(userDevice, callback, ratchet, generateKey(KEYTYPE, KEYLENGTH), generateIv(), message);
     }
 
     /**
-     * Create an AES messageKey and use it to encrypt the message.
-     * Optionally append the Auth Tag of the encrypted message to the messageKey afterwards.
+     * Encrypt the message with the aes key.
+     * Move the AuthTag from the end of the cipherText to the end of the messageKey afterwards.
+     * This prevents an attacker which compromised one recipient device to switch out the cipherText for other recipients.
+     * @see <a href="https://conversations.im/omemo/audit.pdf">OMEMO security audit</a>.
      *
-     * @param message content of the message
-     * @throws NoSuchPaddingException               When no Cipher could be instantiated.
-     * @throws NoSuchAlgorithmException             when no Cipher could be instantiated.
-     * @throws NoSuchProviderException              when BouncyCastle could not be found.
-     * @throws InvalidAlgorithmParameterException   when the Cipher could not be initialized
-     * @throws InvalidKeyException                  when the generated key is invalid
-     * @throws UnsupportedEncodingException         when UTF8 is unavailable
-     * @throws BadPaddingException                  when cipher.doFinal gets wrong padding
-     * @throws IllegalBlockSizeException            when cipher.doFinal gets wrong Block size.
+     * @param message plaintext message
+     * @throws NoSuchPaddingException
+     * @throws NoSuchAlgorithmException
+     * @throws NoSuchProviderException
+     * @throws InvalidAlgorithmParameterException
+     * @throws InvalidKeyException
+     * @throws UnsupportedEncodingException
+     * @throws BadPaddingException
+     * @throws IllegalBlockSizeException
      */
-    public void setMessage(String message) throws NoSuchPaddingException, NoSuchAlgorithmException, NoSuchProviderException, InvalidAlgorithmParameterException, InvalidKeyException, UnsupportedEncodingException, BadPaddingException, IllegalBlockSizeException {
+    private void setMessage(String message) throws NoSuchPaddingException, NoSuchAlgorithmException, NoSuchProviderException, InvalidAlgorithmParameterException, InvalidKeyException, UnsupportedEncodingException, BadPaddingException, IllegalBlockSizeException {
         if (message == null) {
             return;
         }
@@ -159,50 +178,71 @@ public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_
         byte[] clearKeyWithAuthTag = new byte[messageKey.length + 16];
         byte[] cipherTextWithoutAuthTag = new byte[ciphertext.length - 16];
 
-        System.arraycopy(messageKey, 0, clearKeyWithAuthTag, 0, 16);
-        System.arraycopy(ciphertext, 0, cipherTextWithoutAuthTag, 0, cipherTextWithoutAuthTag.length);
-        System.arraycopy(ciphertext, ciphertext.length - 16, clearKeyWithAuthTag, 16, 16);
+        moveAuthTag(messageKey, ciphertext, clearKeyWithAuthTag, cipherTextWithoutAuthTag);
 
         ciphertextMessage = cipherTextWithoutAuthTag;
         messageKey = clearKeyWithAuthTag;
     }
 
     /**
-     * Add a new recipient device to the message.
+     * Move the auth tag from the end of the cipherText to the messageKey.
      *
-     * @param device recipient device
-     * @throws CryptoFailedException                when encrypting the messageKey fails
-     * @throws UndecidedOmemoIdentityException
-     * @throws CorruptedOmemoKeyException
+     * @param messageKey source messageKey without authTag
+     * @param cipherText source cipherText with authTag
+     * @param messageKeyWithAuthTag destination messageKey with authTag
+     * @param cipherTextWithoutAuthTag destination cipherText without authTag
      */
-    public void addRecipient(OmemoDevice device) throws CryptoFailedException, UndecidedOmemoIdentityException, CorruptedOmemoKeyException {
-        addRecipient(device, false);
+    static void moveAuthTag(byte[] messageKey,
+                            byte[] cipherText,
+                            byte[] messageKeyWithAuthTag,
+                            byte[] cipherTextWithoutAuthTag) {
+        // Check dimensions of arrays
+        if (messageKeyWithAuthTag.length != messageKey.length + 16) {
+            throw new IllegalArgumentException("Length of messageKeyWithAuthTag must be length of messageKey + " +
+                    "length of AuthTag (16)");
+        }
+
+        if (cipherTextWithoutAuthTag.length != cipherText.length - 16) {
+            throw new IllegalArgumentException("Length of cipherTextWithoutAuthTag must be length of cipherText " +
+            "- length of AuthTag (16)");
+        }
+
+        // Move auth tag from cipherText to messageKey
+        System.arraycopy(messageKey, 0, messageKeyWithAuthTag, 0, 16);
+        System.arraycopy(cipherText, 0, cipherTextWithoutAuthTag, 0, cipherTextWithoutAuthTag.length);
+        System.arraycopy(cipherText, cipherText.length - 16, messageKeyWithAuthTag, 16, 16);
     }
 
     /**
      * Add a new recipient device to the message.
-     * @param device recipient device
-     * @param ignoreTrust ignore current trust state? Useful for keyTransportMessages that are sent to repair a session
-     * @throws CryptoFailedException
-     * @throws UndecidedOmemoIdentityException
-     * @throws CorruptedOmemoKeyException
+     *
+     * @param contactsDevice device of the recipient
+     * @throws NoIdentityKeyException if we have no identityKey of that device. Can be fixed by fetching and
+     *                                processing the devices bundle.
+     * @throws CorruptedOmemoKeyException if the identityKey of that device is corrupted.
+     * @throws UndecidedOmemoIdentityException if the user hasn't yet decided whether to trust that device or not.
+     * @throws UntrustedOmemoIdentityException if the user has decided not to trust that device.
      */
-    public void addRecipient(OmemoDevice device, boolean ignoreTrust) throws
-            CryptoFailedException, UndecidedOmemoIdentityException, CorruptedOmemoKeyException {
-        OmemoSession<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> session =
-                omemoStore.getOmemoSessionOf(omemoManager, device);
+    public void addRecipient(OmemoDevice contactsDevice)
+            throws NoIdentityKeyException, CorruptedOmemoKeyException, UndecidedOmemoIdentityException,
+            UntrustedOmemoIdentityException {
 
-        if (session != null) {
-            if (!ignoreTrust && !omemoStore.isDecidedOmemoIdentity(omemoManager, device, session.getIdentityKey())) {
-                // Warn user of undecided device
-                throw new UndecidedOmemoIdentityException(device);
-            }
+        OmemoFingerprint fingerprint;
+        fingerprint = OmemoService.getInstance().getOmemoStoreBackend().getFingerprint(userDevice, contactsDevice);
+
+        switch (trustCallback.getTrust(contactsDevice, fingerprint)) {
+
+            case undecided:
+                throw new UndecidedOmemoIdentityException(contactsDevice);
+
+            case trusted:
+                CiphertextTuple encryptedKey = ratchet.doubleRatchetEncrypt(contactsDevice, messageKey);
+                keys.add(new OmemoKeyElement(encryptedKey.getCiphertext(), contactsDevice.getDeviceId(), encryptedKey.isPreKeyMessage()));
+                break;
+
+            case untrusted:
+                throw new UntrustedOmemoIdentityException(contactsDevice, fingerprint);
 
-            if (!ignoreTrust && omemoStore.isTrustedOmemoIdentity(omemoManager, device, session.getIdentityKey())) {
-                // Encrypt key and save to header
-                CiphertextTuple encryptedKey = session.encryptMessageKey(messageKey);
-                keys.add(new OmemoVAxolotlElement.OmemoHeader.Key(encryptedKey.getCiphertext(), device.getDeviceId(), encryptedKey.isPreKeyMessage()));
-            }
         }
     }
 
@@ -211,24 +251,26 @@ public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_
      *
      * @return OmemoMessageElement
      */
-    public OmemoVAxolotlElement finish() {
-        OmemoVAxolotlElement.OmemoHeader header = new OmemoVAxolotlElement.OmemoHeader(
-                omemoManager.getDeviceId(),
+    public OmemoElement finish() {
+        OmemoHeaderElement_VAxolotl header = new OmemoHeaderElement_VAxolotl(
+                userDevice.getDeviceId(),
                 keys,
                 initializationVector
         );
-        return new OmemoVAxolotlElement(header, ciphertextMessage);
+        return new OmemoElement_VAxolotl(header, ciphertextMessage);
     }
 
     /**
      * Generate a new AES key used to encrypt the message.
      *
+     * @param keyType Key Type
+     * @param keyLength Key Length in bit
      * @return new AES key
      * @throws NoSuchAlgorithmException
      */
-    public static byte[] generateKey() throws NoSuchAlgorithmException {
-        KeyGenerator generator = KeyGenerator.getInstance(KEYTYPE);
-        generator.init(KEYLENGTH);
+    public static byte[] generateKey(String keyType, int keyLength) throws NoSuchAlgorithmException {
+        KeyGenerator generator = KeyGenerator.getInstance(keyType);
+        generator.init(keyLength);
         return generator.generateKey().getEncoded();
     }
 
@@ -243,12 +285,4 @@ public class OmemoMessageBuilder<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_
         random.nextBytes(iv);
         return iv;
     }
-
-    public byte[] getCiphertextMessage() {
-        return ciphertextMessage;
-    }
-
-    public byte[] getMessageKey() {
-        return messageKey;
-    }
 }

smack-omemo/src/test/java/org/jivesoftware/smack/omemo/OmemoKeyUtilTest.java

@@ -1,47 +0,0 @@
-/**
- *
- * Copyright 2017 Paul Schaub
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jivesoftware.smack.omemo;
-
-import static junit.framework.TestCase.assertEquals;
-
-import org.jivesoftware.smackx.omemo.util.OmemoKeyUtil;
-
-import org.junit.Test;
-
-/**
- * Test KeyUtil functions.
- *
- * @author Paul Schaub
- */
-public class OmemoKeyUtilTest {
-
-    @Test
-    public void testAddInBounds() {
-        int high = Integer.MAX_VALUE - 2;
-        int max = Integer.MAX_VALUE;
-        assertEquals(OmemoKeyUtil.addInBounds(high, 3), 1);
-        assertEquals(OmemoKeyUtil.addInBounds(1,2), 3);
-        assertEquals(OmemoKeyUtil.addInBounds(max, 5), 5);
-    }
-
-    @Test
-    public void testPrettyFingerprint() {
-        String ugly = "FFFFFFFFEEEEEEEEDDDDDDDDCCCCCCCCBBBBBBBBAAAAAAAA9999999988888888";
-        String pretty = OmemoKeyUtil.prettyFingerprint(ugly);
-        assertEquals(pretty, "FFFFFFFF EEEEEEEE DDDDDDDD CCCCCCCC BBBBBBBB AAAAAAAA 99999999 88888888");
-    }
-}

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/DeviceListTest.java

@@ -14,15 +14,16 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 import java.util.HashSet;
 import java.util.Set;
 
-import org.jivesoftware.smackx.omemo.internal.CachedDeviceList;
+import org.jivesoftware.smackx.omemo.internal.OmemoCachedDeviceList;
 
 import org.junit.Test;
 
@@ -42,7 +43,7 @@ public class DeviceListTest {
      */
     @Test
     public void mergeDeviceListsTest() {
-        CachedDeviceList cached = new CachedDeviceList();
+        OmemoCachedDeviceList cached = new OmemoCachedDeviceList();
         assertNotNull(cached.getActiveDevices());
         assertNotNull(cached.getInactiveDevices());
 
@@ -67,5 +68,11 @@ public class DeviceListTest {
                 !cached.getInactiveDevices().contains(4));
 
         assertTrue(cached.getAllDevices().size() == 4);
+
+        assertFalse(cached.contains(17));
+        cached.addDevice(17);
+        assertTrue(cached.getActiveDevices().contains(17));
+
+        assertNotNull(cached.toString());
     }
 }

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoBundleVAxolotlElementTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static junit.framework.TestCase.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -27,7 +27,7 @@ import org.jivesoftware.smack.test.util.TestUtils;
 import org.jivesoftware.smack.util.StringUtils;
 import org.jivesoftware.smack.util.stringencoder.Base64;
 
-import org.jivesoftware.smackx.omemo.element.OmemoBundleVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoBundleElement_VAxolotl;
 import org.jivesoftware.smackx.omemo.provider.OmemoBundleVAxolotlProvider;
 
 import org.junit.Test;
@@ -50,7 +50,7 @@ public class OmemoBundleVAxolotlElementTest extends SmackTestSuite {
         preKeysB64.put(preKeyId1, preKey1B64);
         preKeysB64.put(preKeyId2, preKey2B64);
 
-        OmemoBundleVAxolotlElement bundle = new OmemoBundleVAxolotlElement(signedPreKeyId,
+        OmemoBundleElement_VAxolotl bundle = new OmemoBundleElement_VAxolotl(signedPreKeyId,
                 signedPreKeyB64, signedPreKeySigB64, identityKeyB64, preKeysB64);
 
         assertEquals("ElementName must match.", "bundle", bundle.getElementName());
@@ -85,7 +85,7 @@ public class OmemoBundleVAxolotlElementTest extends SmackTestSuite {
         byte[] firstPreKey = "FirstPreKey".getBytes(StringUtils.UTF8);
         byte[] secondPreKey = "SecondPreKey".getBytes(StringUtils.UTF8);
 
-        OmemoBundleVAxolotlElement parsed = new OmemoBundleVAxolotlProvider().parse(TestUtils.getParser(actual));
+        OmemoBundleElement_VAxolotl parsed = new OmemoBundleVAxolotlProvider().parse(TestUtils.getParser(actual));
 
         assertTrue("B64-decoded signedPreKey must match.", Arrays.equals(signedPreKey, parsed.getSignedPreKey()));
         assertEquals("SignedPreKeyId must match", signedPreKeyId, parsed.getSignedPreKeyId());

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoConfigurationTest.java

@@ -14,14 +14,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
-
-import java.io.File;
-
-import org.jivesoftware.smackx.omemo.OmemoConfiguration;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 import junit.framework.TestCase;
 import org.junit.Test;
@@ -34,25 +31,6 @@ public class OmemoConfigurationTest {
     @Test
     public void omemoConfigurationTest() {
         @SuppressWarnings("unused") OmemoConfiguration configuration = new OmemoConfiguration();
-        // Default Store Path
-        File storePath = new File("test");
-        assertNull("getFileBasedOmemoStoreDefaultPath MUST return null at this point.",
-                OmemoConfiguration.getFileBasedOmemoStoreDefaultPath());
-        OmemoConfiguration.setFileBasedOmemoStoreDefaultPath(storePath);
-        assertEquals("FileBasedOmemoStoreDefaultPath must equal the one we set.", storePath.getAbsolutePath(),
-                OmemoConfiguration.getFileBasedOmemoStoreDefaultPath().getAbsolutePath());
-
-        // EME
-        OmemoConfiguration.setAddEmeEncryptionHint(false);
-        assertEquals(false, OmemoConfiguration.getAddEmeEncryptionHint());
-        OmemoConfiguration.setAddEmeEncryptionHint(true);
-        assertEquals(true, OmemoConfiguration.getAddEmeEncryptionHint());
-
-        // MAM
-        OmemoConfiguration.setAddMAMStorageProcessingHint(false);
-        assertEquals(false, OmemoConfiguration.getAddMAMStorageProcessingHint());
-        OmemoConfiguration.setAddMAMStorageProcessingHint(true);
-        assertEquals(true, OmemoConfiguration.getAddMAMStorageProcessingHint());
 
         // Body hint
         OmemoConfiguration.setAddOmemoHintBody(false);
@@ -109,5 +87,17 @@ public class OmemoConfigurationTest {
         } catch (IllegalArgumentException e) {
             // Expected
         }
+
+        // Repair broken sessions
+        OmemoConfiguration.setRepairBrokenSessionsWithPrekeyMessages(false);
+        assertFalse(OmemoConfiguration.getRepairBrokenSessionsWithPreKeyMessages());
+        OmemoConfiguration.setRepairBrokenSessionsWithPrekeyMessages(true);
+        assertTrue(OmemoConfiguration.getRepairBrokenSessionsWithPreKeyMessages());
+
+        // Complete fresh sessions
+        OmemoConfiguration.setCompleteSessionWithEmptyMessage(false);
+        assertFalse(OmemoConfiguration.getCompleteSessionWithEmptyMessage());
+        OmemoConfiguration.setCompleteSessionWithEmptyMessage(true);
+        assertTrue(OmemoConfiguration.getCompleteSessionWithEmptyMessage());
     }
 }

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoDeviceListVAxolotlElementTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -24,7 +24,7 @@ import java.util.HashSet;
 import org.jivesoftware.smack.test.util.SmackTestSuite;
 import org.jivesoftware.smack.test.util.TestUtils;
 
-import org.jivesoftware.smackx.omemo.element.OmemoDeviceListVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoDeviceListElement_VAxolotl;
 import org.jivesoftware.smackx.omemo.provider.OmemoDeviceListVAxolotlProvider;
 
 import org.junit.Test;
@@ -41,11 +41,11 @@ public class OmemoDeviceListVAxolotlElementTest extends SmackTestSuite {
         ids.add(1234);
         ids.add(9876);
 
-        OmemoDeviceListVAxolotlElement element = new OmemoDeviceListVAxolotlElement(ids);
+        OmemoDeviceListElement_VAxolotl element = new OmemoDeviceListElement_VAxolotl(ids);
         String xml = element.toXML(null).toString();
 
         XmlPullParser parser = TestUtils.getParser(xml);
-        OmemoDeviceListVAxolotlElement parsed = new OmemoDeviceListVAxolotlProvider().parse(parser);
+        OmemoDeviceListElement_VAxolotl parsed = new OmemoDeviceListVAxolotlProvider().parse(parser);
 
         assertTrue("Parsed element must equal the original.", parsed.getDeviceIds().equals(element.getDeviceIds()));
         assertEquals("Generated XML must match.",

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoDeviceTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoExceptionsTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static junit.framework.TestCase.fail;
 import static org.junit.Assert.assertEquals;

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoFingerprintTest.java

@@ -14,12 +14,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static junit.framework.TestCase.assertEquals;
 import static junit.framework.TestCase.assertNotSame;
 
-import org.jivesoftware.smackx.omemo.OmemoFingerprint;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
 
 import org.junit.Test;
 

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoServiceTest.java

@@ -0,0 +1,91 @@
+/**
+ *
+ * Copyright 2017 Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo;
+
+import static junit.framework.TestCase.assertFalse;
+import static junit.framework.TestCase.assertTrue;
+
+import java.util.Date;
+import java.util.HashSet;
+
+import org.jivesoftware.smack.test.util.SmackTestSuite;
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+
+import org.junit.Test;
+import org.jxmpp.jid.impl.JidCreate;
+import org.jxmpp.stringprep.XmppStringprepException;
+
+public class OmemoServiceTest extends SmackTestSuite {
+
+    private static final long ONE_HOUR = 1000L * 60 * 60;
+    private static final int IGNORE_STALE = OmemoConfiguration.getIgnoreStaleDevicesAfterHours();
+    private static final int DELETE_STALE = OmemoConfiguration.getDeleteStaleDevicesAfterHours();
+
+    @Test(expected = IllegalStateException.class)
+    public void getInstanceFailsWhenNullTest() {
+        OmemoService.getInstance();
+    }
+
+    @Test
+    public void isServiceRegisteredTest() {
+        assertFalse(OmemoService.isServiceRegistered());
+    }
+
+    /**
+     * Test correct functionality of isStale method.
+     * @throws XmppStringprepException
+     */
+    @Test
+    public void isStaleDeviceTest() throws XmppStringprepException {
+        OmemoDevice user = new OmemoDevice(JidCreate.bareFrom("alice@wonderland.lit"), 123);
+        OmemoDevice other = new OmemoDevice(JidCreate.bareFrom("bob@builder.tv"), 444);
+
+        Date now = new Date();
+        Date ignoreMe = new Date(now.getTime() - ((IGNORE_STALE + 1) * ONE_HOUR));
+        Date deleteMe = new Date(now.getTime() - ((DELETE_STALE + 1) * ONE_HOUR));
+        Date imFine = new Date(now.getTime() - ONE_HOUR);
+
+        // One hour "old" devices are (probably) not not stale
+        assertFalse(OmemoService.isStale(user, other, imFine, IGNORE_STALE));
+
+        // Devices one hour "older" than max ages are stale
+        assertTrue(OmemoService.isStale(user, other, ignoreMe, IGNORE_STALE));
+        assertTrue(OmemoService.isStale(user, other, deleteMe, DELETE_STALE));
+
+        // Own device is never stale, no matter how old
+        assertFalse(OmemoService.isStale(user, user, deleteMe, DELETE_STALE));
+
+        // Always return false if date is null.
+        assertFalse(OmemoService.isStale(user, other, null, DELETE_STALE));
+    }
+
+    @Test
+    public void removeOurDeviceTest() throws XmppStringprepException {
+        OmemoDevice a = new OmemoDevice(JidCreate.bareFrom("a@b.c"), 123);
+        OmemoDevice b = new OmemoDevice(JidCreate.bareFrom("a@b.c"), 124);
+
+        HashSet<OmemoDevice> devices = new HashSet<>();
+        devices.add(a); devices.add(b);
+
+        assertTrue(devices.contains(a));
+        assertTrue(devices.contains(b));
+        OmemoService.removeOurDevice(a, devices);
+
+        assertFalse(devices.contains(a));
+        assertTrue(devices.contains(b));
+    }
+}

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoStoreTest.java

@@ -0,0 +1,345 @@
+/**
+ *
+ * Copyright Paul Schaub
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo;
+
+import static junit.framework.TestCase.assertEquals;
+import static junit.framework.TestCase.assertFalse;
+import static junit.framework.TestCase.assertNotNull;
+import static junit.framework.TestCase.assertNotSame;
+import static junit.framework.TestCase.assertNull;
+import static junit.framework.TestCase.assertSame;
+import static junit.framework.TestCase.assertTrue;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.TreeMap;
+
+import org.jivesoftware.smackx.omemo.exceptions.CorruptedOmemoKeyException;
+import org.jivesoftware.smackx.omemo.internal.OmemoCachedDeviceList;
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
+
+import org.junit.AfterClass;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.jxmpp.jid.impl.JidCreate;
+import org.jxmpp.stringprep.XmppStringprepException;
+
+
+public abstract class OmemoStoreTest<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> {
+
+    protected final OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> store;
+    private final OmemoDevice alice, bob;
+
+    private static final TemporaryFolder tmp = initStaticTemp();
+
+    OmemoStoreTest(OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Sess, T_Addr, T_ECPub, T_Bundle, T_Ciph> store)
+            throws XmppStringprepException {
+        this.store = store;
+        alice = new OmemoDevice(JidCreate.bareFrom("alice@wonderland.lit"), 123);
+        bob = new OmemoDevice(JidCreate.bareFrom("bob@builder.tv"), 987);
+    }
+
+    // Tests
+
+    @Test
+    public void keyUtilNotNull() {
+        assertNotNull(store.keyUtil());
+    }
+
+    @Test
+    public void generateOmemoIdentityKeyPairDoesNotReturnNull() {
+        assertNotNull(store.generateOmemoIdentityKeyPair());
+    }
+
+    @Test
+    public void identityKeyFromIdentityKeyPairIsNotNull() {
+        T_IdKeyPair pair = store.generateOmemoIdentityKeyPair();
+        assertNotNull(store.keyUtil().identityKeyFromPair(pair));
+    }
+
+    @Test
+    public void storeLoadRemoveOmemoIdentityKeyPair()
+            throws IOException, CorruptedOmemoKeyException {
+
+        T_IdKeyPair before = store.generateOmemoIdentityKeyPair();
+
+        assertNull(store.loadOmemoIdentityKeyPair(alice));
+        store.storeOmemoIdentityKeyPair(alice, before);
+
+        T_IdKeyPair after = store.loadOmemoIdentityKeyPair(alice);
+        assertNotNull(after);
+
+        // Fingerprints equal
+        assertEquals(store.keyUtil().getFingerprintOfIdentityKeyPair(before),
+                store.keyUtil().getFingerprintOfIdentityKeyPair(after));
+
+        // Byte-representation equals
+        assertTrue(Arrays.equals(
+                store.keyUtil().identityKeyPairToBytes(before),
+                store.keyUtil().identityKeyPairToBytes(after)));
+
+        // Non-existing keypair
+        assertNull("Must return null for non-existing key pairs.", store.loadOmemoIdentityKeyPair(bob));
+
+        // Deleting works
+        store.removeOmemoIdentityKeyPair(alice);
+        assertNull(store.loadOmemoIdentityKeyPair(alice));
+    }
+
+    @Test
+    public void storeLoadRemoveOmemoIdentityKey()
+            throws IOException, CorruptedOmemoKeyException {
+
+        // Create IdentityKeys and get bytes
+        T_IdKey keyA1 = store.keyUtil().identityKeyFromPair(store.generateOmemoIdentityKeyPair());
+        T_IdKey keyB1 = store.keyUtil().identityKeyFromPair(store.generateOmemoIdentityKeyPair());
+        byte[] bytesA1 = store.keyUtil().identityKeyToBytes(keyA1);
+        byte[] bytesB = store.keyUtil().identityKeyToBytes(keyB1);
+
+        // Not null and not of length 0
+        assertNotNull("Serialized identityKey cannot be null.", bytesA1);
+        assertNotNull("Serialized identityKey cannot be null.", bytesB);
+        assertNotSame("Serialized identityKey must be of length > 0.", 0, bytesA1.length);
+        assertNotSame("Serialized identityKey must be of length > 0.", 0, bytesB.length);
+
+        // Keys do not equal
+        assertFalse("Generated IdentityKeys must not be equal (ULTRA unlikely).",
+                Arrays.equals(bytesA1, bytesB));
+
+        // Loading must return null before and not null after saving
+        assertNull("Must return null, the store could not have this key by now.",
+                store.loadOmemoIdentityKey(alice, bob));
+        store.storeOmemoIdentityKey(alice, bob, keyA1);
+        T_IdKey keyA2 = store.loadOmemoIdentityKey(alice, bob);
+        assertNotNull(keyA2);
+
+        // Loaded key must equal stored one
+        byte[] bytesA2 = store.keyUtil().identityKeyToBytes(keyA2);
+        assertTrue("Serialized loaded key must equal serialized stored one.",
+                Arrays.equals(bytesA1, bytesA2));
+
+        // Non-existing keys must return null
+        assertNull("Non-existing keys must be returned as null.", store.loadOmemoIdentityKey(bob, alice));
+
+        // Key must vanish when deleted.
+        store.removeOmemoIdentityKey(alice, bob);
+        assertNull(store.loadOmemoIdentityKey(alice, bob));
+    }
+
+    @Test
+    public void generateOmemoPreKeys() {
+        TreeMap<Integer, T_PreKey> keys = store.generateOmemoPreKeys(31, 49);
+        assertNotNull("Generated data structure must not be null.", keys);
+
+        byte[] lastKey = null;
+
+        for (int i = 31; i <= 79; i++) {
+            assertEquals("Key ids must be ascending order, starting at 31.", Integer.valueOf(i), keys.firstKey());
+            assertNotNull("Every id must match to a key.", keys.get(keys.firstKey()));
+            byte[] bytes = store.keyUtil().preKeyToBytes(keys.get(keys.firstKey()));
+            assertNotNull("Serialized preKey must not be null.", bytes);
+            assertNotSame("Serialized preKey must not be of length 0.", 0, bytes.length);
+
+            if (lastKey != null) {
+                assertFalse("PreKeys MUST NOT be equal.", Arrays.equals(lastKey, bytes));
+            }
+            lastKey = bytes;
+
+            keys.remove(keys.firstKey());
+
+        }
+
+        assertEquals("After deleting 49 keys, there must be no keys left.", 0, keys.size());
+    }
+
+    @Test
+    public void storeLoadRemoveOmemoPreKeys()
+            throws IOException, InterruptedException {
+
+        TreeMap<Integer, T_PreKey> before = store.generateOmemoPreKeys(1, 10);
+        assertEquals("The store must have no prekeys before this test.", 0, store.loadOmemoPreKeys(alice).size());
+
+        store.storeOmemoPreKeys(alice, before);
+        TreeMap<Integer, T_PreKey> after = store.loadOmemoPreKeys(alice);
+        assertNotNull("Loaded preKeys must not be null.", after);
+        assertEquals("Loaded preKey count must equal stored count.", before.size(), after.size());
+
+        // Non-existing key must be returned as null
+        assertNull("Non-existing preKey must be returned as null.", store.loadOmemoPreKey(alice, 10000));
+
+        int last = after.size();
+        for (int i = 1; i <= last; i++) {
+            T_PreKey bKey = before.get(i);
+            T_PreKey aKey = after.get(i);
+
+            assertTrue("Loaded keys must equal stored ones.", Arrays.equals(
+                    store.keyUtil().preKeyToBytes(bKey),
+                    store.keyUtil().preKeyToBytes(aKey)));
+
+            T_PreKey rKey = store.loadOmemoPreKey(alice, i);
+            assertNotNull("Randomly accessed preKeys must not be null.", rKey);
+            assertTrue("Randomly accessed preKeys must equal the stored ones.", Arrays.equals(
+                    store.keyUtil().preKeyToBytes(aKey),
+                    store.keyUtil().preKeyToBytes(rKey)));
+
+            store.removeOmemoPreKey(alice, i);
+            assertNull("PreKey must be null after deletion.", store.loadOmemoPreKey(alice, i));
+        }
+
+        TreeMap<Integer, T_PreKey> postDeletion = store.loadOmemoPreKeys(alice);
+        assertSame("PreKey count must equal 0 after deletion of all keys.", 0, postDeletion.size());
+    }
+
+    @Test
+    public void storeLoadRemoveOmemoSignedPreKeys()
+            throws IOException, CorruptedOmemoKeyException {
+
+        TreeMap<Integer, T_SigPreKey> before = store.loadOmemoSignedPreKeys(alice);
+        assertEquals("At this stage, there must be no signed prekeys in the store.", 0, before.size());
+
+        T_IdKeyPair idp = store.generateOmemoIdentityKeyPair();
+        T_SigPreKey spk = store.generateOmemoSignedPreKey(idp, 125);
+
+        assertNotNull("SignedPreKey must not be null.", spk);
+        assertEquals("ID of signedPreKey must match.", 125, store.keyUtil().signedPreKeyIdFromKey(spk));
+        byte[] bytes = store.keyUtil().signedPreKeyToBytes(spk);
+        assertNotNull("Serialized signedPreKey must not be null", bytes);
+        assertNotSame("Serialized signedPreKey must not be of length 0.", 0, bytes.length);
+
+        // Stored key must equal loaded key
+        store.storeOmemoSignedPreKey(alice, 125, spk);
+        TreeMap<Integer, T_SigPreKey> after = store.loadOmemoSignedPreKeys(alice);
+        assertEquals("We must have exactly 1 signedPreKey now.", 1, after.size());
+        T_SigPreKey spk2 = after.get(after.firstKey());
+        assertEquals("Id of the stored signedPreKey must match the one we stored.",
+                125, store.keyUtil().signedPreKeyIdFromKey(spk2));
+        assertTrue("Serialization of stored and loaded signed preKey must equal.", Arrays.equals(
+                store.keyUtil().signedPreKeyToBytes(spk), store.keyUtil().signedPreKeyToBytes(spk2)));
+
+        // Random access
+        T_SigPreKey rspk = store.loadOmemoSignedPreKey(alice, 125);
+        assertTrue("Serialization of stored and randomly accessed signed preKey must equal.", Arrays.equals(
+                store.keyUtil().signedPreKeyToBytes(spk), store.keyUtil().signedPreKeyToBytes(rspk)));
+        assertNull("Non-existing signedPreKey must be returned as null.",
+                store.loadOmemoSignedPreKey(alice, 10000));
+
+        // Deleting
+        store.removeOmemoSignedPreKey(alice, 125);
+        assertNull("Deleted key must be returned as null.", store.loadOmemoSignedPreKey(alice, 125));
+        assertEquals(0, store.loadOmemoSignedPreKeys(alice).size());
+    }
+
+    @Test
+    public void loadStoreDateOfLastSignedPreKeyRenewal() throws IOException {
+        assertNull("The date of last signed preKey renewal must be null at this stage.",
+                store.getDateOfLastSignedPreKeyRenewal(alice));
+        Date before = new Date();
+        store.setDateOfLastSignedPreKeyRenewal(alice, before);
+        Date after = store.getDateOfLastSignedPreKeyRenewal(alice);
+        assertEquals("Dates must equal.", after, before);
+    }
+
+    @Test
+    public void loadStoreDateOfLastMessageReceived() throws IOException {
+        assertNull("The date of last message received must be null at this stage.",
+                store.getDateOfLastReceivedMessage(alice, bob));
+        Date before = new Date();
+        store.setDateOfLastReceivedMessage(alice, bob, before);
+        Date after = store.getDateOfLastReceivedMessage(alice, bob);
+        assertEquals("Dates must equal.", after, before);
+    }
+
+    @Test
+    public void loadStoreCachedDeviceList() throws IOException {
+        Integer[] active = new Integer[] {1,5,999,10};
+        Integer[] inactive = new Integer[] {6,7,8};
+        OmemoCachedDeviceList before = new OmemoCachedDeviceList(
+                new HashSet<>(Arrays.asList(active)),
+                new HashSet<>(Arrays.asList(inactive)));
+
+        assertNotNull("Loading a non-existent cached deviceList must return an empty list.",
+                store.loadCachedDeviceList(alice, bob.getJid()));
+
+        store.storeCachedDeviceList(alice, bob.getJid(), before);
+        OmemoCachedDeviceList after = store.loadCachedDeviceList(alice, bob.getJid());
+        assertTrue("Loaded deviceList must not be empty", after.getAllDevices().size() != 0);
+
+        assertEquals("Number of entries in active devices must match.", active.length, after.getActiveDevices().size());
+        assertEquals("Number of entries in inactive devices must match.", inactive.length, after.getInactiveDevices().size());
+        assertEquals("Number of total entries must match.", active.length + inactive.length, after.getAllDevices().size());
+
+        for (Integer a : active) {
+            assertTrue(after.getActiveDevices().contains(a));
+            assertTrue(after.getAllDevices().contains(a));
+        }
+
+        for (Integer i : inactive) {
+            assertTrue(after.getInactiveDevices().contains(i));
+            assertTrue(after.getAllDevices().contains(i));
+        }
+
+        store.storeCachedDeviceList(alice, bob.getJid(), new OmemoCachedDeviceList());
+        assertEquals("DeviceList must be empty after overwriting it with empty list.", 0,
+                store.loadCachedDeviceList(alice, bob.getJid()).getAllDevices().size());
+    }
+
+    @Test
+    public void loadAllRawSessionsReturnsEmptyMapTest() {
+        HashMap<Integer, T_Sess> sessions = store.loadAllRawSessionsOf(alice, bob.getJid());
+        assertNotNull(sessions);
+        assertEquals(0, sessions.size());
+    }
+
+    @Test
+    public void loadNonExistentRawSessionReturnsNullTest() {
+        T_Sess session = store.loadRawSession(alice, bob);
+        assertNull(session);
+    }
+
+    @Test
+    public void getFingerprint() throws IOException, CorruptedOmemoKeyException {
+        assertNull("Method must return null for a non-existent fingerprint.", store.getFingerprint(alice));
+        store.storeOmemoIdentityKeyPair(alice, store.generateOmemoIdentityKeyPair());
+        OmemoFingerprint fingerprint = store.getFingerprint(alice);
+        assertNotNull("fingerprint must not be null", fingerprint);
+        assertEquals("Fingerprint must be of length 64", 64, fingerprint.length());
+
+        store.removeOmemoIdentityKeyPair(alice); //clean up
+    }
+
+    // ##############################################################
+    // Workaround for https://github.com/junit-team/junit4/issues/671
+
+    static TemporaryFolder initStaticTemp() {
+        try {
+            return new TemporaryFolder() { { before(); } };
+        } catch (Throwable t) {
+            throw new RuntimeException(t);
+        }
+    }
+
+    @AfterClass
+    public static void cleanup() throws Exception {
+        FileBasedOmemoStore.deleteDirectory(tmp.getRoot());
+    }
+
+    // ##############################################################
+}

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoVAxolotlElementTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static org.junit.Assert.assertEquals;
 
@@ -24,9 +24,9 @@ import org.jivesoftware.smack.test.util.SmackTestSuite;
 import org.jivesoftware.smack.test.util.TestUtils;
 import org.jivesoftware.smack.util.StringUtils;
 import org.jivesoftware.smack.util.stringencoder.Base64;
-
-import org.jivesoftware.smackx.omemo.element.OmemoElement;
-import org.jivesoftware.smackx.omemo.element.OmemoVAxolotlElement;
+import org.jivesoftware.smackx.omemo.element.OmemoElement_VAxolotl;
+import org.jivesoftware.smackx.omemo.element.OmemoHeaderElement_VAxolotl;
+import org.jivesoftware.smackx.omemo.element.OmemoKeyElement;
 import org.jivesoftware.smackx.omemo.provider.OmemoVAxolotlProvider;
 import org.jivesoftware.smackx.omemo.util.OmemoMessageBuilder;
 
@@ -47,12 +47,12 @@ public class OmemoVAxolotlElementTest extends SmackTestSuite {
         int sid = 12131415;
         byte[] iv = OmemoMessageBuilder.generateIv();
 
-        ArrayList<OmemoElement.OmemoHeader.Key> keys = new ArrayList<>();
-        keys.add(new OmemoElement.OmemoHeader.Key(keyData1, keyId1));
-        keys.add(new OmemoElement.OmemoHeader.Key(keyData2, keyId2, true));
+        ArrayList<OmemoKeyElement> keys = new ArrayList<>();
+        keys.add(new OmemoKeyElement(keyData1, keyId1));
+        keys.add(new OmemoKeyElement(keyData2, keyId2, true));
 
-        OmemoVAxolotlElement.OmemoHeader header = new OmemoElement.OmemoHeader(sid, keys, iv);
-        OmemoVAxolotlElement element = new OmemoVAxolotlElement(header, payload);
+        OmemoHeaderElement_VAxolotl header = new OmemoHeaderElement_VAxolotl(sid, keys, iv);
+        OmemoElement_VAxolotl element = new OmemoElement_VAxolotl(header, payload);
 
         String expected =
                 "<encrypted xmlns='eu.siacs.conversations.axolotl'>" +
@@ -69,7 +69,9 @@ public class OmemoVAxolotlElementTest extends SmackTestSuite {
         String actual = element.toXML(null).toString();
         assertEquals("Serialized xml of OmemoElement must match.", expected, actual);
 
-        OmemoVAxolotlElement parsed = new OmemoVAxolotlProvider().parse(TestUtils.getParser(actual));
-        assertEquals("Parsed OmemoElement must equal the original.", element.toXML(null).toString(), parsed.toXML(null).toString());
+        OmemoElement_VAxolotl parsed = new OmemoVAxolotlProvider().parse(TestUtils.getParser(actual));
+        assertEquals("Parsed OmemoElement must equal the original.",
+                element.toXML(null).toString(),
+                parsed.toXML(null).toString());
     }
 }

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/WrapperObjectsTest.java

@@ -14,10 +14,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.jivesoftware.smack.omemo;
+package org.jivesoftware.smackx.omemo;
 
 import static junit.framework.TestCase.assertEquals;
 import static junit.framework.TestCase.assertTrue;
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.KEYLENGTH;
+import static org.jivesoftware.smackx.omemo.util.OmemoConstants.Crypto.KEYTYPE;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -25,35 +27,20 @@ import static org.junit.Assert.assertNotNull;
 import java.security.NoSuchAlgorithmException;
 import java.security.Security;
 
-import org.jivesoftware.smack.packet.Message;
-
 import org.jivesoftware.smackx.omemo.element.OmemoElement;
 import org.jivesoftware.smackx.omemo.exceptions.CryptoFailedException;
 import org.jivesoftware.smackx.omemo.internal.CipherAndAuthTag;
 import org.jivesoftware.smackx.omemo.internal.CiphertextTuple;
-import org.jivesoftware.smackx.omemo.internal.ClearTextMessage;
-import org.jivesoftware.smackx.omemo.internal.IdentityKeyWrapper;
-import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
-import org.jivesoftware.smackx.omemo.internal.OmemoMessageInformation;
 import org.jivesoftware.smackx.omemo.util.OmemoMessageBuilder;
 
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
 import org.junit.Test;
-import org.jxmpp.jid.BareJid;
-import org.jxmpp.jid.impl.JidCreate;
 
 /**
  * Test the identityKeyWrapper.
  */
 public class WrapperObjectsTest {
 
-    @Test
-    public void identityKeyWrapperTest() {
-        Object pseudoKey = new Object();
-        IdentityKeyWrapper wrapper = new IdentityKeyWrapper(pseudoKey);
-        assertEquals(pseudoKey, wrapper.getIdentityKey());
-    }
-
     @Test
     public void ciphertextTupleTest() {
         byte[] c = OmemoMessageBuilder.generateIv();
@@ -67,41 +54,19 @@ public class WrapperObjectsTest {
         assertEquals(OmemoElement.TYPE_OMEMO_MESSAGE, c2.getMessageType());
     }
 
-    @Test
-    public void clearTextMessageTest() throws Exception {
-        Object pseudoKey = new Object();
-        IdentityKeyWrapper wrapper = new IdentityKeyWrapper(pseudoKey);
-        BareJid senderJid = JidCreate.bareFrom("bob@server.tld");
-        OmemoDevice sender = new OmemoDevice(senderJid, 1234);
-        OmemoMessageInformation information = new OmemoMessageInformation(wrapper, sender, OmemoMessageInformation.CARBON.NONE);
-
-        assertTrue("OmemoInformation must state that the message is an OMEMO message.",
-                information.isOmemoMessage());
-        assertEquals(OmemoMessageInformation.CARBON.NONE, information.getCarbon());
-        assertEquals(sender, information.getSenderDevice());
-        assertEquals(wrapper, information.getSenderIdentityKey());
-
-        String body = "Decrypted Body";
-        Message message = new Message(senderJid, body);
-        ClearTextMessage c = new ClearTextMessage(body, message, information);
-
-        assertEquals(message, c.getOriginalMessage());
-        assertEquals(information, c.getMessageInformation());
-        assertEquals(body, c.getBody());
-    }
-
     @Test
     public void cipherAndAuthTagTest() throws NoSuchAlgorithmException, CryptoFailedException {
         Security.addProvider(new BouncyCastleProvider());
-        byte[] key = OmemoMessageBuilder.generateKey();
+        byte[] key = OmemoMessageBuilder.generateKey(KEYTYPE, KEYLENGTH);
         byte[] iv = OmemoMessageBuilder.generateIv();
         byte[] authTag = OmemoMessageBuilder.generateIv();
 
-        CipherAndAuthTag cat = new CipherAndAuthTag(key, iv, authTag);
+        CipherAndAuthTag cat = new CipherAndAuthTag(key, iv, authTag, true);
 
         assertNotNull(cat.getCipher());
         assertArrayEquals(key, cat.getKey());
         assertArrayEquals(iv, cat.getIv());
         assertArrayEquals(authTag, cat.getAuthTag());
+        assertTrue(cat.wasPreKeyEncrypted());
     }
 }

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/util/EphemeralTrustCallback.java

@@ -0,0 +1,59 @@
+/**
+ *
+ * Copyright the original author or authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.util;
+
+import java.util.HashMap;
+
+import org.jivesoftware.smackx.omemo.internal.OmemoDevice;
+import org.jivesoftware.smackx.omemo.trust.OmemoFingerprint;
+import org.jivesoftware.smackx.omemo.trust.OmemoTrustCallback;
+import org.jivesoftware.smackx.omemo.trust.TrustState;
+
+/**
+ * Ephemera Trust Callback used to make trust decisions in tests.
+ */
+public class EphemeralTrustCallback implements OmemoTrustCallback {
+
+    private final HashMap<OmemoDevice, HashMap<OmemoFingerprint, TrustState>> trustStates = new HashMap<>();
+
+    @Override
+    public TrustState getTrust(OmemoDevice device, OmemoFingerprint fingerprint) {
+        HashMap<OmemoFingerprint, TrustState> states = trustStates.get(device);
+
+        if (states != null) {
+            TrustState state = states.get(fingerprint);
+
+            if (state != null) {
+                return state;
+            }
+        }
+
+        return TrustState.undecided;
+    }
+
+    @Override
+    public void setTrust(OmemoDevice device, OmemoFingerprint fingerprint, TrustState state) {
+        HashMap<OmemoFingerprint, TrustState> states = trustStates.get(device);
+
+        if (states == null) {
+            states = new HashMap<>();
+            trustStates.put(device, states);
+        }
+
+        states.put(fingerprint, state);
+    }
+}

smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/util/OmemoMessageBuilderTest.java

@@ -0,0 +1,79 @@
+/**
+ *
+ * Copyright the original author or authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jivesoftware.smackx.omemo.util;
+
+import static org.junit.Assert.assertTrue;
+
+import java.util.Arrays;
+import java.util.Random;
+
+import org.jivesoftware.smack.test.util.SmackTestSuite;
+
+import org.junit.Test;
+
+public class OmemoMessageBuilderTest extends SmackTestSuite {
+
+    private static final byte[] messageKey = new byte[16];
+    private static final byte[] cipherTextWithAuthTag = new byte[35 + 16];
+
+    public OmemoMessageBuilderTest() {
+        Random random = new Random();
+        random.nextBytes(messageKey);
+        random.nextBytes(cipherTextWithAuthTag);
+    }
+
+    @Test
+    public void testMoveAuthTag() {
+        // Extract authTag for testing purposes
+        byte[] authTag = new byte[16];
+        System.arraycopy(cipherTextWithAuthTag, 35, authTag, 0, 16);
+
+        byte[] messageKeyWithAuthTag = new byte[16 + 16];
+        byte[] cipherTextWithoutAuthTag = new byte[35];
+
+        OmemoMessageBuilder.moveAuthTag(messageKey, cipherTextWithAuthTag, messageKeyWithAuthTag, cipherTextWithoutAuthTag);
+
+        // Check if first n - 16 bytes of cipherText got copied over to cipherTextWithoutAuthTag correctly
+        byte[] checkCipherText = new byte[35];
+        System.arraycopy(cipherTextWithAuthTag, 0, checkCipherText, 0, 35);
+        assertTrue(Arrays.equals(checkCipherText, cipherTextWithoutAuthTag));
+
+        byte[] checkMessageKey = new byte[16];
+        System.arraycopy(messageKeyWithAuthTag, 0, checkMessageKey, 0, 16);
+        assertTrue(Arrays.equals(checkMessageKey, messageKey));
+
+        byte[] checkAuthTag = new byte[16];
+        System.arraycopy(messageKeyWithAuthTag, 16, checkAuthTag, 0, 16);
+        assertTrue(Arrays.equals(checkAuthTag, authTag));
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testCheckIllegalMessageKeyWithAuthTagLength() {
+        byte[] illegalMessageKey = new byte[16 + 15]; // too short
+        byte[] cipherTextWithoutAuthTag = new byte[35]; // ok
+
+        OmemoMessageBuilder.moveAuthTag(messageKey, cipherTextWithAuthTag, illegalMessageKey, cipherTextWithoutAuthTag);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testCheckIllegalCipherTextWithoutAuthTagLength() {
+        byte[] messageKeyWithAuthTag = new byte[16 + 16]; // ok
+        byte[] illegalCipherTextWithoutAuthTag = new byte[39]; // too long
+
+        OmemoMessageBuilder.moveAuthTag(messageKey, cipherTextWithAuthTag, messageKeyWithAuthTag, illegalCipherTextWithoutAuthTag);
+    }
+}