/*
 * Note - you need to rename this file to BurpExtender.java before compiling it
 */

package burp;

import java.net.URL;
import java.util.*;
import java.util.regex.*;
import java.io.*;

public class BurpExtender
{
    public IBurpExtenderCallbacks mCallbacks;

    public byte[] processProxyMessage(
            int messageReference,
            boolean messageIsRequest,
            String remoteHost,
            int remotePort,
            boolean serviceIsHttps,
            String httpMethod,
            String url,
            String resourceType,
            String statusCode,
            String responseContentType,
            byte[] message,
            int[] interceptAction)
    {
        if (!messageIsRequest)
        {
            try
            {
                URL uUrl = new URL(serviceIsHttps ? "HTTPS" : "HTTP", remoteHost, remotePort, url);
                if (mCallbacks.isInScope(uUrl))
                {
                    ProcessHTMLComments(new String(message), remoteHost, url);
                }
            }
            catch (Exception e)
            {
                e.printStackTrace();
            }
        }

        return message;
    }

    private void ProcessHTMLComments(String message, String host, String url)
    {
        try
        {
            // Create matcher 
            Pattern pattern = Pattern.compile("<!--\n{0,}.+?\n{0,}-->");
            Matcher matcher = pattern.matcher(message);
            boolean printed = false;

            // Find all matches and print the url only one time
            while (matcher.find())
            {
                if (!printed)
                {
                    String header = "HTML COMMENT IN:" + host + url + "\r\n==========================";
                    System.out.println(header);
                    SaveToFile(host, header, true);
                    printed = true;
                }
                // Get the matching string
                String comment = matcher.group();
                System.out.println(comment);
                SaveToFile(host, comment, false);
            }
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }
    }

    private void SaveToFile(String fileName, String st2write, boolean printTime)
    {
        File aFile = new File(fileName + ".txt");
        Date now = new Date();
        try
        {
            BufferedWriter out = new BufferedWriter(new FileWriter(aFile, aFile.exists()));
            if (printTime)
            {
                out.write("\r\n\r\n" + now.toString() + "\r\n");
            }
            out.write(st2write + "\r\n");
            out.close();
        }
        catch (IOException e)
        {
            e.printStackTrace();
        }
    }

    public void registerExtenderCallbacks(IBurpExtenderCallbacks callbacks)
    {
        mCallbacks = callbacks;
    }
}