add extension, fix spec showing

backend/models/openapi-spec.ts

@@ -1,4 +1,5 @@
 import { BaseEntity, Entity, PrimaryColumn, Column } from "typeorm";
+import { SpecExtension } from "../src/enums";
 
 @Entity()
 export class OpenApiSpec extends BaseEntity {
@@ -8,6 +9,9 @@ export class OpenApiSpec extends BaseEntity {
   @Column()
   spec: string
 
+  @Column({ type: "enum", enum: SpecExtension, default: SpecExtension.JSON})
+  extension: SpecExtension
+
   @Column({ type: "bool", default: false })
   isAutoGenerated: boolean
 }

backend/src/api/spec/index.ts

@@ -6,6 +6,7 @@ import ApiResponseHandler from "../../api-response-handler";
 import Error400BadRequest from "../../errors/error-400-bad-request";
 import { AppDataSource } from "../../data-source";
 import { OpenApiSpec } from "../../../models";
+import { SpecExtension } from "../../enums";
 
 export const getSpecListHandler = async (req: Request, res: Response) => {
   try {
@@ -26,6 +27,16 @@ export const uploadNewSpecHandler = async (req: Request, res: Response) => {
       ?.split(".json")[0]
       ?.split(".yaml")[0]
       ?.split(".yml")[0];
+    let extension = undefined;
+    if (specFile.mimetype === "application/json") {
+      extension = SpecExtension.JSON;
+    } else if (specFile.mimetype === "text/yaml") {
+      extension = SpecExtension.YAML;
+    } else {
+      throw new Error400BadRequest(
+        "Only .json, .yaml, and .yml format allowed."
+      );
+    }
     if (!fileName) {
       throw new Error400BadRequest("No filename provided.");
     }
@@ -39,7 +50,12 @@ export const uploadNewSpecHandler = async (req: Request, res: Response) => {
     const specObject: JSONValue = yaml.load(
       specFile.buffer.toString()
     ) as JSONValue;
-    await SpecService.uploadNewSpec(specObject, fileName);
+    await SpecService.uploadNewSpec(
+      specObject,
+      fileName,
+      extension,
+      specFile.buffer.toString()
+    );
     await ApiResponseHandler.success(res, null);
   } catch (err) {
     await ApiResponseHandler.error(res, err);
@@ -63,13 +79,35 @@ export const updateSpecHandler = async (req: Request, res: Response) => {
       ?.split(".json")[0]
       ?.split(".yaml")[0]
       ?.split(".yml")[0];
+    let extension = undefined;
+    if (specFile.mimetype === "application/json") {
+      extension = SpecExtension.JSON;
+    } else if (specFile.mimetype === "text/yaml") {
+      extension = SpecExtension.YAML;
+    } else {
+      throw new Error400BadRequest(
+        "Only .json, .yaml, and .yml format allowed."
+      );
+    }
     if (!fileName) {
       throw new Error400BadRequest("No filename provided.");
     }
+    const openApiSpecRepository = AppDataSource.getRepository(OpenApiSpec);
+    const exisitingSpec = await openApiSpecRepository.findOneBy({
+      name: fileName,
+    });
+    if (!exisitingSpec) {
+      throw new Error400BadRequest("Spec file does not exist, cannot update.");
+    }
     const specObject: JSONValue = yaml.load(
       specFile.buffer.toString()
     ) as JSONValue;
-    await SpecService.updateSpec(specObject, fileName);
+    await SpecService.updateSpec(
+      specObject,
+      fileName,
+      extension,
+      specFile.buffer.toString()
+    );
     await ApiResponseHandler.success(res, null);
   } catch (err) {
     await ApiResponseHandler.error(res, err);

backend/src/enums.ts

@@ -28,3 +28,8 @@ export enum RiskScore {
   MEDIUM = "medium",
   HIGH = "high",
 }
+
+export enum SpecExtension {
+  JSON = "json",
+  YAML = "yaml",
+}

backend/src/services/endpoints/index.ts

@@ -2,7 +2,7 @@ import { IsNull, Raw } from "typeorm";
 import { isSuspectedParamater } from "../../utils";
 import { ApiEndpoint, ApiTrace, OpenApiSpec } from "../../../models";
 import { AppDataSource } from "../../data-source";
-import { RestMethod } from "../../enums";
+import { RestMethod, SpecExtension } from "../../enums";
 import { ScannerService } from "../scanner/scan";
 
 interface GenerateEndpoint {
@@ -218,6 +218,7 @@ export class EndpointsService {
         endpoint.openapiSpec = spec;
       }
       spec.spec = JSON.stringify(openApiSpec, null, 2);
+      spec.extension = SpecExtension.JSON;
       await openApiSpecRepository.save(spec);
       await apiEndpointRepository.save(endpoints);
     });

backend/src/services/spec/index.ts

@@ -1,6 +1,6 @@
 import { v4 as uuidv4 } from "uuid";
 import { Not } from "typeorm";
-import { RestMethod } from "../../enums";
+import { RestMethod, SpecExtension } from "../../enums";
 import {
   ApiEndpoint,
   ApiTrace,
@@ -20,9 +20,14 @@ export class SpecService {
     return specList;
   }
 
-  static async updateSpec(specObject: JSONValue, fileName: string) {
+  static async updateSpec(
+    specObject: JSONValue,
+    fileName: string,
+    extension: SpecExtension,
+    specString: string
+  ) {
     await this.deleteSpec(fileName);
-    await this.uploadNewSpec(specObject, fileName);
+    await this.uploadNewSpec(specObject, fileName, extension, specString);
   }
 
   static async deleteSpec(fileName: string) {
@@ -48,7 +53,12 @@ export class SpecService {
     await openApiSpecRepository.remove(openApiSpec);
   }
 
-  static async uploadNewSpec(specObject: JSONValue, fileName: string) {
+  static async uploadNewSpec(
+    specObject: JSONValue,
+    fileName: string,
+    extension: SpecExtension,
+    specString: string
+  ) {
     const servers: any[] = specObject["servers"];
     const paths: JSONValue = specObject["paths"];
 
@@ -67,8 +77,9 @@ export class SpecService {
     if (!existingSpec) {
       existingSpec = new OpenApiSpec();
       existingSpec.name = fileName;
+      existingSpec.extension = extension;
     }
-    existingSpec.spec = JSON.stringify(specObject);
+    existingSpec.spec = specString;
     const pathKeys = Object.keys(paths);
     const endpoints: {
       similarEndpoints: ApiEndpoint[];

common/src/enums.ts

@@ -39,3 +39,8 @@ export enum ConnectionType {
   AWS = "AWS",
   GCP = "GCP",
 }
+
+export enum SpecExtension {
+  JSON = "json",
+  YAML = "yaml",
+}

common/src/types.ts

@@ -1,4 +1,4 @@
-import { AlertType, ConnectionType, RestMethod, RiskScore } from "./enums";
+import { AlertType, ConnectionType, RestMethod, RiskScore, SpecExtension } from "./enums";
 
 export interface Meta {
   incoming: boolean;
@@ -127,6 +127,7 @@ export interface OpenApiSpec {
   name: string;
   spec: string;
   isAutoGenerated: boolean;
+  extension: SpecExtension;
 }
 
 export interface Connection {

frontend/src/components/Endpoint/Overview.tsx

@@ -84,7 +84,7 @@ const EndpointOverview: React.FC<EndpointOverviewProps> = React.memo(
           overflowY={{ base: "unset", lg: "scroll" }}
           h={{ base: "unset", lg: "full" }}
         >
-          <Highlight {...defaultProps} theme={theme} code={openAPISpec} language="yaml">
+          <Highlight {...defaultProps} theme={theme} code={endpoint?.openapiSpec?.spec || "No spec."} language={endpoint?.openapiSpec?.extension || "yaml"}>
             {({ className, style, tokens, getLineProps, getTokenProps }) => (
               <pre
                 className={className}

frontend/src/pages/endpoint/[endpointUUID].tsx

@@ -1,5 +1,6 @@
 import { useEffect, useState } from "react";
 import { useRouter } from "next/router";
+import ErrorPage from "next/error";
 import { SideNavLinkDestination } from "../../components/Sidebar/NavLinkUtils";
 import { SidebarLayoutShell } from "../../components/SidebarLayoutShell";
 import EndpointPage from "../../components/Endpoint";
@@ -11,16 +12,23 @@ const Endpoint = () => {
   const { endpointUUID } = router.query;
   const [fetching, setFetching] = useState<boolean>(true);
   const [endpoint, setEndpoint] = useState<ApiEndpointDetailed>();
+  const [notFound, setNotFound] = useState<boolean>(false);
   useEffect(() => {
     const fetchEndpoint = async () => {
       const res = await getEndpoint(endpointUUID?.toString());
+      if (!res) {
+        setNotFound(true);
+      }
       setEndpoint(res);
       setFetching(false);
     }
     if (endpointUUID) {
       fetchEndpoint();
     }
   }, [endpointUUID])
+  if (notFound) {
+    return <ErrorPage statusCode={404} />
+  }
   return (
     <SidebarLayoutShell currentTab={SideNavLinkDestination.Endpoints}>
       <EndpointPage endpoint={endpoint} fetching={fetching} />