proxmox-tools/download.go

package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func downloadImage(config *Config) error {
	filename := getFilenameFromURL(config.ImageURL)
	filepath := filepath.Join("/tmp", filename)

	if _, err := os.Stat(filepath); err == nil {
		fmt.Printf("Image already exists at %s\n", filepath)

		if err := verifyImage(filepath); err != nil {
			fmt.Printf("Image verification failed: %v\n", err)
			fmt.Println("Removing corrupted image and re-downloading...")
			os.Remove(filepath)
		} else {
			fmt.Println("Image verification passed, skipping download")
			config.ImageURL = filepath
			return nil
		}
	}

	maxRetries := 3
	for attempt := 1; attempt <= maxRetries; attempt++ {
		if attempt > 1 {
			fmt.Printf("\nRetry attempt %d/%d...\n", attempt, maxRetries)
			time.Sleep(time.Second * 2)
		}

		fmt.Printf("Downloading image from %s...\n", config.ImageURL)

		resp, err := http.Get(config.ImageURL)
		if err != nil {
			if attempt == maxRetries {
				return fmt.Errorf("failed to download image after %d attempts: %w", maxRetries, err)
			}
			fmt.Printf("Download failed: %v\n", err)
			continue
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			if attempt == maxRetries {
				return fmt.Errorf("bad status: %s", resp.Status)
			}
			fmt.Printf("Bad status: %s\n", resp.Status)
			continue
		}

		out, err := os.Create(filepath)
		if err != nil {
			return fmt.Errorf("failed to create file: %w", err)
		}
		defer out.Close()

		counter := &WriteCounter{Total: resp.ContentLength}
		_, err = io.Copy(out, io.TeeReader(resp.Body, counter))
		if err != nil {
			out.Close()
			os.Remove(filepath)
			if attempt == maxRetries {
				return fmt.Errorf("failed to save image after %d attempts: %w", maxRetries, err)
			}
			fmt.Printf("\nDownload failed: %v\n", err)
			continue
		}
		out.Close()

		fmt.Println("\nDownload completed! Verifying image integrity...")

		if err := verifyImage(filepath); err != nil {
			os.Remove(filepath)
			if attempt == maxRetries {
				return fmt.Errorf("image verification failed after %d attempts: %w", maxRetries, err)
			}
			fmt.Printf("Verification failed: %v\n", err)
			continue
		}

		fmt.Println("Image verification passed!")
		config.ImageURL = filepath
		return nil
	}

	return fmt.Errorf("failed to download and verify image after %d attempts", maxRetries)
}

func verifyImage(filepath string) error {
	fmt.Printf("Checking image with qemu-img...\n")

	cmd := exec.Command("qemu-img", "check", filepath)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("qemu-img check failed: %w\nOutput: %s", err, string(output))
	}

	fmt.Printf("Computing SHA256 checksum...\n")
	file, err := os.Open(filepath)
	if err != nil {
		return fmt.Errorf("failed to open file for checksum: %w", err)
	}
	defer file.Close()

	hash := sha256.New()
	if _, err := io.Copy(hash, file); err != nil {
		return fmt.Errorf("failed to compute checksum: %w", err)
	}

	checksum := fmt.Sprintf("%x", hash.Sum(nil))
	fmt.Printf("SHA256: %s\n", checksum)

	return nil
}

func getFilenameFromURL(url string) string {
	parts := strings.Split(url, "/")
	return parts[len(parts)-1]
}

type WriteCounter struct {
	Total      int64
	Downloaded int64
}

func (wc *WriteCounter) Write(p []byte) (int, error) {
	n := len(p)
	wc.Downloaded += int64(n)
	wc.printProgress()
	return n, nil
}

func (wc *WriteCounter) printProgress() {
	fmt.Printf("\r")
	if wc.Total > 0 {
		percent := float64(wc.Downloaded) / float64(wc.Total) * 100
		fmt.Printf("Downloading... %.0f%% (%d/%d MB)", percent, wc.Downloaded/1024/1024, wc.Total/1024/1024)
	} else {
		fmt.Printf("Downloading... %d MB", wc.Downloaded/1024/1024)
	}
}