fix endpoint, data field, and spec generation (#40)

metlo-labs · Oct 19, 2022 · 2c12831 · 2c12831
1 parent 1b190e8
commit 2c12831
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 36 deletions.
diff --git a/backend/src/services/data-field/index.ts b/backend/src/services/data-field/index.ts
@@ -52,7 +52,7 @@ export class DataFieldService {
     apiEndpoint: ApiEndpoint,
     dataValue: any,
   ): void {
-    const existingMatch = `${dataSection}.${dataPath}`
+    const existingMatch = `${dataSection}${dataPath ? `.${dataPath}` : ""}`
     const dataType = getDataType(dataValue)
     if (!this.dataFields[existingMatch]) {
       const dataField = new DataField()
@@ -214,31 +214,33 @@ export class DataFieldService {
     this.dataFields = apiEndpoint.dataFields.reduce((obj, item) => {
       return {
         ...obj,
-        [`${item.dataSection}.${item.dataPath}`]: item,
+        [`${item.dataSection}${item.dataPath ? `.${item.dataPath}` : ""}`]: item,
       }
     }, {})
     this.updatedFields = {}
     this.findPathDataFields(apiTrace.path, apiEndpoint)
-    this.findPairObjectDataFields(
-      DataSection.REQUEST_QUERY,
-      apiTrace.requestParameters,
-      apiEndpoint,
-    )
-    this.findPairObjectDataFields(
-      DataSection.REQUEST_HEADER,
-      apiTrace.requestHeaders,
-      apiEndpoint,
-    )
+    if (apiTrace.responseStatus < 400) {
+      this.findPairObjectDataFields(
+        DataSection.REQUEST_QUERY,
+        apiTrace.requestParameters,
+        apiEndpoint,
+      )
+      this.findPairObjectDataFields(
+        DataSection.REQUEST_HEADER,
+        apiTrace.requestHeaders,
+        apiEndpoint,
+      )
+      this.findBodyDataFields(
+        DataSection.REQUEST_BODY,
+        apiTrace.requestBody,
+        apiEndpoint,
+      )
+    }
     this.findPairObjectDataFields(
       DataSection.RESPONSE_HEADER,
       apiTrace.responseHeaders,
       apiEndpoint,
     )
-    this.findBodyDataFields(
-      DataSection.REQUEST_BODY,
-      apiTrace.requestBody,
-      apiEndpoint,
-    )
     this.findBodyDataFields(
       DataSection.RESPONSE_BODY,
       apiTrace.responseBody,

diff --git a/backend/src/services/jobs/generate-endpoints-traces.ts b/backend/src/services/jobs/generate-endpoints-traces.ts
@@ -112,7 +112,7 @@ const generateEndpointsFromTraces = async (): Promise<void> => {
               }
             }
             if (pathRegex.length > 0) {
-              pathRegex = String.raw`^${pathRegex}$`
+              pathRegex = String.raw`^${pathRegex}(/)*$`
               const regexKey = `${trace.host}-${trace.method}-${pathRegex}`
               if (regexToTracesMap[regexKey]) {
                 regexToTracesMap[regexKey].traces.push(trace)

diff --git a/backend/src/services/jobs/generate-openapi-spec.ts b/backend/src/services/jobs/generate-openapi-spec.ts
@@ -120,21 +120,23 @@ const generateOpenApiSpec = async (): Promise<void> => {
               )
             }
           }
-          for (const requestParameter of requestParamters) {
-            const key = `${requestParameter.name}<>query`
-            parameters[key] = parseSchema(
-              parameters[key] ?? {},
-              parsedJsonNonNull(requestParameter.value, true),
-            )
-          }
-          for (const requestHeader of requestHeaders) {
-            const key = `${requestHeader.name}<>header`
-            parameters[key] = parseSchema(
-              parameters[key] ?? {},
-              parsedJsonNonNull(requestHeader.value, true),
-            )
-            if (requestHeader.name.toLowerCase() === "content-type") {
-              requestContentType = requestHeader.value.toLowerCase()
+          if (trace.responseStatus < 400) {
+            for (const requestParameter of requestParamters) {
+              const key = `${requestParameter.name}<>query`
+              parameters[key] = parseSchema(
+                parameters[key] ?? {},
+                parsedJsonNonNull(requestParameter.value, true),
+              )
+            }
+            for (const requestHeader of requestHeaders) {
+              const key = `${requestHeader.name}<>header`
+              parameters[key] = parseSchema(
+                parameters[key] ?? {},
+                parsedJsonNonNull(requestHeader.value, true),
+              )
+              if (requestHeader.name.toLowerCase() === "content-type") {
+                requestContentType = requestHeader.value.toLowerCase()
+              }
             }
           }
           for (const responseHeader of responseHeaders) {
@@ -155,7 +157,9 @@ const generateOpenApiSpec = async (): Promise<void> => {
             )
           }
 
-          parseContent(requestBodySpec, requestBody, requestContentType)
+          if (trace.responseStatus < 400) {
+            parseContent(requestBodySpec, requestBody, requestContentType)
+          }
           if (responseBody) {
             if (!responses[responseStatusString]?.content) {
               responses[responseStatusString] = {

diff --git a/backend/src/utils/index.ts b/backend/src/utils/index.ts
@@ -70,12 +70,13 @@ export const getRiskScore = (dataFields: DataField[]): RiskScore => {
   if (!dataFields) {
     return RiskScore.NONE
   }
-  let numRiskySensitiveDataClasses = 0
+  let uniqueSensitiveDataClasses = new Set<string>()
   for (const dataField of dataFields) {
     if (dataField.dataClasses) {
-      numRiskySensitiveDataClasses += dataField.dataClasses?.length ?? 0
+      dataField.dataClasses.forEach(e => uniqueSensitiveDataClasses.add(e))
     }
   }
+  const numRiskySensitiveDataClasses = uniqueSensitiveDataClasses.size
   switch (true) {
     case numRiskySensitiveDataClasses >= 3:
       return RiskScore.HIGH