image.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	log "github.com/sirupsen/logrus"
)

const sha256Prefix = "sha256:"

// authenticateClient authenticates with the registry and returns the client
func authenticateClient(ref ImageReference) (*RegistryClient, error) {
	client := NewRegistryClient()

	log.WithField("registry", ref.Registry).Info("Authenticating with registry")
	if err := client.Authenticate(ref); err != nil {
		return nil, fmt.Errorf("authentication failed: %w", err)
	}
	log.WithField("user", client.GetAuthenticatedUser()).Info("Authenticated successfully")

	return client, nil
}

// fetchManifest retrieves the manifest for the image for the given platform
func fetchManifest(client *RegistryClient, ref ImageReference, platform Platform) (*ManifestV2, error) {
	log.WithFields(log.Fields{
		"repository": ref.Repository,
		"tag":        ref.Tag,
		"platform":   platform,
	}).Info("Fetching manifest")
	manifest, err := client.getManifest(ref, platform)
	if err != nil {
		return nil, fmt.Errorf("failed to get manifest: %w", err)
	}
	return manifest, nil
}

// downloadImageConfig downloads and parses the image configuration
func downloadImageConfig(client *RegistryClient, ref ImageReference, manifest *ManifestV2, tempDir string) (*ImageConfig, string, error) {
	log.Info("Downloading image config")
	configDigest := strings.TrimPrefix(manifest.Config.Digest, sha256Prefix)
	configPath := filepath.Join(tempDir, configDigest+".json")
	if err := client.DownloadBlob(ref, manifest.Config.Digest, configPath); err != nil {
		return nil, "", fmt.Errorf("failed to download config: %w", err)
	}

	configData, err := os.ReadFile(configPath)
	if err != nil {
		return nil, "", err
	}

	var imageConfig ImageConfig
	if err := json.Unmarshal(configData, &imageConfig); err != nil {
		return nil, "", err
	}

	return &imageConfig, configDigest, nil
}

// downloadAndProcessLayer downloads a single layer and creates its metadata files
func downloadAndProcessLayer(client *RegistryClient, ref ImageReference, layerDigestFull string, index int, totalLayers int, imageConfig *ImageConfig, tempDir string) (string, error) {
	log.WithFields(log.Fields{
		"layer_index":  index + 1,
		"total_layers": totalLayers,
		"digest":       layerDigestFull[:19] + "...",
	}).Info("Downloading layer")
	layerDigest := strings.TrimPrefix(layerDigestFull, sha256Prefix)

	compressedPath := filepath.Join(tempDir, layerDigest+".tar.gz")
	if err := client.DownloadBlob(ref, layerDigestFull, compressedPath); err != nil {
		return "", fmt.Errorf("failed to download layer: %w", err)
	}

	diffID := strings.TrimPrefix(imageConfig.RootFS.DiffIDs[index], sha256Prefix)
	layerDir := filepath.Join(tempDir, diffID)
	if err := os.MkdirAll(layerDir, 0755); err != nil {
		return "", err
	}

	layerTarPath := filepath.Join(layerDir, "layer.tar")
	if err := decompressGzip(compressedPath, layerTarPath); err != nil {
		return "", fmt.Errorf("failed to decompress layer: %w", err)
	}

	if err := createLayerMetadata(layerDir, diffID, index, imageConfig); err != nil {
		return "", err
	}

	return diffID, nil
}

// createLayerMetadata creates VERSION and json files for a layer
func createLayerMetadata(layerDir, diffID string, index int, imageConfig *ImageConfig) error {
	if err := os.WriteFile(filepath.Join(layerDir, "VERSION"), []byte("1.0"), 0644); err != nil {
		return err
	}

	layerJSON := map[string]interface{}{
		"id":      diffID,
		"created": "0001-01-01T00:00:00Z",
	}
	if index > 0 {
		prevDiffID := strings.TrimPrefix(imageConfig.RootFS.DiffIDs[index-1], sha256Prefix)
		layerJSON["parent"] = prevDiffID
	}

	return marshalJSONToFile(layerJSON, layerDir, "json")
}

// downloadAllLayers downloads all layers and returns their diff IDs
func downloadAllLayers(client *RegistryClient, ref ImageReference, manifest *ManifestV2, imageConfig *ImageConfig, tempDir string) ([]string, error) {
	layerPaths := make([]string, len(manifest.Layers))

	for i, layer := range manifest.Layers {
		diffID, err := downloadAndProcessLayer(client, ref, layer.Digest, i, len(manifest.Layers), imageConfig, tempDir)
		if err != nil {
			return nil, err
		}
		layerPaths[i] = diffID
	}

	return layerPaths, nil
}

// createDockerManifest creates the manifest.json file for docker load
func createDockerManifest(ref ImageReference, configDigest string, layerPaths []string, tempDir string) error {
	repoTag := ref.Repository + ":" + ref.Tag
	if ref.Registry != "registry-1.docker.io" {
		repoTag = ref.Registry + "/" + repoTag
	}

	layers := make([]string, len(layerPaths))
	for i, p := range layerPaths {
		layers[i] = p + "/layer.tar"
	}

	manifestJSON := []map[string]interface{}{
		{
			"Config":   configDigest + ".json",
			"RepoTags": []string{repoTag},
			"Layers":   layers,
		},
	}

	return marshalJSONToFile(manifestJSON, tempDir, "manifest.json")
}

// createRepositoriesFile creates the repositories file for docker load
func createRepositoriesFile(ref ImageReference, layerPaths []string, tempDir string) error {
	imageName := filepath.Base(ref.Repository)
	topLayer := layerPaths[len(layerPaths)-1]

	repositories := map[string]map[string]string{
		imageName: {ref.Tag: topLayer},
	}

	return marshalJSONToFile(repositories, tempDir, "repositories")
}

// createOutputTar creates the final tar archive
func createOutputTar(ref ImageReference, tempDir, outputDir string, platform Platform) (string, error) {
	if err := os.MkdirAll(outputDir, 0755); err != nil {
		return "", err
	}

	outputPath := filepath.Join(outputDir, imageFilename(ref, platform))

	// Defense-in-depth: confirm the assembled path stays within the output directory.
	cleanOut := filepath.Clean(outputDir)
	cleanPath := filepath.Clean(outputPath)
	if !strings.HasPrefix(cleanPath, cleanOut+string(filepath.Separator)) {
		return "", fmt.Errorf("output path escapes cache directory: %s", cleanPath)
	}

	log.Info("Creating tar archive")
	if err := createTar(tempDir, outputPath); err != nil {
		return "", fmt.Errorf("failed to create tar: %w", err)
	}

	log.WithField("path", outputPath).Info("Image saved")
	return outputPath, nil
}

// DownloadImage downloads a Docker image and saves it as a tar file
func DownloadImage(imageRef string, outputDir string, platform Platform) (string, error) {
	ref := ParseImageReference(imageRef)

	// Validate the image reference to prevent SSRF and other attacks
	if err := ValidateImageReference(ref); err != nil {
		return "", fmt.Errorf("invalid image reference: %w", err)
	}

	client, err := authenticateClient(ref)
	if err != nil {
		return "", err
	}

	manifest, err := fetchManifest(client, ref, platform)
	if err != nil {
		return "", err
	}

	tempDir, err := os.MkdirTemp("", "docker-image-*")
	if err != nil {
		return "", err
	}
	defer func(path string) {
		if err := os.RemoveAll(path); err != nil {
			log.WithError(err).Warn("Failed to remove temp dir")
		}
	}(tempDir)

	imageConfig, configDigest, err := downloadImageConfig(client, ref, manifest, tempDir)
	if err != nil {
		return "", err
	}

	layerPaths, err := downloadAllLayers(client, ref, manifest, imageConfig, tempDir)
	if err != nil {
		return "", err
	}

	if err := createDockerManifest(ref, configDigest, layerPaths, tempDir); err != nil {
		return "", err
	}

	if err := createRepositoriesFile(ref, layerPaths, tempDir); err != nil {
		return "", err
	}

	return createOutputTar(ref, tempDir, outputDir, platform)
}

// GetImagePlatforms returns the available platforms for a multi-arch image.
// Returns nil, nil if the image is single-arch.
func GetImagePlatforms(imageRef string) ([]Platform, error) {
	ref := ParseImageReference(imageRef)

	if err := ValidateImageReference(ref); err != nil {
		return nil, fmt.Errorf("invalid image reference: %w", err)
	}

	client, err := authenticateClient(ref)
	if err != nil {
		return nil, err
	}

	return client.GetPlatforms(ref)
}

// marshalJSONToFile marshals v to JSON and writes it to dir/filename.
func marshalJSONToFile(v interface{}, dir, filename string) error {
	data, err := json.Marshal(v)
	if err != nil {
		return fmt.Errorf("failed to marshal %s: %w", filename, err)
	}
	return os.WriteFile(filepath.Join(dir, filename), data, 0644)
}