[omemo] Handle auth tag XEP compliant

3 files changed, 45 insertions, 15 deletions

diff --git a/omemo/omemo/aes_gcm_fallback.py b/omemo/omemo/aes_gcm_fallback.py
index a34255a..334eab2 100644
--- a/omemo/omemo/aes_gcm_fallback.py
+++ b/omemo/omemo/aes_gcm_fallback.py
@@ -29,11 +29,14 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import logging
 from struct import pack, unpack
 
 from Crypto.Cipher import AES
 from Crypto.Util import strxor
 
+log = logging.getLogger('gajim.plugin_system.omemo')
+
 
 def gcm_rightshift(vec):
     for x in range(15, 0, -1):
@@ -140,13 +143,20 @@ def gcm_encrypt(k, iv, plaintext, auth_data):
 
 def aes_encrypt(key, nonce, plaintext):
     """ Use AES128 GCM with the given key and iv to encrypt the payload. """
-    c, t = gcm_encrypt(key, nonce, plaintext, '')
-    result = c + t
-    return result
-
+    return gcm_encrypt(key, nonce, plaintext, '')
 
-def aes_decrypt(key, nonce, payload):
+def aes_decrypt(_key, nonce, payload):
     """ Use AES128 GCM with the given key and iv to decrypt the payload. """
-    ciphertext = payload[:-16]
-    mac = payload[-16:]
+    if len(_key) >= 32:
+        # XEP-0384
+        log.debug('XEP Compliant Key/Tag')
+        ciphertext = payload
+        key = _key[:16]
+        mac = _key[16:]
+    else:
+        # Legacy
+        log.debug('Legacy Key/Tag')
+        ciphertext = payload[:-16]
+        key = _key
+        mac = payload[-16:]
     return gcm_decrypt(key, nonce, ciphertext, '', mac)
diff --git a/omemo/omemo/aes_gcm_native.py b/omemo/omemo/aes_gcm_native.py
index 40bf127..28c7c43 100644
--- a/omemo/omemo/aes_gcm_native.py
+++ b/omemo/omemo/aes_gcm_native.py
@@ -19,6 +19,7 @@
 
 
 import os
+import logging
 from cryptography.hazmat.primitives.ciphers import Cipher
 from cryptography.hazmat.primitives.ciphers import algorithms
 from cryptography.hazmat.primitives.ciphers.modes import GCM
@@ -32,11 +33,22 @@ if os.name == 'nt':
 else:
     from cryptography.hazmat.backends import default_backend
 
+log = logging.getLogger('gajim.plugin_system.omemo')
 
-def aes_decrypt(key, iv, payload):
+def aes_decrypt(_key, iv, payload):
     """ Use AES128 GCM with the given key and iv to decrypt the payload. """
-    data = payload[:-16]
-    tag = payload[-16:]
+    if len(_key) >= 32:
+        # XEP-0384
+        log.debug('XEP Compliant Key/Tag')
+        data = payload
+        key = _key[:16]
+        tag = _key[16:]
+    else:
+        # Legacy
+        log.debug('Legacy Key/Tag')
+        data = payload[:-16]
+        key = _key
+        tag = payload[-16:]
     if os.name == 'nt':
         _backend = backend
     else:
@@ -58,4 +70,4 @@ def aes_encrypt(key, iv, plaintext):
         algorithms.AES(key),
         GCM(iv),
         backend=_backend).encryptor()
-    return encryptor.update(plaintext) + encryptor.finalize() + encryptor.tag
+    return encryptor.update(plaintext) + encryptor.finalize(), encryptor.tag
diff --git a/omemo/omemo/state.py b/omemo/omemo/state.py
index 6db7f1a..c0aeb93 100644
--- a/omemo/omemo/state.py
+++ b/omemo/omemo/state.py
@@ -226,6 +226,12 @@ class OmemoState:
             log.error('No known devices')
             return
 
+        payload, tag = encrypt(key, iv, plaintext)
+
+        # for XEP-384 Compliance uncomment
+        # key += tag
+        payload += tag
+
         # Encrypt the message key with for each of receivers devices
         for device in devices_list:
             try:
@@ -255,8 +261,6 @@ class OmemoState:
             except:
                 log.warning('Failed to find key for device ' + str(device))
 
-        payload = encrypt(key, iv, plaintext)
-
         result = {'sid': self.own_device_id,
                   'keys': encrypted_keys,
                   'jid': jid,
@@ -279,6 +283,12 @@ class OmemoState:
             log.error('No known devices')
             return
 
+        payload, tag = encrypt(key, iv, plaintext)
+
+        # for XEP-384 Compliance uncomment
+        # key += tag
+        payload += tag
+
         for tup in devices_list:
             self.get_session_cipher(tup[0], tup[1])
 
@@ -321,8 +331,6 @@ class OmemoState:
                 log.exception('ERROR:')
                 log.warning('Failed to find key for device ' + str(dev))
 
-        payload = encrypt(key, iv, plaintext)
-
         result = {'sid': self.own_device_id,
                   'keys': encrypted_keys,
                   'jid': jid,