Simplify API key validity check

RevEngAI · Oct 12, 2024 · 0282a33 · 0282a33
1 parent 718c4fd
commit 0282a33
Show file tree

Hide file tree

Showing 10 changed files with 34 additions and 43 deletions.
diff --git a/...ai/reveng/toolkit/ghidra/binarysimularity/ui/recentanalyses/RecentAnalysesTableModel.java b/...ai/reveng/toolkit/ghidra/binarysimularity/ui/recentanalyses/RecentAnalysesTableModel.java
@@ -1,8 +1,8 @@
 package ai.reveng.toolkit.ghidra.binarysimularity.ui.recentanalyses;
 
-import ai.reveng.toolkit.ghidra.binarysimularity.ui.autoanalysis.AutoAnalysisResultsTableModel;
 import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService;
 import ai.reveng.toolkit.ghidra.core.services.api.types.*;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.APIAuthenticationException;
 import docking.widgets.table.AbstractDynamicTableColumn;
 import docking.widgets.table.TableColumnDescriptor;
 import docking.widgets.table.threaded.ThreadedTableModelStub;

diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java
@@ -4,6 +4,7 @@
 import ai.reveng.toolkit.ghidra.core.RevEngAIAnalysisStatusChanged;
 import ai.reveng.toolkit.ghidra.core.services.api.types.*;
 import ai.reveng.toolkit.ghidra.core.services.api.types.Collection;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.APIAuthenticationException;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import ghidra.app.util.opinion.ElfLoader;

diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/MockApi.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/MockApi.java
@@ -1,6 +1,7 @@
 package ai.reveng.toolkit.ghidra.core.services.api;
 
 import ai.reveng.toolkit.ghidra.core.services.api.types.*;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.APIAuthenticationException;
 import org.json.JSONObject;
 
 import java.io.FileNotFoundException;
@@ -206,6 +207,10 @@ public String getAnalysisLogs(BinaryID binID) {
         return null;
     }
 
+    @Override
+    public void authenticate() throws APIAuthenticationException {
+    }
+
     @Override
     public List<ModelName> models() {
         return null;

diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiImplementation.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiImplementation.java
@@ -2,6 +2,8 @@
 
 import ai.reveng.toolkit.ghidra.core.services.api.types.*;
 import ai.reveng.toolkit.ghidra.core.services.api.types.Collection;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.APIAuthenticationException;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.InvalidAPIInfoException;
 import com.google.common.primitives.Bytes;
 import org.json.JSONObject;
 
@@ -382,28 +384,16 @@ public List<ModelName> models(){
         return result;
     }
 
-    public boolean checkCredentials(){
+    @Override
+    public void authenticate() throws InvalidAPIInfoException {
         var request = requestBuilderForEndpoint("authenticate")
-                .GET()
                 .build();
-        HttpResponse<String> response  = null;
         try {
-            response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        } catch (InterruptedException e) {
-            throw new RuntimeException(e);
-        }
-
-        switch (response.statusCode()){
-            case 200:
-
-                return true;
-            case 401:
-                return false;
-            default:
-                throw new RuntimeException("Request with unexpected status code: " + response.statusCode() + " and message: " + response.body());
+            sendRequest(request);
+        } catch (APIAuthenticationException e) {
+            throw new InvalidAPIInfoException("Invalid API key", e);
         }
     }
+
 }
 
diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiInterface.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/TypedApiInterface.java
@@ -5,6 +5,8 @@
 import java.util.List;
 
 import ai.reveng.toolkit.ghidra.core.services.api.types.*;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.APIAuthenticationException;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.InvalidAPIInfoException;
 
 
 /**
@@ -107,5 +109,6 @@ List<FunctionMatch> annSymbolsForBinary(
 
     String getAnalysisLogs(BinaryID binID);
 
+    void authenticate() throws InvalidAPIInfoException;
 }
 
diff --git a/...ain/java/ai/reveng/toolkit/ghidra/core/services/api/types/APIAuthenticationException.java b/...ain/java/ai/reveng/toolkit/ghidra/core/services/api/types/APIAuthenticationException.java
diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/types/ApiInfo.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/types/ApiInfo.java
@@ -2,16 +2,15 @@
 
 import ai.reveng.toolkit.ghidra.core.models.ReaiConfig;
 import ai.reveng.toolkit.ghidra.core.services.api.TypedApiImplementation;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.InvalidAPIInfoException;
 import com.google.gson.Gson;
 import org.json.JSONException;
 
 import java.io.FileNotFoundException;
 import java.io.FileReader;
-import java.io.IOException;
 import java.net.URI;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.Optional;
 
 public record ApiInfo(
         URI hostURI,
@@ -42,18 +41,8 @@ public void checkCredentials() throws InvalidAPIInfoException {
             throw new InvalidAPIInfoException("Server health check failed: " + health.getString("message"));
         }
 
-        Boolean credentialsValid = null;
-        try {
-            credentialsValid = api.checkCredentials();
-
-        } catch (JSONException e) {
-            throw new InvalidAPIInfoException("Invalid JSON response from server " + hostURI, e);
-        } catch (Exception e) {
-            throw new InvalidAPIInfoException("Failed to validate credentials", e);
-        }
-        if (!credentialsValid){
-            throw new InvalidAPIInfoException("Invalid API key");
-        }
+        // Throws InvalidAPIInfoException if authentication fails
+        api.authenticate();
 
     }
 
@@ -78,5 +67,4 @@ public static ApiInfo fromConfig() throws FileNotFoundException {
         return fromConfig(configFilePath);
 
     }
-
 }
diff --git a/.../reveng/toolkit/ghidra/core/services/api/types/exceptions/APIAuthenticationException.java b/.../reveng/toolkit/ghidra/core/services/api/types/exceptions/APIAuthenticationException.java
@@ -0,0 +1,11 @@
+package ai.reveng.toolkit.ghidra.core.services.api.types.exceptions;
+
+/**
+ * This exception indicates an unexpected case of the API returning an authentication error
+ * this can happen when attempting to retrieve information about a private analysis ID owned by a different account
+ */
+public class APIAuthenticationException extends RuntimeException{
+    public APIAuthenticationException(String message) {
+        super(message);
+    }
+}
diff --git a/...es/api/types/InvalidAPIInfoException.java → ...s/exceptions/InvalidAPIInfoException.java b/...es/api/types/InvalidAPIInfoException.java → ...s/exceptions/InvalidAPIInfoException.java
@@ -1,4 +1,4 @@
-package ai.reveng.toolkit.ghidra.core.services.api.types;
+package ai.reveng.toolkit.ghidra.core.services.api.types.exceptions;
 
 public class InvalidAPIInfoException extends Exception {
     public InvalidAPIInfoException(String message) {

diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/ui/wizard/panels/UserCredentialsPanel.java b/src/main/java/ai/reveng/toolkit/ghidra/core/ui/wizard/panels/UserCredentialsPanel.java
@@ -3,7 +3,7 @@
 import javax.swing.*;
 
 import ai.reveng.toolkit.ghidra.core.services.api.types.ApiInfo;
-import ai.reveng.toolkit.ghidra.core.services.api.types.InvalidAPIInfoException;
+import ai.reveng.toolkit.ghidra.core.services.api.types.exceptions.InvalidAPIInfoException;
 import ai.reveng.toolkit.ghidra.core.ui.wizard.SetupWizardStateKey;
 import docking.wizard.AbstractMageJPanel;
 import docking.wizard.IllegalPanelStateException;