From bfbc352624fcca607a71e71ccc73b7c1ce86e4f1 Mon Sep 17 00:00:00 2001
From: Jari Zwarts <jarizw@gmail.com>
Date: Tue, 12 Jul 2016 21:59:15 +0200
Subject: [PATCH] Add ephemeral encryption & add CryptoSwift back
 (extensions...)

---
 .../xcdebugger/Expressions.xcexplist          | 11 +++++++
 Noti/Crypt.swift                              | 33 +++++++++++++++++--
 Noti/Ephemerals.swift                         | 15 ++++++++-
 Noti/PushManager.swift                        |  7 ++--
 Podfile                                       |  1 +
 5 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/Noti.xcworkspace/xcuserdata/Jari.xcuserdatad/xcdebugger/Expressions.xcexplist b/Noti.xcworkspace/xcuserdata/Jari.xcuserdatad/xcdebugger/Expressions.xcexplist
index e34e8d1..e7176ca 100644
--- a/Noti.xcworkspace/xcuserdata/Jari.xcuserdatad/xcdebugger/Expressions.xcexplist
+++ b/Noti.xcworkspace/xcuserdata/Jari.xcuserdatad/xcdebugger/Expressions.xcexplist
@@ -35,6 +35,17 @@
             </PersistentString>
          </PersistentStrings>
       </ContextState>
+      <ContextState
+         contextName = "Ephemerals.sendEphemeral([String : AnyObject]) -&gt; ():Ephemerals.swift">
+         <PersistentStrings>
+            <PersistentString
+               value = "json.rawString()">
+            </PersistentString>
+            <PersistentString
+               value = "json[&quot;push&quot;].rawString()">
+            </PersistentString>
+         </PersistentStrings>
+      </ContextState>
       <ContextState
          contextName = "PushManager.disconnect() -&gt; ():PushManager.swift">
       </ContextState>
diff --git a/Noti/Crypt.swift b/Noti/Crypt.swift
index 7ebe66e..ff27059 100644
--- a/Noti/Crypt.swift
+++ b/Noti/Crypt.swift
@@ -7,6 +7,7 @@
 //
 
 import Foundation
+import CryptoSwift
 
 public class Crypt {
     var key: [UInt8];
@@ -22,8 +23,7 @@ public class Crypt {
     
     func decryptMessage(cipher: String) -> String? {
         let rawData = NSData(base64EncodedString: cipher, options: NSDataBase64DecodingOptions(rawValue: 0))
-        var rawBytes = [UInt8](count: rawData!.length / sizeof(UInt8), repeatedValue: 0)
-        rawData?.getBytes(&rawBytes, length: rawBytes.count)
+        var rawBytes = rawData!.toArray()
         
         let tag = NSData(bytes: [UInt8](rawBytes[1...16]))
         let iv = [UInt8](rawBytes[17...28])
@@ -33,14 +33,41 @@ public class Crypt {
         if res == nil {
             return nil
         } else {
-            //is the resulting tag correct?
+            //verify the resulting tag...
             if tag == res!.1 {
                 return String(data: res!.0, encoding: NSUTF8StringEncoding)
             } else {
                 return nil
             }
         }
+    }
+    
+    func encryptMessage(message: String) -> String? {
+        let iv = CC.generateRandom(12)
+        let messageData = message.dataUsingEncoding(NSUTF8StringEncoding)!
+        let res = try? CC.GCM.crypt(CC.OpMode.encrypt, algorithm: .aes, data: messageData, key: NSData(bytes: key), iv: iv, aData: NSData(), tagLength: 16)
+        if res == nil {
+            return nil
+        }
         
+        let tag = res!.1
+        var data = [UInt8]()
+        data.append(49) // 1
+        data.appendContentsOf(tag.toArray())
+        data.appendContentsOf(iv.toArray())
+        data.appendContentsOf(res!.0.toArray())
+        
+        let out = NSData(bytes: data).base64EncodedStringWithOptions(NSDataBase64EncodingOptions.init(rawValue: 0))
+        
+        return out
     }
     
+}
+
+extension NSData {
+    func toArray() -> [UInt8] {
+        var bytes = [UInt8](count: self.length / sizeof(UInt8), repeatedValue: 0)
+        self.getBytes(&bytes, length: bytes.count)
+        return bytes
+    }
 }
\ No newline at end of file
diff --git a/Noti/Ephemerals.swift b/Noti/Ephemerals.swift
index e7e2b7f..593c516 100644
--- a/Noti/Ephemerals.swift
+++ b/Noti/Ephemerals.swift
@@ -15,17 +15,30 @@ import Alamofire
 import SwiftyJSON
 
 class Ephemerals: NSObject {
-    var token:String;
+    var token:String
+    var crypt:Crypt?
     
     init(token:String) {
         self.token = token
     }
     
     internal func sendEphemeral(body: [String: AnyObject]) {
+        var body = body
         let headers = [
             "Access-Token": token
         ];
         
+        if crypt != nil && body["type"] as? String == "push" {
+            print("Encrypting ephemeral...")
+            let json = JSON.init(body)
+            
+            let cipher = crypt!.encryptMessage(json["push"].rawString()!)
+            body["push"] = [
+                "encrypted": true,
+                "ciphertext": cipher!
+            ]
+        }
+        
         print("Sending ephemeral...")
         print("-------- BODY --------")
         debugPrint(body)
diff --git a/Noti/PushManager.swift b/Noti/PushManager.swift
index 5010026..8e903c5 100644
--- a/Noti/PushManager.swift
+++ b/Noti/PushManager.swift
@@ -12,7 +12,6 @@ import Foundation
 import Starscream
 import SwiftyJSON
 import Alamofire
-import CryptoSwift
 
 class PushManager: NSObject, WebSocketDelegate, NSUserNotificationCenterDelegate {
     var socket:WebSocket?
@@ -55,14 +54,14 @@ class PushManager: NSObject, WebSocketDelegate, NSUserNotificationCenterDelegate
     func initCrypt() {
         let keyData = userDefaults.objectForKey("secureKey") as? NSData
         if keyData != nil {
-            var key = [UInt8](count: keyData!.length / sizeof(UInt8), repeatedValue: 0)
-            keyData!.getBytes(&key, length: keyData!.length)
-            self.crypt = Crypt(key: key)
+            let key = keyData?.toArray()
+            self.crypt = Crypt(key: key!)
             print("Encryption enabled!")
         } else {
             self.crypt = nil
             print("Encryption not enabled")
         }
+        self.ephemerals.crypt = self.crypt
     }
     
     func getUserInfo(callback: (() -> Void)?) {
diff --git a/Podfile b/Podfile
index c09bbae..c7fc5e5 100644
--- a/Podfile
+++ b/Podfile
@@ -5,4 +5,5 @@ target 'Noti' do
     pod 'SwiftyJSON', '~> 2.3.2'
     pod 'Alamofire', '~> 3.4.1'
     pod 'Sparkle', '~> 1.14.0'
+    pod 'CryptoSwift', '~> 0.5.1'
 end
\ No newline at end of file