Created
November 9, 2024 17:34
-
-
Save WalBeh/32164bf26e168270878838644e247cd7 to your computer and use it in GitHub Desktop.
trivy db updater for cache-dir
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"bytes" | |
"encoding/json" | |
"flag" | |
"fmt" | |
"io" | |
"os" | |
"os/exec" | |
"path/filepath" | |
"time" | |
) | |
type MetadataJSON struct { | |
Version int `json:"Version"` | |
NextUpdate time.Time `json:"NextUpdate"` | |
UpdatedAt time.Time `json:"UpdatedAt"` | |
DownloadedAt time.Time `json:"DownloadedAt"` | |
} | |
func readMetadata(metadataPath string) (*MetadataJSON, error) { | |
file, err := os.ReadFile(metadataPath) | |
if err != nil { | |
return nil, fmt.Errorf("failed to read metadata file: %v", err) | |
} | |
var metadata MetadataJSON | |
if err := json.Unmarshal(file, &metadata); err != nil { | |
return nil, fmt.Errorf("failed to parse metadata JSON: %v", err) | |
} | |
return &metadata, nil | |
} | |
// copyDir recursively copies a directory tree | |
func copyDir(src string, dst string) error { | |
if err := os.MkdirAll(dst, 0755); err != nil { | |
return err | |
} | |
entries, err := os.ReadDir(src) | |
if err != nil { | |
return err | |
} | |
for _, entry := range entries { | |
srcPath := filepath.Join(src, entry.Name()) | |
dstPath := filepath.Join(dst, entry.Name()) | |
if entry.IsDir() { | |
if err := copyDir(srcPath, dstPath); err != nil { | |
return err | |
} | |
} else { | |
srcFile, err := os.Open(srcPath) | |
if err != nil { | |
return err | |
} | |
defer srcFile.Close() | |
dstFile, err := os.Create(dstPath) | |
if err != nil { | |
return err | |
} | |
defer dstFile.Close() | |
if _, err := io.Copy(dstFile, srcFile); err != nil { | |
return err | |
} | |
} | |
} | |
return nil | |
} | |
func backupTrivyDB(cacheDir string) error { | |
backupDir := "/tmp/trivy_save" | |
os.RemoveAll(backupDir) | |
if err := os.MkdirAll(backupDir, 0755); err != nil { | |
return fmt.Errorf("failed to create backup directory: %v", err) | |
} | |
if err := copyDir(cacheDir, backupDir); err != nil { | |
return fmt.Errorf("failed to copy directory: %v", err) | |
} | |
fmt.Printf("Successfully backed up %s to %s\n", cacheDir, backupDir) | |
return nil | |
} | |
func updateMetadataNextUpdate(metadataPath string) error { | |
metadata, err := readMetadata(metadataPath) | |
if err != nil { | |
return err | |
} | |
metadata.NextUpdate = time.Now().Add(4 * time.Hour) | |
updatedData, err := json.Marshal(metadata) | |
if err != nil { | |
return fmt.Errorf("failed to marshal updated metadata: %v", err) | |
} | |
if err := os.WriteFile(metadataPath, updatedData, 0644); err != nil { | |
return fmt.Errorf("failed to write updated metadata: %v", err) | |
} | |
fmt.Printf("Delay NextUpdate timestamp to: %s\n", metadata.NextUpdate.Format(time.RFC3339)) | |
return nil | |
} | |
func restoreTrivyDB(cacheDir string) error { | |
backupDir := "/tmp/trivy_save" | |
if err := os.RemoveAll(cacheDir); err != nil { | |
return fmt.Errorf("failed to remove current directory: %v", err) | |
} | |
if err := copyDir(backupDir, cacheDir); err != nil { | |
return fmt.Errorf("failed to restore from backup: %v", err) | |
} | |
metadataPath := filepath.Join(cacheDir, "db", "metadata.json") | |
if err := updateMetadataNextUpdate(metadataPath); err != nil { | |
return fmt.Errorf("failed to update metadata timestamp: %v", err) | |
} | |
fmt.Printf("Successfully restored %s from %s\n", cacheDir, backupDir) | |
return nil | |
} | |
func runTrivyUpdateCommand(cacheDir string) error { | |
cmd := exec.Command("trivy", "image", "--cache-dir", cacheDir, "--download-db-only") | |
var stderr bytes.Buffer | |
cmd.Stderr = &stderr | |
if err := cmd.Run(); err != nil { | |
return fmt.Errorf("trivy command failed: %v\nError output:\n%s", err, stderr.String()) | |
} | |
return nil | |
} | |
func main() { | |
cacheDir := flag.String("cache-dir", "/tmp/trivy", "Directory to store Trivy cache") | |
flag.Parse() | |
metadataFile := fmt.Sprintf("%s/db/metadata.json", *cacheDir) | |
metadata, err := readMetadata(metadataFile) | |
if err != nil { | |
fmt.Println("Error reading metadata.json:", err) | |
fmt.Println("Executing Trivy DB download...") | |
if err := backupTrivyDB(*cacheDir); err != nil { | |
fmt.Println("Error backing up Trivy DB:", err) | |
return | |
} | |
if err := runTrivyUpdateCommand(*cacheDir); err != nil { | |
fmt.Println(err) | |
fmt.Println("Update failed, restoring from backup...") | |
if restoreErr := restoreTrivyDB(*cacheDir); restoreErr != nil { | |
fmt.Printf("Error restoring from backup: %v\n", restoreErr) | |
} | |
return | |
} | |
// Read the new metadata after update | |
if newMetadata, err := readMetadata(metadataFile); err == nil { | |
fmt.Printf("Next Trivy DB update will happen at: %s\n", newMetadata.NextUpdate.Format(time.RFC3339)) | |
} | |
return | |
} | |
if time.Now().After(metadata.NextUpdate) { | |
fmt.Println("Updating Trivy DB...") | |
if err := backupTrivyDB(*cacheDir); err != nil { | |
fmt.Println("Error backing up Trivy DB:", err) | |
return | |
} | |
if err := runTrivyUpdateCommand(*cacheDir); err != nil { | |
fmt.Println(err) | |
fmt.Println("Update failed, restoring from backup...") | |
if restoreErr := restoreTrivyDB(*cacheDir); restoreErr != nil { | |
fmt.Printf("Error restoring from backup: %v\n", restoreErr) | |
} | |
return | |
} | |
fmt.Println("Trivy DB update complete.") | |
// Read the new metadata after update | |
if newMetadata, err := readMetadata(metadataFile); err == nil { | |
fmt.Printf("Next Trivy DB update will happen at: %s\n", newMetadata.NextUpdate.Format(time.RFC3339)) | |
} | |
} else { | |
fmt.Println("Trivy DB is up-to-date. No update needed.") | |
fmt.Printf("Next Trivy DB update will happen at: %s\n", metadata.NextUpdate.Format(time.RFC3339)) | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment