Setup.java

package servlets;

import com.mongodb.MongoClient;
import dbProcs.Constants;
import dbProcs.Database;
import dbProcs.MongoDatabase;
import dbProcs.Setter;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServlet;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.LinkOption;
import java.nio.file.NoSuchFileException;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Locale;
import java.util.Properties;
import java.util.ResourceBundle;
import java.util.UUID;
import org.apache.commons.io.FileUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import servlets.module.challenge.XxeChallenge1;
import servlets.module.lesson.XxeLesson;
import utils.Validate;

public class Setup extends HttpServlet {

  private static final Logger log = LogManager.getLogger(Setup.class);
  private static final long serialVersionUID = -892181347446991016L;

  private static volatile Boolean installedCached = null;

  public void doPost(HttpServletRequest request, HttpServletResponse response)
      throws ServletException, IOException {
    // Translation Stuff
    Locale locale = new Locale(Validate.validateLanguage(request.getSession()));

    ResourceBundle.getBundle("i18n.servlets.errors", locale);
    ResourceBundle bundle = ResourceBundle.getBundle("i18n.text", locale);
    response.setContentType("text/html; charset=UTF-8");
    request.setCharacterEncoding("UTF-8");

    // Output Stuff
    PrintWriter out = response.getWriter();
    String htmlOutput = "";
    boolean success = false;
    boolean validateInput = true;
    boolean saveMysqlProperties = false;
    boolean hasDBFile = false;

    // Parameters From Form
    String dbHost = request.getParameter("dbhost");
    String dbPort = request.getParameter("dbport");
    String dbUser = request.getParameter("dbuser");
    String dbPass = request.getParameter("dbpass");

    String dbOptions;
    String connectionURL;
    String driverType;

    String dbOverride = request.getParameter("dboverride");

    Properties mysql_props = Setup.getDBProps();
    Properties mongo_props = new Properties();

    hasDBFile = (mysql_props != null);

    if (hasDBFile) {
      // Db auth file exists, try to load from it

      if (dbHost.isEmpty() != dbPort.isEmpty()) {
        // Only one of db host and db port provided, we need both or neither

        htmlOutput += "If you override db host and db port, both must be entered!";
        validateInput = false;
        connectionURL = "";
      } else if (dbHost.isEmpty() && dbPort.isEmpty()) {
        // Both db host and db port are missing, load from props file instead
        connectionURL = mysql_props.getProperty("databaseConnectionURL");
        String databaseSchema = mysql_props.getProperty("databaseSchema");

        if (connectionURL == null || databaseSchema == null) {
          // Nothing useful given in user input or from properties file, bail out.
          validateInput = false;
        }
      } else {
        // Override db properties from request parameters
        connectionURL = "jdbc:mariadb://" + dbHost + ":" + dbPort + "/";

        // Store the overridden data in properties file
        saveMysqlProperties = true;
      }

      dbOptions = mysql_props.getProperty("databaseOptions");
      if (dbOptions == null) {
        dbOptions = "useUnicode=true&character_set_server=utf8mb4";
      }
      driverType = mysql_props.getProperty("DriverType");
      if (driverType == null) {
        driverType = "org.mariadb.jdbc.Driver";
      }
      if (dbUser.isEmpty()) {
        dbUser = mysql_props.getProperty("databaseUsername");
        if (dbUser == null) {
          validateInput = false;
        }
      }
      if (dbPass.isEmpty()) {
        dbPass = mysql_props.getProperty("databasePassword");
        if (dbPass == null) {
          validateInput = false;
        }
      }
    } else {
      connectionURL = "jdbc:mariadb://" + dbHost + ":" + dbPort + "/";
      driverType = "org.mariadb.jdbc.Driver";
      dbOptions = "useUnicode=true&character_set_server=utf8mb4";
      validateInput = true;
      saveMysqlProperties = true;
    }

    if (!validateInput) {
      htmlOutput += "Data validation failed.";
      success = false;
    } else {

      String dbAuth = request.getParameter("dbauth");

      String mongodbHost = request.getParameter("mhost");
      String mongodbPort = request.getParameter("mport");
      String nosqlprops =
          new File(Database.class.getResource("/challenges/NoSqlInjection1.properties").getFile())
              .getAbsolutePath();

      try (InputStream mongo_input = new FileInputStream(nosqlprops)) {

        mongo_props.load(mongo_input);
      }

      String mongodbName = mongo_props.getProperty("databaseName");
      if (mongodbName == null) {
        String message = "Could not find databaseName in nosql properties file";
        log.fatal(message);
        throw new RuntimeException(message);
      }

      log.debug("Starting database setup...");

      String auth = "";

      String enableMongoChallenge = request.getParameter("enableMongoChallenge");

      String enableUnsafeLevels = request.getParameter("unsafeLevels");

      // Mongo DB properties
      StringBuffer mongoProp = new StringBuffer();
      mongoProp.append("connectionHost=" + mongodbHost);
      mongoProp.append("\n");
      mongoProp.append("connectionPort=" + mongodbPort);
      mongoProp.append("\n");
      mongoProp.append("databaseName=" + mongodbName);
      mongoProp.append("\n");
      mongoProp.append("connectTimeout=10000");
      mongoProp.append("\n");
      mongoProp.append("socketTimeout=0");
      mongoProp.append("\n");
      mongoProp.append("serverSelectionTimeout=30000");
      mongoProp.append("\n");

      try {
        auth = new String(Files.readAllBytes(Paths.get(Constants.SETUP_AUTH)));
      } catch (NoSuchFileException e) {
        // Auth file could not be found.
        htmlOutput += "Auth file could not be found";
        log.error("Auth file could not be found: " + e.toString());
      }

      if (auth == "") {
        // No auth loaded, could be because user never reloaded setup page after an
        // error. Generate it again
        log.debug("Generating auth file");

        generateAuth();
      }

      if (!auth.equals(dbAuth)) {
        log.debug("Invalid auth supplied");

        // The supplied auth data was incorrect
        htmlOutput += bundle.getString("generic.text.setup.authentication.failed");
        log.error("Authorization mismatch: " + auth + " does not equal " + dbAuth);

      } else {
        // Test the user's entered database properties. Use DriverManager directly instead of
        // the pool: setup is a one-shot credential check that runs before database.properties
        // exists, and routing it through a pooled DataSource keyed on (url, user) would
        // silently reuse a stale pool when the password changes between setup attempts.
        //
        // Append connectTimeout=5000 so a typo'd host or unreachable port fails fast (matching
        // ConnectionPool's 5s default) instead of hanging on the OS's default TCP connect
        // timeout — DriverManager has no built-in timeout and the setup page would otherwise
        // appear to freeze.
        Boolean connectionSuccess = false;
        log.debug("Attempting to connect to database");

        String testJdbcUrl;
        if (dbOptions != null && !dbOptions.isEmpty()) {
          testJdbcUrl = connectionURL + "?" + dbOptions + "&connectTimeout=5000";
        } else {
          testJdbcUrl = connectionURL + "?connectTimeout=5000";
        }

        try (Connection conn = DriverManager.getConnection(testJdbcUrl, dbUser, dbPass)) {
          connectionSuccess = true;
          log.debug("Database connection successful");

        } catch (SQLException e) {
          htmlOutput += bundle.getString("generic.text.setup.connection.failed") + e.getMessage();

          log.error("DB connection error: " + e.toString());
          connectionSuccess = false;
        }

        if (!connectionSuccess) {
          htmlOutput += bundle.getString("generic.text.setup.connection.failed");
        } else {
          // Write the user's entered mysql database properties to file

          if (saveMysqlProperties) {

            try (OutputStream mysql_output = new FileOutputStream(Constants.MYSQL_DB_PROP)) {

              mysql_props = new Properties();

              mysql_props.setProperty("databaseConnectionURL", connectionURL);
              mysql_props.setProperty("DriverType", driverType);
              mysql_props.setProperty("databaseOptions", dbOptions);
              mysql_props.setProperty("databaseSchema", "core");
              mysql_props.setProperty("databaseUsername", dbUser);
              mysql_props.setProperty("databasePassword", dbPass);

              // save properties to project root folder
              mysql_props.store(mysql_output, null);
              success = true;

            } catch (IOException e) {

              success = false;

              htmlOutput = bundle.getString("generic.text.setup.failed") + ": " + e.getMessage();

              log.error("Could not save mysql properties file: " + e.toString());
            }

          } else {
            success = true;
          }

          if (success) {
            // Writing db file succeeded

            try {
              if (dbOverride.equalsIgnoreCase("override")) {
                executeSqlScript();
                htmlOutput =
                    bundle.getString("generic.text.setup.success")
                        + " "
                        + bundle.getString("generic.text.setup.success.overwrittendb");
              } else if (dbOverride.equalsIgnoreCase("upgrade")) {
                executeUpdateScript();
                htmlOutput =
                    bundle.getString("generic.text.setup.success")
                        + " "
                        + bundle.getString("generic.text.setup.success.updatedb");
              } else {
                htmlOutput = bundle.getString("generic.text.setup.success");
              }
              success = true;
            } catch (SQLException e) {
              htmlOutput = bundle.getString("generic.text.setup.failed") + ": " + e.getMessage();
              log.error(bundle.getString("generic.text.setup.failed") + ": " + e.getMessage());
              if (!hasDBFile) {
                FileUtils.deleteQuietly(new File(Constants.MYSQL_DB_PROP));
              }
            }
            // Clean up File as it is not needed anymore. Will Cause a new one to be
            // generated next time too
            removeAuthFile();
            resetInstalledCache();
          }

          if (enableMongoChallenge.equalsIgnoreCase("enable")) {
            if (!Validate.isValidPortNumber(mongodbPort)) {
              htmlOutput = bundle.getString("generic.text.setup.error.valid.port");
              FileUtils.deleteQuietly(new File(Constants.MYSQL_DB_PROP));
            } else {
              Files.write(
                  Paths.get(Constants.MONGO_DB_PROP),
                  mongoProp.toString().getBytes(),
                  StandardOpenOption.CREATE);
              if (MongoDatabase.getMongoDbConnection(null).listDatabaseNames() == null) {
                htmlOutput = bundle.getString("generic.text.setup.connection.mongo.failed");
                if (!hasDBFile) {
                  FileUtils.deleteQuietly(new File(Constants.MYSQL_DB_PROP));
                }
              } else {
                try {
                  executeMongoScript();
                } catch (IOException e) {
                  htmlOutput =
                      bundle.getString("generic.text.setup.failed") + ": " + e.getMessage();
                  if (!hasDBFile) {
                    FileUtils.deleteQuietly(new File(Constants.MYSQL_DB_PROP));
                  }
                }
              }
            }
          }

          if (enableUnsafeLevels.equalsIgnoreCase("enable")) {
            openUnsafeLevels();
            if (!executeCreateChallengeFile()) {
              htmlOutput = bundle.getString("generic.text.setup.file.failed");
              if (!hasDBFile) {
                FileUtils.deleteQuietly(new File(Constants.MYSQL_DB_PROP));
              }
            }
          }
        }
      }
    }
    if (success) {
      htmlOutput =
          "<h2 class=\"title\" id=\"login_title\">"
              + bundle.getString("generic.text.setup.response.success")
              + "</h2><p>"
              + htmlOutput
              + " "
              + bundle.getString("generic.text.setup.response.success.redirecting")
              + "</p>";
    } else {
      log.error("Could not create database...");
      if (!hasDBFile) {
        FileUtils.deleteQuietly(new File(Constants.MYSQL_DB_PROP));
      }
      htmlOutput =
          "<h2 class=\"title\" id=\"login_title\">"
              + bundle.getString("generic.text.setup.response.failed")
              + "</h2><p>"
              + htmlOutput
              + "</p>";
    }
    out.write(htmlOutput);

    out.close();
  }

  /**
   * Validates that db host and port are either both provided or both empty. Returns null if valid,
   * or an error message if invalid.
   */
  static String validateHostPort(String dbHost, String dbPort) {
    if (dbHost == null) dbHost = "";
    if (dbPort == null) dbPort = "";
    if (dbHost.isEmpty() != dbPort.isEmpty()) {
      return "If you override db host and db port, both must be entered!";
    }
    return null;
  }

  public static boolean isInstalled() {
    Boolean cached = installedCached;
    if (cached != null) {
      return cached;
    }

    synchronized (Setup.class) {
      cached = installedCached;
      if (cached != null) {
        return cached;
      }

      boolean installed = false;

      Properties prop = getDBProps();

      if (prop != null) {
        try (Connection coreConnection = Database.getCoreConnection(null)) {
          if (coreConnection != null) {
            installed = true;
          }
        } catch (SQLException e) {
          log.info("isInstalled got SQL exception " + e.toString() + ", assuming not installed.");
        }
      }

      if (installed) {
        installedCached = true;
      } else {
        generateAuth();
      }

      return installed;
    }
  }

  /** Clear the cached installation status so the next call re-evaluates. */
  public static void resetInstalledCache() {
    installedCached = null;
  }

  public static Properties getDBProps() {

    Properties prop = new Properties();

    // Pull Driver and DB URL out of database.properties

    String mysql_props = Constants.MYSQL_DB_PROP;

    try (InputStream mysql_input = new FileInputStream(mysql_props)) {

      prop.load(mysql_input);

      return prop;

    } catch (IOException e) {
      log.info("Could not load properties file, assuming doesn't exist: " + e.toString());
      return null;
    }
  }

  private static void generateAuth() {
    try {
      if (!Files.exists(Paths.get(Constants.SETUP_AUTH), LinkOption.NOFOLLOW_LINKS)) {
        UUID randomUUID = UUID.randomUUID();
        log.info("Auth file not found, creating: " + Constants.SETUP_AUTH);

        Files.write(
            Paths.get(Constants.SETUP_AUTH),
            randomUUID.toString().getBytes(),
            StandardOpenOption.CREATE);
        log.info("Generated UUID " + randomUUID + " in " + Constants.SETUP_AUTH);
      }
    } catch (IOException e) {
      log.fatal("Unable to generate auth: " + e.getMessage());
      throw new RuntimeException(e);
    }
  }

  private static void removeAuthFile() {
    if (!Files.exists(Paths.get(Constants.SETUP_AUTH), LinkOption.NOFOLLOW_LINKS)) {
      log.info("Could not find " + Constants.SETUP_AUTH);
    } else {
      FileUtils.deleteQuietly(new File(Constants.SETUP_AUTH));
    }
  }

  private synchronized void executeSqlScript() throws IOException, SQLException {

    File file =
        new File(getClass().getClassLoader().getResource("/database/coreSchema.sql").getFile());
    String data = FileUtils.readFileToString(file, Charset.defaultCharset());

    log.debug("Initializing core database");
    try (Connection databaseConnection = Database.getDatabaseConnection(null, true)) {
      try (Statement psProcToexecute = databaseConnection.createStatement()) {
        psProcToexecute.executeUpdate(data);
      }

      file =
          new File(
              getClass().getClassLoader().getResource("/database/moduleSchemas.sql").getFile());
      data = FileUtils.readFileToString(file, Charset.defaultCharset());
      log.debug("Initializing module database");

      try (Statement psProcToexecute = databaseConnection.createStatement()) {
        psProcToexecute.executeUpdate(data);
      }
    }
  }

  private synchronized void executeMongoScript() throws IOException {

    MongoClient mongoConnection = null;

    try {
      File file =
          new File(getClass().getClassLoader().getResource("/mongodb/moduleSchemas.js").getFile());
      mongoConnection = MongoDatabase.getMongoDbConnection(null);
      MongoDatabase.executeMongoScript(file, mongoConnection);
    } catch (IOException e) {
      throw e;
    } finally {
      MongoDatabase.closeConnection(mongoConnection);
    }
  }

  private synchronized void executeUpdateScript() throws IOException, SQLException {

    File file =
        new File(
            getClass().getClassLoader().getResource("/database/updatev3_0tov3_1.sql").getFile());

    String data;

    data = FileUtils.readFileToString(file, Charset.defaultCharset());

    try (Connection databaseConnection = Database.getDatabaseConnection(null, true);
        Statement psProcToexecute = databaseConnection.createStatement()) {
      psProcToexecute.executeUpdate(data);
    }
  }

  private synchronized void openUnsafeLevels() {
    String ApplicationRoot = getServletContext().getRealPath("");
    Setter.openAllModules(ApplicationRoot, true);
  }

  private synchronized Boolean executeCreateChallengeFile() {
    if (XxeLesson.createXxeLessonSolutionFile()
        && XxeChallenge1.createXxeChallenge1SolutionFile()) {
      return true;
    }

    return false;
  }
}