From f530472a26aeca741be50d589d4344c0a2eb0db5 Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Tue, 7 May 2019 13:59:34 +1000
Subject: [PATCH] Added FallBackSessionCipher.

---
 Pods                                          |  2 +-
 .../src/Loki/FallbackSessionCipher.swift      | 93 +++++++++++++++++++
 .../src/Messages/OWSMessageSender.m           | 66 ++++++++++++-
 SignalServiceKit/src/TSConstants.h            |  4 +
 4 files changed, 162 insertions(+), 3 deletions(-)
 create mode 100644 SignalServiceKit/src/Loki/FallbackSessionCipher.swift

diff --git a/Pods b/Pods
index 56186c3b3..f48c85741 160000
--- a/Pods
+++ b/Pods
@@ -1 +1 @@
-Subproject commit 56186c3b395dffd12fc22b675461720dd5a67059
+Subproject commit f48c857414ea68734bc786ffcaad77a35ac216a2
diff --git a/SignalServiceKit/src/Loki/FallbackSessionCipher.swift b/SignalServiceKit/src/Loki/FallbackSessionCipher.swift
new file mode 100644
index 000000000..473874063
--- /dev/null
+++ b/SignalServiceKit/src/Loki/FallbackSessionCipher.swift
@@ -0,0 +1,93 @@
+import Curve25519Kit
+import CryptoSwift
+
+private extension String {
+    // Convert hex string to Data
+    var hexData: Data {
+        var hex = self
+        var data = Data()
+        while(hex.count > 0) {
+            let subIndex = hex.index(hex.startIndex, offsetBy: 2)
+            let c = String(hex[..<subIndex])
+            hex = String(hex[subIndex...])
+            var ch: UInt32 = 0
+            Scanner(string: c).scanHexInt32(&ch)
+            var char = UInt8(ch)
+            data.append(&char, count: 1)
+        }
+        return data
+    }
+}
+
+/// A fallback session cipher which uses the the recipients pubkey to encrypt data
+@objc public final class FallBackSessionCipher: NSObject {
+    
+    // The length of the iv
+    private let ivLength: Int32 = 16;
+    
+    // The pubkey hex string of the recipient
+    private let recipientId: String
+    
+    // The pubkey representation of the hex id
+    private lazy var recipientPubKey: Data = {
+        var recipientId = self.recipientId
+        
+        // We need to check here if the id is prefix with '05'
+        // We only need to do this if the length is 66
+        if (recipientId.count == 66 && recipientId.hasPrefix("05")) {
+            recipientId = recipientId.substring(from: 2)
+        }
+
+        return recipientId.hexData
+    }()
+    
+    // The identity manager
+    private let identityKeyStore: OWSIdentityManager
+    
+    // Our identity key
+    private lazy var myIdentityKeyPair: ECKeyPair? = {
+        return identityKeyStore.identityKeyPair()
+    }()
+    
+    /// Creare a FallBackSessionCipher.
+    /// This is a very basic cipher and should only be used in special cases such as Friend Requests.
+    ///
+    /// - Parameters:
+    ///   - recipientId: The pubkey string of the recipient
+    ///   - identityKeyStore: The identity manager
+    @objc public init(recipientId: String, identityKeyStore: OWSIdentityManager) {
+        self.recipientId = recipientId
+        self.identityKeyStore = identityKeyStore
+        super.init()
+    }
+    
+    /// Encrypt a message
+    ///
+    /// - Parameter message: The message to encrypt
+    /// - Returns: The encypted message or nil if it failed
+    @objc public func encrypt(message: Data) -> Data? {
+        do {
+            let myIdentityKeyPair = self.myIdentityKeyPair!
+            let symmetricKey = try Curve25519.generateSharedSecret(fromPublicKey: recipientPubKey, privateKey: myIdentityKeyPair.privateKey)
+            return try diffieHellmanEncrypt(message: message, symmetricKey: symmetricKey)
+        } catch {
+            Logger.warn("FallBackSessionCipher: Failed to encrypt message")
+            return nil
+        }
+    }
+    
+    // Encypt the message with the symmetric key and a 16 bit iv
+    private func diffieHellmanEncrypt(message: Data, symmetricKey: Data) throws -> Data {
+        let iv = Randomness.generateRandomBytes(ivLength)!
+        let ivBytes = [UInt8](iv)
+        
+        let symmetricKeyBytes = [UInt8](symmetricKey)
+        let messageBytes = [UInt8](message)
+        
+        let blockMode = CBC(iv: ivBytes)
+        let aes = try AES(key: symmetricKeyBytes, blockMode: blockMode)
+        let cipherText = try aes.encrypt(messageBytes)
+        let ivAndCipher = ivBytes + cipherText
+        return Data(bytes: ivAndCipher, count: ivAndCipher.count)
+    }
+}
diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m
index d7159bec0..661736867 100644
--- a/SignalServiceKit/src/Messages/OWSMessageSender.m
+++ b/SignalServiceKit/src/Messages/OWSMessageSender.m
@@ -1059,8 +1059,19 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
         if (messageSend.isUDSend) {
             hasValidMessageType = [messageType isEqualToNumber:@(TSUnidentifiedSenderMessageType)];
         } else {
+            NSArray *validMessageTypes = @[
+                                           @(TSEncryptedWhisperMessageType),
+                                           @(TSPreKeyWhisperMessageType),
+                                           @(TSFriendRequestMessageType) // Loki friend request
+                                           ];
+            hasValidMessageType = [validMessageTypes containsObject:messageType];
+
+            /* Loki: Original code:
+             * ================
             hasValidMessageType = ([messageType isEqualToNumber:@(TSEncryptedWhisperMessageType)] ||
-                [messageType isEqualToNumber:@(TSPreKeyWhisperMessageType)]);
+                                   [messageType isEqualToNumber:@(TSPreKeyWhisperMessageType)]);
+             */
+
         }
 
         if (!hasValidMessageType) {
@@ -1509,7 +1520,11 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
         @try {
             // This may involve blocking network requests, so we do it _before_
             // we open a transaction.
-            [self throws_ensureRecipientHasSessionForMessageSend:messageSend deviceId:deviceId];
+            // TODO: Replace this when we add in friend request stuff
+            Boolean isFriendRequest = true;
+            if (!isFriendRequest) {
+                [self throws_ensureRecipientHasSessionForMessageSend:messageSend deviceId:deviceId];
+            }
 
             __block NSDictionary *_Nullable messageDict;
             __block NSException *encryptionException;
@@ -1684,6 +1699,47 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
             }) retainUntilComplete];
 }
 
+- (nullable NSDictionary *)throws_encryptedFriendMessageForMessageSend:(OWSMessageSend *)messageSend
+                                                        deviceId:(NSNumber *)deviceId
+                                                       plainText:(NSData *)plainText
+{
+    OWSAssertDebug(messageSend);
+    OWSAssertDebug(deviceId);
+    OWSAssertDebug(plainText);
+    
+    SignalRecipient *recipient = messageSend.recipient;
+    NSString *recipientId = recipient.recipientId;
+    
+    FallBackSessionCipher *cipher = [[FallBackSessionCipher alloc] initWithRecipientId:recipientId
+                                                                    identityKeyStore:self.identityManager];
+    
+    // This will return nil if encryption failed
+    NSData *_Nullable serializedMessage = [cipher encryptWithMessage:[plainText paddedMessageBody]];
+    if (!serializedMessage) {
+        OWSFailDebug(@"Failed to encrypt friend message to: %@", recipientId);
+        return nil;
+    }
+    
+    OWSMessageServiceParams *messageParams =
+    [[OWSMessageServiceParams alloc] initWithType:TSFriendRequestMessageType
+                                      recipientId:recipientId
+                                           device:[deviceId intValue]
+                                          content:serializedMessage
+                                         isSilent:false
+                                         isOnline:false
+                                   registrationId:0];
+    
+    NSError *error;
+    NSDictionary *jsonDict = [MTLJSONAdapter JSONDictionaryFromModel:messageParams error:&error];
+    
+    if (error) {
+        OWSProdError([OWSAnalyticsEvents messageSendErrorCouldNotSerializeMessageJson]);
+        return nil;
+    }
+    
+    return jsonDict;
+}
+
 - (nullable NSDictionary *)throws_encryptedMessageForMessageSend:(OWSMessageSend *)messageSend
                                                         deviceId:(NSNumber *)deviceId
                                                        plainText:(NSData *)plainText
@@ -1699,6 +1755,12 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
     SignalRecipient *recipient = messageSend.recipient;
     NSString *recipientId = recipient.recipientId;
     OWSAssertDebug(recipientId.length > 0);
+    
+    // TODO: Change this when we have friend request support
+    Boolean isFriendRequest = true;
+    if (isFriendRequest) {
+        return [self throws_encryptedFriendMessageForMessageSend:messageSend deviceId:deviceId plainText:plainText];
+    }
 
     // This may throw an exception.
     if (![storage containsSession:recipientId deviceId:[deviceId intValue] protocolContext:transaction]) {
diff --git a/SignalServiceKit/src/TSConstants.h b/SignalServiceKit/src/TSConstants.h
index bce163b0b..8cbdb84b0 100644
--- a/SignalServiceKit/src/TSConstants.h
+++ b/SignalServiceKit/src/TSConstants.h
@@ -14,6 +14,10 @@ typedef NS_ENUM(NSInteger, TSWhisperMessageType) {
     TSPreKeyWhisperMessageType = 3,
     TSUnencryptedWhisperMessageType = 4,
     TSUnidentifiedSenderMessageType = 6,
+    
+    // Loki: contains prekeys + message and is using simple encryption
+    // Must match the protobuf type
+    TSFriendRequestMessageType = 101,
 };
 
 #pragma mark Server Address