Skip to content

Instantly share code, notes, and snippets.

@WalBeh
Created November 9, 2024 17:34
Show Gist options
  • Save WalBeh/32164bf26e168270878838644e247cd7 to your computer and use it in GitHub Desktop.
Save WalBeh/32164bf26e168270878838644e247cd7 to your computer and use it in GitHub Desktop.
trivy db updater for cache-dir
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