Merge pull request #92 from Nizernizer/feature/issue-181

update API Sitemap optimization
HXSecurity · Sep 8, 2021 · d6641e8 · d6641e8
2 parents 8bccf08 + 9fe53ba
commit d6641e8
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 22 deletions.
diff --git a/iast-core/src/main/java/com/secnium/iast/core/enhance/plugins/api/SpringApplicationImpl.java b/iast-core/src/main/java/com/secnium/iast/core/enhance/plugins/api/SpringApplicationImpl.java
@@ -2,14 +2,13 @@
 
 import com.secnium.iast.core.handler.IastClassLoader;
 import com.secnium.iast.core.handler.controller.impl.HttpImpl;
-import com.secnium.iast.core.handler.models.ApiDataModel;
 import com.secnium.iast.core.handler.models.MethodEvent;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
-import java.util.*;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static com.secnium.iast.core.report.ApiReport.sendReport;
@@ -28,18 +27,15 @@ public static void getWebApplicationContext(MethodEvent event, AtomicInteger inv
         Object applicationContext = event.returnValue;
         createClassLoader(applicationContext);
         loadApplicationContext();
-        List<ApiDataModel> invoke = null;
-        String apiList = null;
+        Map<String, Object> invoke = null;
         try {
-            invoke = (List<ApiDataModel>) getAPI.invoke(null, applicationContext);
-            apiList = invoke.toString();
-            apiList = apiList.replace("=", ":").replace("{", "{\"").replace("}", "\"}").replace(" ", "").replace(":", "\":\"").replace("'", "").replace(",", "\",\"").replace("\"[", "[\"").replace("]\"", "\"]").replace("null", "").replace("}\"", "}").replace("\"{", "{").replace("/{\"", "/{").replace("\"},\"met", "}\",\"met").replace("clazz", "class");
+            invoke = (Map<String, Object>) getAPI.invoke(null, applicationContext);
+            sendReport(invoke);
         } catch (IllegalAccessException e) {
             e.printStackTrace();
         } catch (InvocationTargetException e){
             e.printStackTrace();
         }
-        sendReport(apiList);
         isSend = true;
         }
     }

diff --git a/iast-core/src/main/java/com/secnium/iast/core/report/ApiReport.java b/iast-core/src/main/java/com/secnium/iast/core/report/ApiReport.java
@@ -2,11 +2,8 @@
 
 import com.secnium.iast.core.EngineManager;
 import com.secnium.iast.core.handler.vulscan.ReportConstant;
-import org.json.JSONArray;
 import org.json.JSONObject;
-import com.secnium.iast.core.handler.models.ApiDataModel;
 
-import java.util.List;
 import java.util.Map;
 
 /**
@@ -16,22 +13,20 @@
  */
 public class ApiReport {
 
-    public static void sendReport(String apiList) {
+    public static void sendReport(Map<String, Object> apiList) {
          String report = createReport(apiList);
          EngineManager.sendNewReport(report);
     }
 
-    private static String createReport(String apiList) {
+    private static String createReport(Map<String, Object> apiList) {
         JSONObject report = new JSONObject();
         JSONObject detail = new JSONObject();
-        JSONArray api = new JSONArray();
         report.put(ReportConstant.REPORT_KEY, ReportConstant.REPORT_API);
         report.put(ReportConstant.REPORT_VALUE_KEY, detail);
         detail.put(ReportConstant.AGENT_ID, AgentRegisterReport.getAgentFlag());
-        detail.put(ReportConstant.API_DATA, api);
-        String result = report.toString();
-        result = result.replace("\"api_data\":[]", "\"api_data\":" + apiList).replace("url","uri").replace("returnType","return_type").replace("\"method\":[]","\"method\":[\"GET\",\"POST\"]");
-        return result;
+        JSONObject apiListJson = new JSONObject(apiList);
+        detail.put(ReportConstant.API_DATA, apiListJson.get(ReportConstant.API_DATA));
+        return report.toString();
     }
 
 }
diff --git a/iast-servlet/src/main/java/cn/huoxian/iast/servlet/SpringApplicationContext.java b/iast-servlet/src/main/java/cn/huoxian/iast/servlet/SpringApplicationContext.java
@@ -14,14 +14,14 @@
 
 public class SpringApplicationContext {
 
-    public static List<Object> getAPI(Object applicationContext) {
-        return getAPIList((ApplicationContext) applicationContext);
+    public static Map<String, Object> getAPI(Object applicationContext) {
+        return createReport(getAPIList((ApplicationContext) applicationContext));
     }
 
-    public static List<Object> getAPIList(ApplicationContext applicationContext) {
+    public static List<ApiDataModel> getAPIList(ApplicationContext applicationContext) {
         RequestMappingHandlerMapping mapping = applicationContext.getBean(RequestMappingHandlerMapping.class);
         Map<RequestMappingInfo, HandlerMethod> methodMap = mapping.getHandlerMethods();
-        List<Object> apiList = new ArrayList<>();
+        List<ApiDataModel> apiList = new ArrayList<>();
         for (RequestMappingInfo info : methodMap.keySet()) {
             ApiDataModel apiDataModel = new ApiDataModel();
             HandlerMethod handlerMethod = methodMap.get(info);
@@ -121,4 +121,36 @@ public static List<Object> getAPIList(ApplicationContext applicationContext) {
         return apiList;
     }
 
+    private static Map<String, Object> createReport(List<ApiDataModel> apiList) {
+        Map<String, Object> apiDataReport = new HashMap<>();
+        List<Object> apiData = new ArrayList<>();
+        for (ApiDataModel apiDataModel:apiList
+        ) {
+            Map<String, Object> api = new HashMap<>();
+            apiData.add(api);
+            api.put("uri",apiDataModel.getUrl());
+            String[] methods = apiDataModel.getMethod();
+            List<Object> methodsjson = new ArrayList<>(Arrays.asList(methods));
+            api.put("method",methodsjson);
+            api.put("class",apiDataModel.getClazz());
+            List<Map<String, String>> parameters = apiDataModel.getParameters();
+            List<Object> parametersJson = new ArrayList<>();
+            api.put("parameters",parametersJson);
+            for (Map<String,String> parameter:parameters
+            ) {
+                Map<String, Object> parameterjson = new HashMap<>();
+                parametersJson.add(parameterjson);
+                parameterjson.put("name",parameter.get("name"));
+                parameterjson.put("type",parameter.get("type"));
+                parameterjson.put("annotation",parameter.get("annotation"));
+            }
+            api.put("return_type",apiDataModel.getReturnType());
+            api.put("file",apiDataModel.getFile());
+            api.put("controller",apiDataModel.getController());
+            api.put("description",apiDataModel.getDescription());
+        }
+        apiDataReport.put("api_data",apiData);
+        return apiDataReport;
+    }
+
 }