import java.util.HashMap;
import java.util.Map;
import org.apache.gravitino.catalog.jdbc.config.JdbcConfig;
import org.apache.gravitino.catalog.jdbc.utils.DataSourceUtils;
import org.apache.gravitino.exceptions.GravitinoRuntimeException;

/**
 * Standalone verification of the CVE-2026-41042 H2-block fix in
 * DataSourceUtils.createDataSource(JdbcConfig). It is compiled and run against the
 * REAL patched DataSourceUtils / JdbcConfig source (on the classpath of the
 * official Gravitino 1.2.0 server libs), so it exercises the actual guard code.
 */
public class H2BlockTest {

  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> props(String url, String driver) {
    Map<String, String> p = new HashMap<>();
    p.put("jdbc-url", url);
    p.put("jdbc-driver", driver);
    p.put("jdbc-user", "test");
    p.put("jdbc-password", "test");
    p.put("jdbc-database", "test");
    return p;
  }

  /**
   * Returns the exact H2-guard message thrown, or null if the H2 guard did not fire (either the
   * datasource was created, or a different, non-H2 exception was thrown).
   */
  static String h2GuardMessage(String url, String driver) {
    try {
      DataSourceUtils.createDataSource(new JdbcConfig(props(url, driver)));
      return null; // created without H2 guard firing
    } catch (GravitinoRuntimeException e) {
      String m = e.getMessage();
      if (URL_MSG.equals(m) || DRV_MSG.equals(m)) {
        return m;
      }
      return null; // e.g. "Error creating datasource" -> not the H2 guard
    } catch (Exception e) {
      return null;
    }
  }

  static int failures = 0;

  static void expectBlocked(String label, String url, String driver, String expectedMsg) {
    String got = h2GuardMessage(url, driver);
    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, String url, String driver) {
    String got = h2GuardMessage(url, driver);
    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) {
    // --- CVE-2026-41042 original payload (relational JDBC testConnection path) ---
    String cveUrl =
        "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()";
    expectBlocked("cve-h2-url+org.h2.Driver", cveUrl, "org.h2.Driver", URL_MSG);

    // --- lakehouse-paimon variant payload (same H2 URL sink) ---
    String variantUrl =
        "jdbc:h2:mem:paimonrce;INIT=CREATE ALIAS EXEC AS "
            + "'String f() throws Exception { new ProcessBuilder(new String[]{"
            + "\"sh\",\"-c\",\"id > /tmp/m\"}).start().waitFor(); return \"ok\" }'\\;CALL EXEC()";
    expectBlocked("variant-h2-url+org.h2.Driver", variantUrl, "org.h2.Driver", URL_MSG);

    // --- URL-encoded prefix evasion (single + double encoding) ---
    expectBlocked("encoded-jdbc%3Ah2", "jdbc%3Ah2%3Amem%3Atest", "org.h2.Driver", URL_MSG);
    expectBlocked("double-encoded-jdbc%253Ah2", "jdbc%253Ah2%3Amem%3Atest", "org.h2.Driver",
        URL_MSG);

    // --- case-insensitivity ---
    expectBlocked("upper-JDBC:H2", "JDBC:H2:mem:test", "org.h2.Driver", URL_MSG);
    expectBlocked("mixed-Jdbc:H2", "Jdbc:H2:mem:test", "org.h2.Driver", URL_MSG);
    expectBlocked("driver-case-ORG.h2.Driver", "jdbc:postgresql://localhost/db", "ORG.h2.Driver",
        DRV_MSG);

    // --- H2 driver blocked even with a non-H2 URL ---
    expectBlocked("org.h2.Driver-with-mysql-url", "jdbc:mysql://localhost/db", "org.h2.sub.Driver",
        DRV_MSG);
    expectBlocked("org.h2.Driver-plain", "jdbc:mysql://localhost/db", "org.h2.Driver", DRV_MSG);

    // --- legitimate (non-H2) configs must NOT be blocked by the H2 guard ---
    expectAllowed("postgresql", "jdbc:postgresql://localhost/db", "org.postgresql.Driver");
    expectAllowed("mysql", "jdbc:mysql://localhost/db", "com.mysql.cj.jdbc.Driver");
    expectAllowed("mariadb", "jdbc:mariadb://localhost/db", "org.mariadb.jdbc.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);
  }
}