diff --git a/src/main/java/com/nccgroup/loggerplusplus/LoggerPlusPlus.java b/src/main/java/com/nccgroup/loggerplusplus/LoggerPlusPlus.java index 29d37deb..74c1419e 100644 --- a/src/main/java/com/nccgroup/loggerplusplus/LoggerPlusPlus.java +++ b/src/main/java/com/nccgroup/loggerplusplus/LoggerPlusPlus.java @@ -23,8 +23,6 @@ import java.awt.*; import java.util.Arrays; import java.util.List; -import java.util.stream.Collectors; -import java.util.stream.Collectors; import static com.nccgroup.loggerplusplus.util.Globals.PREF_RESTRICT_TO_SCOPE; @@ -93,34 +91,21 @@ public void initialize(MontoyaApi montoya) { montoya.proxy().registerResponseHandler(logProcessor.getProxyHttpResponseHandler()); //Add menu item to Burp's frame menu. - JFrame rootFrame = null; + Frame rootFrame = null; try { - rootFrame = getBurpFrame(); - JMenuBar menuBar = rootFrame.getJMenuBar(); - if(menuBar != null) { - loggerMenu = new LoggerMenu(LoggerPlusPlus.this); - menuBar.add(loggerMenu, menuBar.getMenuCount() - 1); + rootFrame = montoya.userInterface().swingUtils().suiteFrame(); + if (rootFrame instanceof JFrame) { + JMenuBar menuBar = ((JFrame) rootFrame).getJMenuBar(); + if (menuBar != null) { + loggerMenu = new LoggerMenu(LoggerPlusPlus.this); + menuBar.add(loggerMenu, menuBar.getMenuCount() - 1); + } } } catch (Exception e) { log.error("Could not find root frame. Window JMenu will not be added"); } } - private JFrame getBurpFrame() throws Exception { - // Get all frames - Frame[] allFrames = JFrame.getFrames(); - // Filter the stream find the main burp window frame, and convert to a list - List filteredFrames = Arrays.stream(allFrames).filter(f -> - f.getTitle().startsWith("Burp Suite") && f.isVisible() - ).collect(Collectors.toList()); - // If size is 1, we have the main burp frame. Otherwise fails - if (filteredFrames.size() == 1) { - return (JFrame) filteredFrames.get(0); - } else { - throw new Exception("Expected one burp pane, but found " + filteredFrames.size()); - } - } - public void unloadExtension() { if(loggerMenu != null && loggerMenu.getParent() != null){ loggerMenu.getParent().remove(loggerMenu);