added options to inject userScripts and messageHandlers

jakejcheng · May 28, 2020 · a776d60 · a776d60
1 parent fcbba84
commit a776d60
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/Sources/WKZombie/RenderOperation.swift b/Sources/WKZombie/RenderOperation.swift
@@ -43,6 +43,8 @@ internal class RenderOperation : Operation {
     var authenticationBlock : AuthenticationHandler?
     var postAction: PostAction = .none
 
+    var messageHandlers: [String:(_ message: WKScriptMessage)->Void] = [:]
+
     internal fileprivate(set) var result : Data?
     internal fileprivate(set) var response : URLResponse?
     internal fileprivate(set) var error : Error?
@@ -160,6 +162,10 @@ internal class RenderOperation : Operation {
 
     fileprivate func setupReferences() {
         dispatch_sync_on_main_thread {
+            for name in messageHandlers.keys {
+                webView?.configuration.userContentController.removeScriptMessageHandler(forName: name)
+                webView?.configuration.userContentController.add(self, name: name)
+            }
             webView?.configuration.userContentController.add(self, name: "doneLoading")
             webView?.navigationDelegate = self
         }
@@ -191,6 +197,12 @@ extension RenderOperation : WKScriptMessageHandler {
                 webView.stopLoading()
                 self.webView(webView, didFinish: nil)
             }
+
+            for (name, messageHandler) in self.messageHandlers {
+                if message.name == name {
+                    messageHandler(message)
+                }
+            }
         }
     }
 }

diff --git a/Sources/WKZombie/Renderer.swift b/Sources/WKZombie/Renderer.swift
@@ -57,9 +57,9 @@ internal class Renderer {
     }()
 
     fileprivate var webView : WKWebView!
+    fileprivate var messageHandlers: [String:(_ message: WKScriptMessage)->Void]
 
-
-    init(processPool: WKProcessPool? = nil) {
+    init(processPool: WKProcessPool? = nil, userScripts: [WKUserScript] = [], messageHandlers: [String:(_ message: WKScriptMessage)->Void] = [:]) {
         let doneLoadingWithoutMediaContentScript = "window.webkit.messageHandlers.doneLoading.postMessage(\(Renderer.scrapingCommand));"
         let doneLoadingUserScript = WKUserScript(source: doneLoadingWithoutMediaContentScript, injectionTime: .atDocumentEnd, forMainFrameOnly: true)
 
@@ -71,7 +71,13 @@ internal class Renderer {
         let contentController = WKUserContentController()
         contentController.addUserScript(doneLoadingUserScript)
         contentController.addUserScript(getElementUserScript)
-
+
+        for userScript in userScripts {
+            contentController.addUserScript(userScript)
+        }
+
+        self.messageHandlers = messageHandlers
+
         let config = WKWebViewConfiguration()
         config.processPool = processPool ?? WKProcessPool()
         config.userContentController = contentController
@@ -181,6 +187,7 @@ internal class Renderer {
         }
         operation.requestBlock = requestBlock
         operation.authenticationBlock = authenticationHandler
+        operation.messageHandlers = messageHandlers
         return operation
     }
 

diff --git a/Sources/WKZombie/WKZombie.swift b/Sources/WKZombie/WKZombie.swift
@@ -112,9 +112,9 @@ public class WKZombie {
 
      - returns: A WKZombie instance.
      */
-    public init(name: String? = "WKZombie", processPool: WKProcessPool? = nil) {
+    public init(name: String? = "WKZombie", processPool: WKProcessPool? = nil, userScripts: [WKUserScript] = [], messageHandlers: [String:(_ message: WKScriptMessage)->Void] = [:]) {
         self.name = name
-        self._renderer = Renderer(processPool: processPool)
+        self._renderer = Renderer(processPool: processPool, userScripts: userScripts, messageHandlers: messageHandlers)
         self._fetcher = ContentFetcher()
     }