Update burp frame grabber to use new API

nccgroup · Dec 2, 2022 · a7ed919 · a7ed919
1 parent dfc6c3f
commit a7ed919
Showing 1 changed file with 8 additions and 23 deletions.
diff --git 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<Frame> 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);