Add companion json file for each database installed

GitHubSecurityLab · Jan 24, 2024 · 1d38b62 · 1d38b62
1 parent 170b1d9
commit 1d38b62
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 34 deletions.
diff --git a/cmd/download.go b/cmd/download.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -83,14 +84,24 @@ func download() {
 		}
 
 		// get the commit this DB was created from
-		commitSha, primaryLanguage, err := utils.ExtractDBInfo(body)
+		metadata, err := utils.ExtractDBInfo(body)
 		if err != nil {
 			log.Fatal(err)
 		}
+		metadata["provenance"] = nwoFlag
+		commitSha := metadata["creationMetadata"].(map[string]interface{})["sha"].(string)
+		shortCommitSha := commitSha[:8]
+		primaryLanguage := metadata["primaryLanguage"].(string)
+		fmt.Println()
+		fmt.Println("Commit SHA:", commitSha)
+		fmt.Println("Short Commit SHA:", shortCommitSha)
+		fmt.Println("Primary language:", primaryLanguage)
 
-		filename := fmt.Sprintf("%s.zip", commitSha)
-		dir := filepath.Join(utils.GetPath(nwoFlag), primaryLanguage)
-		path := filepath.Join(dir, filename)
+		zipFilename := fmt.Sprintf("%s-%s.zip", primaryLanguage, shortCommitSha)
+		jsonFilename := fmt.Sprintf("%s-%s.json", primaryLanguage, shortCommitSha)
+		dir := utils.GetPath(nwoFlag)
+		zipPath := filepath.Join(dir, zipFilename)
+		jsonPath := filepath.Join(dir, jsonFilename)
 
 		// create directory if not exists
 		if _, err := os.Stat(dir); errors.Is(err, os.ErrNotExist) {
@@ -100,19 +111,34 @@ func download() {
 			}
 		}
 
-		// create file if not exists
-		if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
+		// create DB file if doesnot exists
+		if _, err := os.Stat(zipPath); errors.Is(err, os.ErrNotExist) {
 			// write the DB to disk
-			err = os.WriteFile(path, body, 0755)
+			err = os.WriteFile(zipPath, body, 0755)
 			if err != nil {
 				log.Fatal(err)
 			}
-			fmt.Printf("Writing DB to %s\n", path)
+			fmt.Printf("Writing DB to %s\n", zipPath)
 
 		} else {
-			fmt.Printf("Aborting, DB %s already exists\n", path)
+			fmt.Printf("Aborting, DB %s already exists\n", zipPath)
 		}
 
+		// create Metadata file if doesnot exists
+		if _, err := os.Stat(jsonPath); errors.Is(err, os.ErrNotExist) {
+			// Convert the map to JSON
+			jsonData, err := json.Marshal(metadata)
+			if err != nil {
+				log.Fatal(err)
+			}
+			// Write the JSON data to a file
+			err = os.WriteFile(jsonPath, jsonData, 0644)
+			if err != nil {
+				log.Fatal(err)
+			}
+		} else {
+			fmt.Printf("Aborting, DB metadata %s already exists\n", jsonPath)
+		}
 	}
 	fmt.Println("Done")
 }
diff --git a/cmd/install.go b/cmd/install.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -32,23 +33,24 @@ func init() {
 }
 
 func install(nwo string, dbPath string, remove bool) {
-	fmt.Printf("Installing '%s' DB for '%s'\n", dbPath, nwo)
+	fmt.Printf("Installing '%s' database for '%s'\n", dbPath, nwo)
 
 	// Check if the path exists
 	fileinfo, err := os.Stat(dbPath)
 	var zipPath string
 	if os.IsNotExist(err) {
-		log.Fatal(errors.New("DB path does not exist"))
+		log.Fatal(errors.New("Database path does not exist"))
 	}
 	if fileinfo.IsDir() {
-		fmt.Printf("Validating %s DB\n", dbPath)
+		fmt.Printf("Validating '%s' database\n", dbPath)
 		err := utils.ValidateDB(dbPath)
 		if err != nil {
-			fmt.Println("DB is not valid")
+			fmt.Println("Database is not valid")
+			return
 		}
 		// Compress DB
 		zipfilename := filepath.Join(os.TempDir(), "qldb.zip")
-		fmt.Println("Compressing DB to", zipfilename)
+		fmt.Println("Compressing database")
 		if err := utils.ZipDirectory(zipfilename, dbPath); err != nil {
 			log.Fatal(err)
 		}
@@ -57,7 +59,7 @@ func install(nwo string, dbPath string, remove bool) {
 	} else {
 		// Check if the file is a zip
 		if !strings.HasSuffix(dbPath, ".zip") {
-			log.Fatal(errors.New("DB path is not a zip file"))
+			log.Fatal(errors.New("Database is not a zip file"))
 		}
 
 		zipPath = dbPath
@@ -78,10 +80,10 @@ func install(nwo string, dbPath string, remove bool) {
 			// if there is one directory in the tmpdir, use that as the tmpdir
 			tmpdir = filepath.Join(tmpdir, dirEntries[0].Name())
 		}
-		fmt.Printf("Validating %s DB\n", tmpdir)
+		fmt.Printf("Validating '%s' database\n", tmpdir)
 		err = utils.ValidateDB(tmpdir)
 		if err != nil {
-			fmt.Println("DB is not valid")
+			fmt.Println("Database is not valid")
 		}
 	}
 
@@ -95,22 +97,35 @@ func install(nwo string, dbPath string, remove bool) {
 	if err != nil {
 		log.Fatal(err)
 	}
-	commitSha, primaryLanguage, err := utils.ExtractDBInfo(zipBytes)
+
+	metadata, err := utils.ExtractDBInfo(zipBytes)
+	if err != nil {
+		log.Fatal(err)
+	}
+	metadata["provenance"] = nwoFlag
+	commitSha := metadata["creationMetadata"].(map[string]interface{})["sha"].(string)
 	shortCommitSha := commitSha[:8]
+	primaryLanguage := metadata["primaryLanguage"].(string)
+	fmt.Println()
 	fmt.Println("Commit SHA:", commitSha)
 	fmt.Println("Short Commit SHA:", shortCommitSha)
 	fmt.Println("Primary language:", primaryLanguage)
 
+	zipFilename := fmt.Sprintf("%s-%s.zip", primaryLanguage, shortCommitSha)
+	jsonFilename := fmt.Sprintf("%s-%s.json", primaryLanguage, shortCommitSha)
+	dir := utils.GetPath(nwoFlag)
+
 	// Destination path
-	filename := fmt.Sprintf("%s-%s.zip", primaryLanguage, shortCommitSha)
-	destPath := filepath.Join(utils.GetPath(nwo), filename)
-	fmt.Println("Installing DB to", destPath)
+	zipDestPath := filepath.Join(dir, zipFilename)
+	jsonDestPath := filepath.Join(dir, jsonFilename)
+
+	fmt.Println("Installing database to '" + zipDestPath + "'")
 
 	// Check if the DB is already installed
-	if _, err := os.Stat(destPath); errors.Is(err, os.ErrNotExist) {
+	if _, err := os.Stat(zipDestPath); errors.Is(err, os.ErrNotExist) {
 
 		// Create the directory if it doesn't exist
-		err = os.MkdirAll(filepath.Dir(destPath), 0755)
+		err = os.MkdirAll(filepath.Dir(zipDestPath), 0755)
 		if err != nil {
 			log.Fatal(err)
 			return
@@ -124,7 +139,7 @@ func install(nwo string, dbPath string, remove bool) {
 		}
 		defer srcFile.Close()
 
-		destFile, err := os.Create(destPath)
+		destFile, err := os.Create(zipDestPath)
 		if err != nil {
 			log.Fatal(err)
 			return
@@ -142,11 +157,28 @@ func install(nwo string, dbPath string, remove bool) {
 			log.Fatal(err)
 		}
 	} else {
-		fmt.Println("DB already installed for same commit")
+		fmt.Println("Database already installed for same commit")
+	}
+
+	if _, err := os.Stat(jsonDestPath); errors.Is(err, os.ErrNotExist) {
+		// Convert the map to JSON
+		jsonData, err := json.Marshal(metadata)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		// Write the JSON data to a file
+		err = os.WriteFile(jsonDestPath, jsonData, 0644)
+		if err != nil {
+			log.Fatal(err)
+		}
+	} else {
+		fmt.Println("Database metadata already exists for same commit")
 	}
+
 	// Remove DB from the current location if -r flag is set
 	if remove {
-		fmt.Println("Removing DB from", dbPath)
+		fmt.Println("Removing database from '" + dbPath + "'")
 		if err := os.RemoveAll(dbPath); err != nil {
 			log.Fatal(err)
 		}

diff --git a/utils/utils.go b/utils/utils.go
@@ -48,12 +48,12 @@ func ValidateDB(dbPath string) error {
 	return nil
 }
 
-func ExtractDBInfo(body []byte) (string, string, error) {
+func ExtractDBInfo(body []byte) (map[string]interface{}, error) {
 	zipReader, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
 	if err != nil {
 		log.Fatal(err)
 	}
-	fmt.Print("Extracting commit version from DB ... ")
+	fmt.Print("Extracting database information ... ")
 	for _, zf := range zipReader.File {
 		if strings.HasSuffix(zf.Name, "codeql-database.yml") {
 			f, err := zf.Open()
@@ -70,14 +70,10 @@ func ExtractDBInfo(body []byte) (string, string, error) {
 			if err != nil {
 				log.Fatal(err)
 			}
-			commitSha := dbData["creationMetadata"].(map[string]interface{})["sha"].(string)
-			fmt.Printf("%s\n", commitSha)
-			// extractionDate := dbData["creationMetadata"].(map[string]interface{})["creationTime"].(string)
-			primaryLanguage := dbData["primaryLanguage"].(string)
-			return commitSha, primaryLanguage, nil
+			return dbData, nil
 		}
 	}
-	return "", "", errors.New("codeql-database.yml not found")
+	return nil, errors.New("codeql-database.yml not found")
 }
 
 // Unzip will decompress a zip archive, moving all files and folders