From 897c5c4d43c2cbbda8ba72cde15ab022f3281788 Mon Sep 17 00:00:00 2001
From: thc202 <thc202@gmail.com>
Date: Fri, 28 Jun 2024 11:35:03 +0100
Subject: [PATCH] Allow to invoke scripts with custom engines

Allow the `ZestRuntime` to provide a `ScriptEngine` based on a file
extension and change `ZestActionInvoke` to use it, allowing for
runtimes to provide their own script engines.

Signed-off-by: thc202 <thc202@gmail.com>
---
 .../zest/core/v1/ZestActionInvoke.java        | 20 ++++++++++++++-----
 .../org/zaproxy/zest/core/v1/ZestRuntime.java | 14 +++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/zaproxy/zest/core/v1/ZestActionInvoke.java b/src/main/java/org/zaproxy/zest/core/v1/ZestActionInvoke.java
index 9c962da..865ed4f 100644
--- a/src/main/java/org/zaproxy/zest/core/v1/ZestActionInvoke.java
+++ b/src/main/java/org/zaproxy/zest/core/v1/ZestActionInvoke.java
@@ -108,11 +108,7 @@ public String invoke(ZestResponse response, ZestRuntime runtime)
                 // Its a Zest script
                 engine = runtime.getScriptEngineFactory().getScriptEngine();
             } else {
-                ScriptEngineManager manager = new ScriptEngineManager();
-                engine = manager.getEngineByExtension(ext);
-                if (engine == null) {
-                    engine = manager.getEngineByName(ext);
-                }
+                engine = getScriptEngine(runtime, ext);
             }
         }
 
@@ -197,6 +193,20 @@ public String invoke(ZestResponse response, ZestRuntime runtime)
         }
     }
 
+    private static ScriptEngine getScriptEngine(ZestRuntime runtime, String ext) {
+        ScriptEngine engine = runtime.getScriptEngine(ext);
+        if (engine != null) {
+            return engine;
+        }
+
+        ScriptEngineManager manager = new ScriptEngineManager();
+        engine = manager.getEngineByExtension(ext);
+        if (engine != null) {
+            return engine;
+        }
+        return manager.getEngineByName(ext);
+    }
+
     private Charset getCharsetImpl() {
         if (charset == null || charset.isEmpty()) {
             return Charset.defaultCharset();
diff --git a/src/main/java/org/zaproxy/zest/core/v1/ZestRuntime.java b/src/main/java/org/zaproxy/zest/core/v1/ZestRuntime.java
index 28bd9f4..6032c1e 100644
--- a/src/main/java/org/zaproxy/zest/core/v1/ZestRuntime.java
+++ b/src/main/java/org/zaproxy/zest/core/v1/ZestRuntime.java
@@ -4,6 +4,7 @@
 package org.zaproxy.zest.core.v1;
 
 import java.util.List;
+import javax.script.ScriptEngine;
 import javax.script.ScriptEngineFactory;
 import org.openqa.selenium.WebDriver;
 
@@ -55,6 +56,19 @@ default void setGlobalVariable(String name, String value) {
         // Nothing to do.
     }
 
+    /**
+     * Gets a {@code ScriptEngine} for the given extension.
+     *
+     * <p>If no engine is returned it is used a default engine, if available.
+     *
+     * @param extension the extension of the script.
+     * @return the {@code ScriptEngine}, or {@code null}.
+     * @since 0.22.0
+     */
+    default ScriptEngine getScriptEngine(String extension) {
+        return null;
+    }
+
     /**
      * Get the last response.
      *