import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import org.apache.gravitino.catalog.lakehouse.paimon.PaimonConfig;
import org.apache.gravitino.catalog.lakehouse.paimon.utils.CatalogUtils;
import org.apache.gravitino.exceptions.GravitinoRuntimeException;

/**
 * Standalone verification of the CVE-2026-41042 H2-block fix in the
 * lakehouse-paimon CatalogUtils.checkPaimonConfig() JDBC branch. This path opens
 * JDBC connections through Paimon's own JdbcCatalog (DriverManager), bypassing
 * DataSourceUtils, so it needs its own H2 guard. We invoke the real (private)
 * checkPaimonConfig via reflection so the actual patched code is exercised
 * without spinning up Paimon's CatalogFactory.
 */
public class PaimonH2BlockTest {

  static final String URL_MSG = "H2 JDBC URL is not allowed in catalog configuration";
  static final String DRV_MSG = "H2 JDBC driver is not allowed in catalog configuration";

  static Map<String, String> jdbcProps(String uri, String driver) {
    Map<String, String> p = new HashMap<>();
    p.put("metastore", "jdbc");
    p.put("warehouse", "/tmp/paimon_wh");
    p.put("uri", uri);
    p.put("jdbc-driver", driver);
    p.put("jdbc-user", "test");
    p.put("jdbc-password", "test");
    return p;
  }

  static Map<String, String> filesystemProps() {
    Map<String, String> p = new HashMap<>();
    p.put("metastore", "filesystem");
    p.put("warehouse", "/tmp/paimon_wh");
    return p;
  }

  /** Returns the H2-guard message thrown by checkPaimonConfig, or null if it did not fire. */
  static String h2GuardMessage(Map<String, String> props) throws Exception {
    PaimonConfig cfg = new PaimonConfig(props);
    Method m = CatalogUtils.class.getDeclaredMethod("checkPaimonConfig", PaimonConfig.class);
    m.setAccessible(true);
    try {
      m.invoke(null, cfg);
      return null; // returned normally -> H2 guard did not fire
    } catch (InvocationTargetException ite) {
      Throwable t = ite.getCause();
      if (t instanceof GravitinoRuntimeException) {
        String msg = t.getMessage();
        if (URL_MSG.equals(msg) || DRV_MSG.equals(msg)) {
          return msg;
        }
      }
      return null; // different exception (e.g. driver not loadable) -> not the H2 guard
    }
  }

  static int failures = 0;

  static void expectBlocked(String label, Map<String, String> props, String expectedMsg)
      throws Exception {
    String got = h2GuardMessage(props);
    if (expectedMsg.equals(got)) {
      System.out.println("PASS [blocked]      " + label + " -> " + got);
    } else {
      failures++;
      System.out.println("FAIL [should block] " + label + " -> expected '" + expectedMsg
          + "' but got '" + got + "'");
    }
  }

  static void expectAllowed(String label, Map<String, String> props) throws Exception {
    String got = h2GuardMessage(props);
    if (got == null) {
      System.out.println("PASS [allowed]      " + label);
    } else {
      failures++;
      System.out.println("FAIL [should allow] " + label + " -> unexpectedly blocked: '" + got + "'");
    }
  }

  public static void main(String[] args) throws Exception {
    String cveH2 =
        "jdbc:h2:mem:rce;INIT=CREATE ALIAS EXEC AS "
            + "'String f() throws Exception { new ProcessBuilder(new String[]{"
            + "\"sh\",\"-c\",\"id\"}).start().waitFor(); return \"ok\" }'\\;CALL EXEC()";

    // --- Paimon JDBC metastore backend with a malicious H2 URI (variant bypass payload) ---
    expectBlocked("paimon-jdbc-h2-uri+org.h2.Driver", jdbcProps(cveH2, "org.h2.Driver"), URL_MSG);
    // --- URL-encoded H2 prefix evasion ---
    expectBlocked("paimon-jdbc-encoded-h2-uri", jdbcProps("jdbc%3Ah2%3Amem%3Atest", "org.h2.Driver"),
        URL_MSG);
    expectBlocked("paimon-jdbc-double-encoded-h2-uri",
        jdbcProps("jdbc%253Ah2%3Amem%3Atest", "org.h2.Driver"), URL_MSG);
    // --- case-insensitive URL ---
    expectBlocked("paimon-jdbc-upper-JDBC:H2", jdbcProps("JDBC:H2:mem:test", "org.h2.Driver"),
        URL_MSG);
    // --- H2 driver blocked even with a non-H2 URI ---
    expectBlocked("paimon-jdbc-org.h2.Driver-with-pg-uri",
        jdbcProps("jdbc:postgresql://localhost/db", "org.h2.Driver"), DRV_MSG);
    expectBlocked("paimon-jdbc-ORG.h2.Driver-case",
        jdbcProps("jdbc:postgresql://localhost/db", "ORG.h2.Driver"), DRV_MSG);

    // --- legitimate configs must NOT be blocked by the H2 guard ---
    expectAllowed("paimon-filesystem-backend", filesystemProps());
    expectAllowed("paimon-jdbc-postgresql-driver",
        jdbcProps("jdbc:postgresql://localhost/db", "org.postgresql.Driver"));

    System.out.println();
    if (failures == 0) {
      System.out.println("RESULT: ALL ASSERTIONS PASSED");
      System.exit(0);
    }
    System.out.println("RESULT: " + failures + " ASSERTION(S) FAILED");
    System.exit(1);
  }
}