com.denimgroup.threadfix.importer.cli.CommandLineMigration.java Source code

Java tutorial

Introduction

Here is the source code for com.denimgroup.threadfix.importer.cli.CommandLineMigration.java

Source

////////////////////////////////////////////////////////////////////////
//
//     Copyright (c) 2009-2015 Denim Group, Ltd.
//
//     The contents of this file are subject to the Mozilla Public License
//     Version 2.0 (the "License"); you may not use this file except in
//     compliance with the License. You may obtain a copy of the License at
//     http://www.mozilla.org/MPL/
//
//     Software distributed under the License is distributed on an "AS IS"
//     basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See the
//     License for the specific language governing rights and limitations
//     under the License.
//
//     The Original Code is ThreadFix.
//
//     The Initial Developer of the Original Code is Denim Group, Ltd.
//     Portions created by Denim Group, Ltd. are Copyright (C)
//     Denim Group, Ltd. All Rights Reserved.
//
//     Contributor(s): Denim Group, Ltd.
//
////////////////////////////////////////////////////////////////////////

package com.denimgroup.threadfix.importer.cli;

import com.denimgroup.threadfix.StringEscapeUtils;
import com.denimgroup.threadfix.importer.config.SpringConfiguration;
import com.denimgroup.threadfix.importer.util.RegexUtils;
import com.denimgroup.threadfix.logging.SanitizedLogger;
import org.apache.commons.io.FileUtils;
import org.springframework.stereotype.Component;

import java.io.*;
import java.util.List;
import java.util.Map;

import static com.denimgroup.threadfix.CollectionUtils.list;
import static com.denimgroup.threadfix.CollectionUtils.map;

@Component
public class CommandLineMigration {

    private static final SanitizedLogger LOGGER = new SanitizedLogger(CommandLineMigration.class);

    private static final String TABLE_PATTERN = "CREATE MEMORY TABLE (.*)";
    private static final String INSERT_PATTERN = "INSERT INTO (.*) VALUES";
    private static final String ACUNETIX_ESCAPE = "Cross-Site Scripting in HTML \\''script\\'' tag";
    private static final String ACUNETIX_ESCAPE_REPLACE = "Cross-Site Scripting in HTML "script" tag";

    private static Map<String, String> tableMap = map();

    private static final int SAME_SET_TRY_LIMIT = 10;

    public static void main(String[] args) {

        if (!check(args))
            return;

        try {
            String inputScript = args[0];
            String inputMySqlConfig = args[1];
            String outputScript = "import.sql";
            String outputMySqlConfigTemp = "jdbc_temp.properties";
            String errorLogFile = "error.log";
            String fixedSqlFile = "error_sql.sql";
            String errorLogAttemp1 = "error1.log";
            String infoLogFile = "info.log";
            String rollbackScript = "rollback.sql";

            deleteFile(outputScript);
            deleteFile(errorLogFile);
            deleteFile(errorLogFile);
            deleteFile(fixedSqlFile);
            deleteFile(errorLogAttemp1);
            deleteFile(infoLogFile);
            deleteFile(rollbackScript);

            PrintStream infoPrintStream = new PrintStream(new FileOutputStream(new File(infoLogFile)));
            System.setOut(infoPrintStream);

            long startTime = System.currentTimeMillis();

            copyFile(inputMySqlConfig, outputMySqlConfigTemp);

            LOGGER.info("Creating threadfix table in mySql database ...");
            ScriptRunner scriptRunner = SpringConfiguration.getContext().getBean(ScriptRunner.class);

            startTime = printTimeConsumed(startTime);
            convert(inputScript, outputScript);

            startTime = printTimeConsumed(startTime);

            PrintStream errPrintStream = new PrintStream(new FileOutputStream(new File(errorLogFile)));
            System.setErr(errPrintStream);

            LOGGER.info("Sending sql script to MySQL server ...");
            scriptRunner.run(outputScript, outputMySqlConfigTemp);

            long errorCount = scriptRunner.checkRunningAndFixStatements(errorLogFile, fixedSqlFile);
            long lastCount = errorCount + 1;
            int times = 1;
            int sameFixedSet = 0;

            // Repeat
            while (errorCount > 0) {
                //Flush error log screen to other file
                errPrintStream = new PrintStream(new FileOutputStream(new File(errorLogAttemp1)));
                System.setErr(errPrintStream);

                times += 1;

                if (errorCount == lastCount) {
                    sameFixedSet++;
                } else {
                    sameFixedSet = 0;
                }

                LOGGER.info("Found " + errorCount + " error statements. Sending fixed sql script to MySQL server "
                        + times + " times ...");
                scriptRunner.run(fixedSqlFile, outputMySqlConfigTemp);
                lastCount = errorCount;
                errorCount = scriptRunner.checkRunningAndFixStatements(errorLogAttemp1, fixedSqlFile);

                if (errorCount > lastCount || sameFixedSet > SAME_SET_TRY_LIMIT)
                    break;
            }

            if (errorCount > 0) {
                LOGGER.error("After " + times + " of trying, still found errors in sql script. "
                        + "Please check error_sql.sql and error1.log for more details.");
                LOGGER.info("Do you want to keep data in MySQL, and then import manually error statements (y/n)? ");
                try (java.util.Scanner in = new java.util.Scanner(System.in)) {
                    String answer = in.nextLine();
                    if (!answer.equalsIgnoreCase("y")) {
                        rollbackData(scriptRunner, outputMySqlConfigTemp, rollbackScript);
                    } else {
                        LOGGER.info(
                                "Data imported to MySQL, but still have some errors. Please check error_sql.sql and error1.log to import manually.");
                    }

                }

            } else {
                printTimeConsumed(startTime);
                LOGGER.info("Migration successfully finished");
            }

            deleteFile(outputMySqlConfigTemp);

        } catch (Exception e) {
            LOGGER.error("Error: ", e);
        }
    }

    private static void convert(String inputScript, String outputScript) {
        File file = new File(inputScript);

        LOGGER.info("Converting threadfix script to mySql script " + outputScript + " ...");

        File outputFile = new File(outputScript);

        FileOutputStream fos = null;
        try {
            fos = new FileOutputStream(outputFile);

            OutputStreamWriter osw = new OutputStreamWriter(fos);

            List<String> lines = FileUtils.readLines(file);

            osw.write("SET FOREIGN_KEY_CHECKS=0;\n");

            String table;
            for (String line : lines) {
                if (line != null && line.toUpperCase().startsWith("CREATE MEMORY TABLE ")) {
                    table = RegexUtils.getRegexResult(line, TABLE_PATTERN);
                    System.out.println("Create new table:" + table);
                    String[] tableName = table.split("\\(", 2);
                    if (tableName.length == 2) {
                        List<String> fieldList = list();
                        String[] fields = tableName[1].trim().replace("(", "").replace(")", "").split(",");
                        for (int i = 0; i < fields.length; i++) {
                            if (!"CONSTRAINT".equalsIgnoreCase(fields[i].trim().split(" ")[0])) {
                                String field = fields[i].trim().split(" ")[0].replace("\"", "");
                                if (!fieldList.contains(field))
                                    fieldList.add(field);
                            }
                        }
                        String fieldsStr = org.apache.commons.lang3.StringUtils.join(fieldList, ",");
                        tableMap.put(tableName[0].toUpperCase(), "(" + fieldsStr + ")");
                    }
                } else if (line != null && line.toUpperCase().startsWith("INSERT INTO ")) {
                    table = RegexUtils.getRegexResult(line, INSERT_PATTERN).toUpperCase();
                    if (tableMap.get(table) != null) {
                        line = line.replaceFirst(" " + table + " ", " " + table + tableMap.get(table) + " ");
                        if (line.contains(ACUNETIX_ESCAPE)) {
                            line = line.replace(ACUNETIX_ESCAPE, ACUNETIX_ESCAPE_REPLACE);
                        }
                        line = escapeString(line) + ";\n";

                        osw.write(line);
                    }
                }
            }
            osw.write("SET FOREIGN_KEY_CHECKS=1;\n");
            osw.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static long printTimeConsumed(long startTime) {
        long endTime = System.currentTimeMillis();
        LOGGER.info("Finished in " + (endTime - startTime) + " ms");
        return endTime;
    }

    private static boolean check(String[] args) {
        if (args.length != 2) {
            LOGGER.warn("This program accepts two argument, threadfix script and mysql config files.");
            return false;
        }

        for (String arg : args) {
            File file = new File(arg);
            if (!file.exists()) {
                LOGGER.warn(arg + ": The file must exist.");
                return false;
            }

            if (file.isDirectory()) {
                LOGGER.warn(arg + ": The file must not be a directory.");
                return false;
            }
        }
        return true;
    }

    private static void copyFile(String oldFilePath, String newFilePath) throws IOException {
        FileUtils.writeStringToFile(new File(newFilePath), FileUtils.readFileToString(new File(oldFilePath)));
        LOGGER.info("Copied from " + oldFilePath + " to " + newFilePath);
    }

    private static void deleteFile(String oldFilePath) throws IOException {

        File fouput = new File(oldFilePath);
        if (fouput.exists() && fouput.isFile())
            if (fouput.delete())
                LOGGER.info("File " + oldFilePath + " has been deleted");
            else
                LOGGER.info("File " + oldFilePath + " has not been deleted");
    }

    private static String escapeString(String line) {

        if (line.toUpperCase().contains("INSERT INTO APPLICATION(")
                || line.toUpperCase().contains("INSERT INTO DEFAULTCONFIGURATION(")
                || line.toUpperCase().contains("INSERT INTO USER(")
                || line.toUpperCase().contains("INSERT INTO FINDING(")) {
            line = StringEscapeUtils.unescapeUnicode(line);
        }

        if (line.toUpperCase().contains("INSERT INTO FINDING")) {
            //Double backslash for multi backslash group
            line = line.replaceAll("(\\\\[\\\\]+)", "$1$1");
            // Double for single backslash
            line = line.replaceAll("([^\\\\])(\\\\[^nrt\\\\])", "$1\\\\$2");
        }

        return line;
    }

    private static void rollbackData(ScriptRunner runner, String jdbcConfig, String rollbackScript) {
        try {
            LOGGER.info("Rolling back data...");
            FileUtils.writeStringToFile(new File(rollbackScript), "DROP DATABASE threadfix;");
            runner.run(rollbackScript, jdbcConfig);
        } catch (IOException e) {
            LOGGER.error("Error", e);
        }
    }
}