Implemented Artifactory upload, download and delete handling

- Added GRABIT_ARTIFACTORY_TOKEN support
- Implemented cache download priority
- Added fallback to source URL
- Added validation at each step
- Improved error handling and logging
- Integrated cache upload functionality

Co-authored-by: Amin Assaid <[email protected]>
Co-authored-by: Amy Druham <[email protected]>

cmd/add.go

@@ -4,6 +4,9 @@
 package cmd
 
 import (
+	"fmt"
+	"os"
+
 	"github.com/cisco-open/grabit/internal"
 	"github.com/spf13/cobra"
 )
@@ -18,6 +21,7 @@ func addAdd(cmd *cobra.Command) {
 	addCmd.Flags().String("algo", internal.RecommendedAlgo, "Integrity algorithm")
 	addCmd.Flags().String("filename", "", "Target file name to use when downloading the resource")
 	addCmd.Flags().StringArray("tag", []string{}, "Resource tags")
+	addCmd.Flags().String("cache", "", "Artifactory cache URL")
 	cmd.AddCommand(addCmd)
 }
 
@@ -26,6 +30,18 @@ func runAdd(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	// Get cache URL
+	cacheURL, err := cmd.Flags().GetString("cache")
+	if err != nil {
+		return err
+	}
+	// Check token if cache is requested
+	if cacheURL != "" {
+		token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN")
+		if token == "" {
+			return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set")
+		}
+	}
 	lock, err := internal.NewLock(lockFile, true)
 	if err != nil {
 		return err
@@ -42,7 +58,7 @@ func runAdd(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	err = lock.AddResource(args, algo, tags, filename)
+	err = lock.AddResource(args, algo, tags, filename, cacheURL)
 	if err != nil {
 		return err
 	}

cmd/add_test.go

@@ -10,6 +10,10 @@ import (
 )
 
 func TestRunAdd(t *testing.T) {
+	// Set the GRABIT_ARTIFACTORY_TOKEN environment variable.
+	t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "test-token")
+
+	// Setup HTTP handler for the resource
 	handler := func(w http.ResponseWriter, r *http.Request) {
 		_, err := w.Write([]byte(`abcdef`))
 		if err != nil {
@@ -18,8 +22,29 @@ func TestRunAdd(t *testing.T) {
 	}
 	port, server := test.HttpHandler(handler)
 	defer server.Close()
+
+	// Setup dummy cache server
+	cacheHandler := func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == "PUT" {
+			w.WriteHeader(http.StatusCreated)
+		}
+	}
+	cachePort, cacheServer := test.HttpHandler(cacheHandler)
+	defer cacheServer.Close()
+
+	// Create empty lockfile
+	lockFile := test.TmpFile(t, "")
+
 	cmd := NewRootCmd()
-	cmd.SetArgs([]string{"-f", test.TmpFile(t, ""), "add", fmt.Sprintf("http://localhost:%d/test.html", port)})
+	// Add cache URL to the command
+	cacheURL := fmt.Sprintf("http://localhost:%d", cachePort)
+	cmd.SetArgs([]string{
+		"-f", lockFile,
+		"add",
+		fmt.Sprintf("http://localhost:%d/test.html", port),
+		"--cache", cacheURL,
+	})
+
 	err := cmd.Execute()
 	assert.Nil(t, err)
 }

cmd/delete_test.go

@@ -15,6 +15,7 @@ func TestRunDelete(t *testing.T) {
 	Tags = ['tag1', 'tag2']
 `)
 	cmd := NewRootCmd()
+	cmd.Flags().String("cache", "", "Artifactory URL for caching")
 	cmd.SetArgs([]string{"-f", testfilepath, "delete", "http://localhost:123456/test.html"})
 	err := cmd.Execute()
 	assert.Nil(t, err)

internal/artifactory_cache_test.go

@@ -0,0 +1,91 @@
+package internal
+
+import (
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/cisco-open/grabit/test"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestAddWithArtifactoryCache verifies adding a resource with caching enabled.
+func TestAddWithArtifactoryCacheTokenNotSet(t *testing.T) {
+	// Setup a simple HTTP handler that always returns "test content".
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`test content`))
+	}
+	// Start the HTTP server and get the port it runs on.
+	port, server := test.HttpHandler(handler)
+	defer server.Close() // Ensure the server is stopped after the test.
+
+	// Create a temporary lock file for testing.
+	path := test.TmpFile(t, "")
+	lock, err := NewLock(path, true)
+	require.NoError(t, err) // Fail the test if the lock cannot be created.
+
+	// Set up URLs for the source file and cache.
+	sourceURL := fmt.Sprintf("http://localhost:%d/test.txt", port)
+	cacheURL := fmt.Sprintf("http://localhost:%d", port)
+
+	// Attempt to add a resource to the lock file.
+	err = lock.AddResource([]string{sourceURL}, "sha256", []string{}, "", cacheURL)
+	// Verify that the error message indicates the token is not set.
+	assert.Contains(t, err.Error(), "GRABIT_ARTIFACTORY_TOKEN environment variable is not set")
+}
+
+// TestDeleteWithArtifactoryCache verifies deleting a resource with caching enabled.
+func TestDeleteWithArtifactoryCache(t *testing.T) {
+	// Set the GRABIT_ARTIFACTORY_TOKEN environment variable.
+	t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "test-token")
+
+	// Setup an HTTP handler to handle DELETE requests.
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == "DELETE" { // Respond with OK for DELETE requests.
+			w.WriteHeader(http.StatusOK)
+		}
+	}
+	// Start the HTTP server and get the port it runs on.
+	port, server := test.HttpHandler(handler)
+	defer server.Close() // Ensure the server is stopped after the test.
+
+	// Set up URLs for the source file and cache.
+	sourceURL := fmt.Sprintf("http://localhost:%d/test.txt", port)
+	cacheURL := fmt.Sprintf("http://localhost:%d", port)
+
+	// Create a lock file with the resource and cache information.
+	lockContent := fmt.Sprintf(`[[Resource]]
+        Urls = ['%s']
+        Integrity = 'sha256-test'
+        CacheUri = '%s'`, sourceURL, cacheURL)
+
+	lockPath := test.TmpFile(t, lockContent)
+	lock, err := NewLock(lockPath, false)
+	require.NoError(t, err) // Fail the test if the lock cannot be created.
+
+	// Save the lock file before modifying it.
+	err = lock.Save()
+	require.NoError(t, err) // Fail the test if saving fails.
+
+	// Delete the resource from the lock file.
+	lock.DeleteResource(sourceURL)
+
+	// Checks that http.StatusOK was returned meaning the delete was successful
+	resp, err := http.Get(server.URL)
+	if err != nil {
+		require.NoError(t, err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("expected 200 but got %d", resp.StatusCode)
+	}
+
+	// Save the lock file again after deletion.
+	err = lock.Save()
+	require.NoError(t, err) // Fail the test if saving fails.
+
+	// Reload the lock file and verify the resource is gone.
+	newLock, err := NewLock(lockPath, false)
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(newLock.conf.Resource)) // Ensure no resources remain.
+}

internal/lock.go

@@ -8,14 +8,26 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/http"
 	"os"
+	"path"
+	"path/filepath"
 	"strconv"
+	"strings"
 
+	"github.com/carlmjohnson/requests"
 	toml "github.com/pelletier/go-toml/v2"
+	"github.com/rs/zerolog/log"
 )
 
 var COMMENT_PREFIX = "//"
 
+// getArtifactoryURL constructs the URL for an artifact in Artifactory
+func getArtifactoryURL(baseURL, integrity string) string {
+	return fmt.Sprintf("%s/%s", baseURL, integrity)
+
+}
+
 // Lock represents a grabit lockfile.
 type Lock struct {
 	path string
@@ -49,30 +61,99 @@ func NewLock(path string, newOk bool) (*Lock, error) {
 	return &Lock{path: path, conf: conf}, nil
 }
 
-func (l *Lock) AddResource(paths []string, algo string, tags []string, filename string) error {
+func (l *Lock) AddResource(paths []string, algo string, tags []string, filename string, cacheURL string) error {
 	for _, u := range paths {
 		if l.Contains(u) {
 			return fmt.Errorf("resource '%s' is already present", u)
 		}
 	}
-	r, err := NewResourceFromUrl(paths, algo, tags, filename)
+	r, err := NewResourceFromUrl(paths, algo, tags, filename, cacheURL)
 	if err != nil {
 		return err
 	}
+
+	// If cache URL is provided, handles Artifactory upload
+	if cacheURL != "" {
+		token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN")
+		if token == "" {
+			return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set")
+		}
+
+		// Add context here
+		ctx := context.Background()
+		path, err := GetUrltoTempFile(paths[0], token, ctx)
+		if err != nil {
+			return fmt.Errorf("failed to get file for cache: %s", err)
+		}
+		defer os.Remove(path)
+
+		// Upload to Artifactory using hash as filename
+		err = uploadToArtifactory(path, cacheURL, r.Integrity)
+		if err != nil {
+			return fmt.Errorf("failed to upload to cache: %v", err)
+		}
+	}
+
 	l.conf.Resource = append(l.conf.Resource, *r)
 	return nil
 }
 
+func uploadToArtifactory(filePath, cacheURL, integrity string) error {
+	token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN")
+	if token == "" {
+		return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set")
+	}
+
+	// Use the integrity value directly for the URL
+	artifactoryURL := getArtifactoryURL(cacheURL, integrity)
+
+	// Upload the file using the requests package
+	err := requests.
+		URL(artifactoryURL).
+		Method(http.MethodPut).
+		Header("Authorization", fmt.Sprintf("Bearer %s", token)).
+		BodyFile(filePath). // Using BodyFile instead of ReadFile
+		Fetch(context.Background())
+
+	if err != nil {
+		return fmt.Errorf("upload failed: %w", err)
+	}
+
+	return nil
+}
 func (l *Lock) DeleteResource(path string) {
 	newStatements := []Resource{}
 	for _, r := range l.conf.Resource {
 		if !r.Contains(path) {
 			newStatements = append(newStatements, r)
+		} else if r.Contains(path) && r.CacheUri != "" {
+			token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN")
+			if token == "" {
+				log.Warn().Msg("Warning: Unable to delete from Artifactory: GRABIT_ARTIFACTORY_TOKEN not set.")
+
+				continue
+			}
+
+			artifactoryURL := getArtifactoryURL(r.CacheUri, r.Integrity)
+
+			err := deleteCache(artifactoryURL, token)
+			if err != nil {
+				log.Warn().Msg("Warning: Unable to delete from Artifactory")
+			}
 		}
 	}
 	l.conf.Resource = newStatements
 }
 
+func deleteCache(url, token string) error {
+	// Create and send a DELETE request with an Authorization header.
+	return requests.
+		URL(url).
+		Method(http.MethodDelete).
+		Header("Authorization", fmt.Sprintf("Bearer %s", token)).
+		Fetch(context.Background())
+}
+
 const NoFileMode = os.FileMode(0)
 
 // strToFileMode converts a string to a os.FileMode.
@@ -165,6 +246,47 @@ func (l *Lock) Download(dir string, tags []string, notags []string, perm string,
 	for i, r := range filteredResources {
 		resource := r
 		go func() {
+			// Try Artifactory first if available
+			if resource.CacheUri != "" {
+				token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN")
+				if token != "" {
+					artifactoryURL := getArtifactoryURL(resource.CacheUri, resource.Integrity)
+					filename := resource.Filename
+					if filename == "" {
+						filename = path.Base(resource.Urls[0])
+					}
+					fullPath := filepath.Join(dir, filename)
+
+					// Use getUrl with bearer token
+					tmpPath, err := getUrl(artifactoryURL, fullPath, token, ctx)
+					if err == nil {
+						//  integrity check
+						algo, err := getAlgoFromIntegrity(resource.Integrity)
+						if err != nil {
+							errorCh <- err
+							return
+						}
+						err = checkIntegrityFromFile(tmpPath, algo, resource.Integrity, artifactoryURL)
+						if err != nil {
+							errorCh <- err
+							return
+						}
+						if mode != NoFileMode {
+							err = os.Chmod(tmpPath, mode.Perm())
+						}
+						if err == nil {
+							errorCh <- nil
+							if statusLine != nil {
+								statusLine.Increment(i)
+							}
+							return
+						}
+					}
+					if strings.Contains(err.Error(), "lookup invalid") || strings.Contains(err.Error(), "dial tcp") {
+						fmt.Printf("Failed to download from Artifactory, falling back to original URL: %v\n", err)
+					}
+				}
+			}
 
 			err := resource.Download(dir, mode, ctx)
 			errorCh <- err

internal/lock_test.go

@@ -51,7 +51,7 @@ func TestLockManipulations(t *testing.T) {
 	port, server := test.HttpHandler(handler)
 	defer server.Close()
 	resource := fmt.Sprintf("http://localhost:%d/test2.html", port)
-	err = lock.AddResource([]string{resource}, "sha512", []string{}, "")
+	err = lock.AddResource([]string{resource}, "sha512", []string{}, "", "")
 	assert.Nil(t, err)
 	assert.Equal(t, 2, len(lock.conf.Resource))
 	err = lock.Save()
@@ -68,7 +68,7 @@ func TestDuplicateResource(t *testing.T) {
 		Integrity = 'sha256-asdasdasd'`, url))
 	lock, err := NewLock(path, false)
 	assert.Nil(t, err)
-	err = lock.AddResource([]string{url}, "sha512", []string{}, "")
+	err = lock.AddResource([]string{url}, "sha512", []string{}, "", "")
 	assert.NotNil(t, err)
 	assert.Contains(t, err.Error(), "already present")
 }