From 31ee61a5cde21d1ad92fb3012c21f84d9f27acf0 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Thu, 4 Aug 2022 16:42:12 +0200
Subject: [PATCH] feat: Add convenience functions

---
 lib/src/omemo/sessionmanager.dart | 40 +++++++++++++++++++++++++++++--
 test/omemo_test.dart              | 28 +++++-----------------
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/lib/src/omemo/sessionmanager.dart b/lib/src/omemo/sessionmanager.dart
index 93f2913..08ce7d3 100644
--- a/lib/src/omemo/sessionmanager.dart
+++ b/lib/src/omemo/sessionmanager.dart
@@ -7,6 +7,7 @@ import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/helpers.dart';
 import 'package:omemo_dart/src/omemo/bundle.dart';
 import 'package:omemo_dart/src/omemo/device.dart';
+import 'package:omemo_dart/src/x3dh/x3dh.dart';
 import 'package:synchronized/synchronized.dart';
 
 /// The info used for when encrypting the AES key for the actual payload.
@@ -75,8 +76,43 @@ class OmemoSessionManager {
     });
   }
 
-  Future<void> addSessionFromBundle(String jid, String deviceId, OmemoBundle bundle) async {
-    // TODO(PapaTutuWawa): Do
+  /// Create a ratchet session initiated by Alice to the user with Jid [jid] and the device
+  /// [deviceId] from the bundle [bundle].
+  Future<X3DHAliceResult> addSessionFromBundle(String jid, String deviceId, OmemoBundle bundle) async {
+    final kexResult = await x3dhFromBundle(
+      bundle,
+      device.ik,
+    );
+    final ratchet = await OmemoDoubleRatchet.initiateNewSession(
+      bundle.spk,
+      kexResult.sk,
+      kexResult.ad,
+    );
+
+    await addSession(jid, deviceId, ratchet);
+
+    return kexResult;
+  }
+
+  /// Build a new session with the user at [jid] with the device [deviceId] using data
+  /// from the key exchange [kex].
+  // TODO(PapaTutuWawa): Use OMEMOKeyExchange
+  // TODO(PapaTutuWawa): Replace the OPK
+  Future<void> addSessionFromKeyExchange(String jid, String deviceId, X3DHMessage kex) async {
+    final kexResult = await x3dhFromInitialMessage(
+      kex,
+      device.spk,
+      // TODO(PapaTutuWawa): Fix
+      device.opks.values.elementAt(0),
+      device.ik,
+    );
+    final ratchet = await OmemoDoubleRatchet.acceptNewSession(
+      device.spk,
+      kexResult.sk,
+      kexResult.ad,
+    );
+
+    await addSession(jid, deviceId, ratchet);
   }
   
   /// Encrypt the key [plaintext] for all known bundles of [jid]. Returns a map that
diff --git a/test/omemo_test.dart b/test/omemo_test.dart
index f60a7cf..2970783 100644
--- a/test/omemo_test.dart
+++ b/test/omemo_test.dart
@@ -12,37 +12,21 @@ void main() {
     final bobSession = await OmemoSessionManager.generateNewIdentity(opkAmount: 1);
 
     // Perform the X3DH
-    final x3dhAliceResult = await x3dhFromBundle(
+    final x3dhAliceResult = await aliceSession.addSessionFromBundle(
+      bobJid,
+      bobSession.device.id,
       await bobSession.device.toBundle(),
-      aliceSession.device.ik,
     );
-    final x3dhBobResult = await x3dhFromInitialMessage(
+    await bobSession.addSessionFromKeyExchange(
+      aliceJid,
+      aliceSession.device.id,
       X3DHMessage(
         aliceSession.device.ik.pk,
         x3dhAliceResult.ek.pk,
         '2',
       ),
-      bobSession.device.spk,
-      bobSession.device.opks.values.elementAt(0),
-      bobSession.device.ik,
     );
 
-    // Build the ratchets
-    final aliceRatchet = await OmemoDoubleRatchet.initiateNewSession(
-      bobSession.device.spk.pk,
-      x3dhAliceResult.sk,
-      x3dhAliceResult.ad,
-    );
-    final bobRatchet = await OmemoDoubleRatchet.acceptNewSession(
-      bobSession.device.spk,
-      x3dhBobResult.sk,
-      x3dhBobResult.ad,
-    );
-
-    // Add the ratchets to the session managers
-    await aliceSession.addSession(bobJid, bobSession.device.id, aliceRatchet);
-    await bobSession.addSession(aliceJid, aliceSession.device.id, bobRatchet);
-
     // Alice encrypts a message for Bob
     const messagePlaintext = 'Hello Bob!';
     final aliceMessage = await aliceSession.encryptToJid(bobJid, messagePlaintext);