package burp;

import burp.application.apitypes.ApiType;

import burp.utils.Executor;

import javax.swing.*;

import java.awt.event.ActionEvent;

import java.awt.event.ActionListener;

import java.util.ArrayList;

import java.util.HashSet;

import java.util.List;

import java.util.concurrent.CompletableFuture;

public class ContextMenu implements IContextMenuFactory {

private static final HashSet<Byte> availableToolFlag = new HashSet<>();

static {

availableToolFlag.add(IContextMenuInvocation.CONTEXT_PROXY_HISTORY);

availableToolFlag.add(IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_REQUEST);

availableToolFlag.add(IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_REQUEST);

}

@Override

public List<JMenuItem> createMenuItems(IContextMenuInvocation invocation) {

if (availableToolFlag.contains(invocation.getInvocationContext())) {

ArrayList<JMenuItem> menuItemList = new ArrayList<>();

JMenuItem menuItem = new JMenuItem("Do API scan");

menuItem.addActionListener(new ContextMenuActionListener(invocation));

menuItemList.add(menuItem);

return menuItemList;

} else {

return null;

}

}

static class ContextMenuActionListener implements ActionListener {

IContextMenuInvocation invocation;

public ContextMenuActionListener(IContextMenuInvocation invocation) {

this.invocation = invocation;

}

@Override

public void actionPerformed(ActionEvent actionEvent) {

CompletableFuture.supplyAsync(() -> {

PassiveScanner passiveScanner = BurpExtender.getPassiveScanner();

IHttpRequestResponse[] httpRequestResponses = invocation.getSelectedMessages();

for (IHttpRequestResponse httpRequestResponse : httpRequestResponses) {

ArrayList<ApiType> apiTypes = passiveScanner.getApiScanner().detect(httpRequestResponse, false);

passiveScanner.parseApiDocument(apiTypes);

}

return null;

}, Executor.getExecutor());

}

}

}