diff --git a/Sources/WKZombie/RenderOperation.swift b/Sources/WKZombie/RenderOperation.swift index 81a8682..65e3f65 100644 --- 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 index 384f8c7..8f958d5 100644 --- 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 index fecb1b6..d2316d7 100644 --- 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() }