From 07f594a7fde4bebb86d75af40d7be7ebb201dbd9 Mon Sep 17 00:00:00 2001
From: Kyle Spearrin <kyle.spearrin@gmail.com>
Date: Fri, 16 Aug 2019 08:43:38 -0400
Subject: [PATCH] handle messages coming from scripts to app

---
 .../safari/SafariExtensionHandler.swift       | 14 ++++++++---
 .../SafariExtensionViewController.swift       | 25 ++++++++++++++++++-
 src/safari/app/desktop/safari/script.js       |  2 +-
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/src/safari/app/desktop/safari/SafariExtensionHandler.swift b/src/safari/app/desktop/safari/SafariExtensionHandler.swift
index efb8e099d8..8655220ba2 100644
--- a/src/safari/app/desktop/safari/SafariExtensionHandler.swift
+++ b/src/safari/app/desktop/safari/SafariExtensionHandler.swift
@@ -12,11 +12,19 @@ class SafariExtensionHandler: SFSafariExtensionHandler {
     
     override func messageReceived(withName messageName: String, from page: SFSafariPage, userInfo: [String : Any]?) {
         // This method will be called when a content script provided by your extension calls safari.extension.dispatchMessage("message").
-        page.getPropertiesWithCompletionHandler { properties in
-            NSLog("The extension received a message (\(messageName)) from a script injected into (\(String(describing: properties?.url))) with userInfo (\(userInfo ?? [:]))")
+        
+        if(messageName == "bitwarden") {
+            page.getPropertiesWithCompletionHandler { properties in
+                NSLog("The extension received a message (\(messageName)) from a script injected into (\(String(describing: properties?.url))) with userInfo (\(userInfo ?? [:]))")
+                
+                DispatchQueue.main.async {
+                    SafariExtensionViewController.shared.replyMessageFromScript(msg: userInfo)
+                }
+                
+            }
         }
         
-        page.dispatchMessageToScript(withName: "getInfo", userInfo: ["hello": "world", "foo": "bar"])
+        // page.dispatchMessageToScript(withName: "getInfo", userInfo: ["hello": "world", "foo": "bar"])
     }
     
     override func toolbarItemClicked(in window: SFSafariWindow) {
diff --git a/src/safari/app/desktop/safari/SafariExtensionViewController.swift b/src/safari/app/desktop/safari/SafariExtensionViewController.swift
index ff342b17dc..03047e6b17 100644
--- a/src/safari/app/desktop/safari/SafariExtensionViewController.swift
+++ b/src/safari/app/desktop/safari/SafariExtensionViewController.swift
@@ -90,10 +90,27 @@ class SafariExtensionViewController: SFSafariExtensionViewController, WKScriptMe
     }
     
     func replyMessage(message: AppMessage) {
+        if(webView == nil) {
+            return;
+        }
         let json = (jsonSerialize(obj: message) ?? "null")
         webView.evaluateJavaScript("window.bitwardenSafariAppMessageReceiver(\(json));", completionHandler: nil)
     }
-
+    
+    func replyMessageFromScript(msg: [String : Any]?) {
+        if(webView == nil) {
+            return;
+        }
+        let newMsg = AppMessage()
+        newMsg.command = "cs_message"
+        do {
+            let jsonData = try JSONSerialization.data(withJSONObject: msg as Any, options: [])
+            newMsg.data = String(data: jsonData, encoding: .utf8)
+        } catch let error {
+            print("error converting to json: \(error)")
+        }
+        replyMessage(message: newMsg)
+    }
 }
 
 func jsonSerialize<T: Encodable>(obj: T?) -> String? {
@@ -120,6 +137,12 @@ func jsonDeserialize<T: Decodable>(json: String?) -> T? {
 }
 
 class AppMessage : Decodable, Encodable {
+    init() {
+        id = ""
+        command = ""
+        data = nil
+        responseData = nil
+    }
     var id: String
     var command: String
     var data: String?
diff --git a/src/safari/app/desktop/safari/script.js b/src/safari/app/desktop/safari/script.js
index 0bd8816292..ff7ef97b6d 100644
--- a/src/safari/app/desktop/safari/script.js
+++ b/src/safari/app/desktop/safari/script.js
@@ -1,5 +1,5 @@
 document.addEventListener("DOMContentLoaded", (event) => {
-                          safari.extension.dispatchMessage("Hello World!", { key: "value2" });
+                          safari.extension.dispatchMessage("bitwarden", { hello: "world!!!" });
                           
                           safari.self.addEventListener("message", (e) => {
                                                        console.log(e);