//
// Burp Suite Logger++
// 
// Released as open source by NCC Group Plc - https://www.nccgroup.trust/
// 
// Originally Developed by Soroush Dalili (@irsdl)
// Maintained by Corey Arthur (@CoreyD97)
//
// Project link: http://www.github.com/nccgroup/BurpSuiteLoggerPlusPlus
//
// Released under AGPL see LICENSE for more information
//

package com.nccgroup.loggerplusplus.imports;

import com.nccgroup.loggerplusplus.LoggerPlusPlus;
import com.nccgroup.loggerplusplus.logentry.EntryImportWorker;

import burp.*;
import java.util.ArrayList;
import java.util.Iterator;
import java.io.*;
import java.net.URL;
import java.net.http.HttpRequest;

import javax.swing.*;

public class LoggerImport {

    public static String getLoadFile() {
        JFileChooser chooser = null;
        chooser = new JFileChooser();
        chooser.setDialogTitle("Import File");
        int val = chooser.showOpenDialog(null);

        if (val == JFileChooser.APPROVE_OPTION) {
            return chooser.getSelectedFile().getAbsolutePath();
        }

        return "";
    }

    public static ArrayList<String> readFile(String filename) {
        BufferedReader reader;
        ArrayList<String> lines = new ArrayList<String>();

        try {
            reader = new BufferedReader(new FileReader(filename));
        } catch (FileNotFoundException e) {
            LoggerPlusPlus.callbacks.printError("LoggerImport-readFile: Error Opening File " + filename);
            return new ArrayList<String>();
        }
        try {
            String line;
            while ( (line = reader.readLine()) != null ) {
                lines.add(line);
            }
        } catch (IOException e) {
            LoggerPlusPlus.callbacks.printError("LoggerImport-readFile: Error Reading Line");
            return new ArrayList<String>();
        }

        return lines;
    }

    public static ArrayList<IHttpRequestResponse> importWStalker() {
        ArrayList<String> lines = new ArrayList<String>();
        ArrayList<IHttpRequestResponse> requests = new ArrayList<IHttpRequestResponse>();
        IExtensionHelpers helpers = LoggerPlusPlus.callbacks.getHelpers();
        
        String filename = getLoadFile();
        lines = readFile(filename);

        Iterator<String> i = lines.iterator();
        while (i.hasNext()) {
            try {
                String line = i.next();
                String[] v = line.split(","); // Format: "base64(request),base64(response),url"

                byte[] request = helpers.base64Decode(v[0]);
                byte[] response = helpers.base64Decode(v[1]);
                String url = v[3];

                LoggerRequestResponse x = new LoggerRequestResponse(url, request, response);
                requests.add(x);

            } catch (Exception e) {
                LoggerPlusPlus.callbacks.printError("LoggerImport-importWStalker: Error Parsing Content");
                return new ArrayList<IHttpRequestResponse>();
            }
        }

        return requests;
    }

    public static ArrayList<IHttpRequestResponse> importZAP() {
        ArrayList<String> lines = new ArrayList<String>();
        ArrayList<IHttpRequestResponse> requests = new ArrayList<IHttpRequestResponse>();
        IExtensionHelpers helpers = LoggerPlusPlus.callbacks.getHelpers();
        
        String filename = getLoadFile();
        lines = readFile(filename);
        Iterator<String> i = lines.iterator();

        // Format:
        // ==== [0-9]+ ==========
        // REQUEST
        // <empty>
        // RESPONSE
        String reSeparator = "^==== [0-9]+ ==========$";
        String reResponse = "^HTTP/[0-9]\\.[0-9] [0-9]+ .*$";

        // Ignore first line, since it should be a separator
        if ( i.hasNext() ) {
            i.next();
        }

        boolean isRequest = true;
        String requestBuffer = "";
        String responseBuffer = "";
        String url = "";

        // Loop lines
        while (i.hasNext()) {
            String line = i.next();

            // Request and Response Ready
            if ( line.matches(reSeparator) || !i.hasNext() ) {
                // TODO: Remove one or two \n at the end of requestBuffer

                byte[] req = helpers.stringToBytes(requestBuffer);
                byte[] res = helpers.stringToBytes(responseBuffer);

                // Add IHttpRequestResponse Object
                LoggerRequestResponse x = new LoggerRequestResponse(url, req, res);
                requests.add(x);

                // Reset content
                isRequest = true;
                requestBuffer = "";
                responseBuffer = "";
                url = "";

                continue;
            }

            // It's the beginning of a request
            if ( requestBuffer.length() == 0 ) {
                try {
                    // Expected format: "GET https://whatever/whatever.html HTTP/1.1"
                    String[] x = line.split(" ");
                    url = x[1];

                    URL u = new URL(url);
                    String path = u.getPath();
                    line = x[0] + " " + path + " " + x[2]; // fix the path in the request

                } catch (Exception e) {
                    LoggerPlusPlus.callbacks.printError("importZAP: Wrong Path Format");
                    return new ArrayList<IHttpRequestResponse>();
                } 
            }

            // It's the beginning of a response
            if ( line.matches(reResponse) ) {
                isRequest = false;
            }

            // Add line to the corresponding buffer
            if ( isRequest ) {
                requestBuffer += line;
                requestBuffer += "\n";
            } else {
                responseBuffer += line;
                responseBuffer += "\n";
            }
        }

        return requests;
    }

    public static boolean loadImported(ArrayList<IHttpRequestResponse> requests) {
        EntryImportWorker importWorker = LoggerPlusPlus.instance.getLogProcessor().createEntryImportBuilder()
            .setOriginatingTool(IBurpExtenderCallbacks.TOOL_EXTENDER)
            .setEntries(requests)
            .setInterimConsumer(integers -> {
                //Optional
                //Outputs chunks of integers representing imported indices
                //May be used to update progress bar for example
            })
            .setCallback(() -> {
                //Optional
                //Called when all entries have been imported.
            }).build();
        importWorker.execute();

        return true;
    }
}
@@ -2184,32 +2184,169 @@
  getLoadFile();%0A
+        if ( filename.length() == 0 ) %7B // exit if no file selected%0A            return new ArrayList%3CIHttpRequestResponse%3E();%0A        %7D%0A%0A
         lines = 
@@ -2365,17 +2365,16 @@
 ename);%0A
-%0A
         
@@ -2408,24 +2408,33 @@
 iterator();%0A
+        %0A
         whil
@@ -3484,32 +3484,169 @@
  getLoadFile();%0A
+        if ( filename.length() == 0 ) %7B // exit if no file selected%0A            return new ArrayList%3CIHttpRequestResponse%3E();%0A        %7D%0A%0A
         lines = 

Fix error when no file selected