From 343403a427b67c36210f0d8834e83c34d921a70b Mon Sep 17 00:00:00 2001
From: Levente Kovacs <levente.kovacs@real-digital.de>
Date: Wed, 24 Apr 2024 13:53:53 +0200
Subject: [PATCH] Add option to load ATOR config automatically on startup

---
 src/main/java/burp/BurpExtender.java | 18 ++++++++-
 src/main/java/burp/ImportATOR.java   | 57 ++++++++++++++--------------
 2 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/src/main/java/burp/BurpExtender.java b/src/main/java/burp/BurpExtender.java
index 8e1fb2a..39a78c0 100644
--- a/src/main/java/burp/BurpExtender.java
+++ b/src/main/java/burp/BurpExtender.java
@@ -1,9 +1,15 @@
 package burp;
 
+import burp.ImportATOR;
 import javax.swing.*;
 import java.awt.*;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Properties;
+import java.util.Map;
+import java.io.File;
+import org.json.simple.JSONObject;
+import org.json.simple.parser.JSONParser;
 
 
 public class BurpExtender implements IBurpExtender, IContextMenuFactory, ITab,  IHttpListener {
@@ -40,8 +46,18 @@ public void registerExtenderCallbacks(IBurpExtenderCallbacks iBurpExtenderCallba
         extension.append("[*] ");
         extension.append(EXTENSION_NAME);
         
+		// Load default JSON config, if available
+		ImportATOR importer = new ImportATOR(callbacks);
+		String filepath = System.getenv().get("ATOR_CONFIG_FILEPATH");
+		if(filepath != null) {
+			File f = new File(filepath);
+			if(f.exists() && !f.isDirectory()) {
+				importer.loadJSONFile(filepath);
+				callbacks.printOutput("ATOR configuration loaded successfully from " + filepath);
+			}
+		}
         callbacks.printOutput("ATOR loaded successfully");
-     }
+    }
 
 	public BurpExtender getComponent() {
 		return BurpExtender.this;
diff --git a/src/main/java/burp/ImportATOR.java b/src/main/java/burp/ImportATOR.java
index 15ecea4..02df95d 100644
--- a/src/main/java/burp/ImportATOR.java
+++ b/src/main/java/burp/ImportATOR.java
@@ -24,36 +24,37 @@ public ImportATOR(IBurpExtenderCallbacks callbacks) {
 	
 	public void readJSONFile() {
 		String filePath = null;
-        JFileChooser fileSelector = new JFileChooser(FileSystemView.getFileSystemView().getHomeDirectory()); 
-        int fileChoosenState = fileSelector.showOpenDialog(null); 
-        if (fileChoosenState == JFileChooser.APPROVE_OPTION) 
-        	filePath = fileSelector.getSelectedFile().getAbsolutePath(); 
-        if(filePath != null)
-        {
-        	SetttingsTab.importATORFile.setText(filePath);
-        	JSONParser jsonParser = new JSONParser();
-         
-	        try (FileReader reader = new FileReader(filePath))
-	        {
-	        	JSONObject jsonObject = (JSONObject) jsonParser.parse(reader);
-	        	
-	        	JSONObject errorCondition = (JSONObject) jsonObject.get("errorCondition");
-	        	
-	        	parseErrorCondition(errorCondition);
-	        	
-	        	JSONObject obtainToken = (JSONObject) jsonObject.get("obtainToken");
-	        	parseObtainToken(obtainToken);
-	        	JSONObject errorConditionReplacement = (JSONObject) jsonObject.get("errorConditionReplacement");
-	        	parseErrorConditionReplacement(errorConditionReplacement);
-	        	
-	        }
-	        catch(Exception exp) {
-	        	callbacks.printOutput("Exception while importing file.."+ exp.getMessage());
-	        }
-       }
-        
+		JFileChooser fileSelector = new JFileChooser(FileSystemView.getFileSystemView().getHomeDirectory()); 
+		int fileChoosenState = fileSelector.showOpenDialog(null); 
+		if (fileChoosenState == JFileChooser.APPROVE_OPTION)
+			filePath = fileSelector.getSelectedFile().getAbsolutePath(); 
+		if(filePath != null)
+		{
+			loadJSONFile(filePath);
+		}
 	}
 	
+	public void loadJSONFile(String filePath) {
+		SetttingsTab.importATORFile.setText(filePath);
+		JSONParser jsonParser = new JSONParser();
+
+		try (FileReader reader = new FileReader(filePath))
+		{
+			JSONObject jsonObject = (JSONObject) jsonParser.parse(reader);
+
+			JSONObject errorCondition = (JSONObject) jsonObject.get("errorCondition");
+
+		    parseErrorCondition(errorCondition);
+
+			JSONObject obtainToken = (JSONObject) jsonObject.get("obtainToken");
+			parseObtainToken(obtainToken);
+			JSONObject errorConditionReplacement = (JSONObject) jsonObject.get("errorConditionReplacement");
+			parseErrorConditionReplacement(errorConditionReplacement);
+	    }
+        catch(Exception exp) {
+			callbacks.printOutput("Exception while importing file.."+ exp.getMessage());
+        }
+	}
 	
 	public void parseErrorCondition(JSONObject jsonObject) {