packer-cn/builder/vmware/step_download_iso.go

package vmware

import (
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"
)

// This step downloads the ISO specified.
//
// Uses:
//   cache packer.Cache
//   config *config
//   ui     packer.Ui
//
// Produces:
//   iso_path string
type stepDownloadISO struct{}

func (s stepDownloadISO) Run(state map[string]interface{}) multistep.StepAction {
	cache := state["cache"].(packer.Cache)
	config := state["config"].(*config)
	ui := state["ui"].(packer.Ui)

	log.Printf("Acquiring lock to download the ISO.")
	cachePath := cache.Lock(config.ISOUrl)
	defer cache.Unlock(config.ISOUrl)

	url, err := url.Parse(config.ISOUrl)
	if err != nil {
		ui.Error(fmt.Sprintf("Error parsing iso_url: %s", err))
		return multistep.ActionHalt
	}

	// Start the download in a goroutine so that we cancel it and such.
	var progress uint
	downloadComplete := make(chan bool, 1)
	go func() {
		ui.Say("Copying or downloading ISO. Progress will be shown periodically.")
		cachePath, err = s.downloadUrl(cachePath, url, &progress)
		downloadComplete <- true
	}()

	progressTimer := time.NewTicker(15 * time.Second)
	defer progressTimer.Stop()

DownloadWaitLoop:
	for {
		select {
		case <-downloadComplete:
			log.Println("Download of ISO completed.")
			break DownloadWaitLoop
		case <-progressTimer.C:
			ui.Say(fmt.Sprintf("Download progress: %d%%", progress))
		case <-time.After(1 * time.Second):
			if _, ok := state[multistep.StateCancelled]; ok {
				ui.Say("Interrupt received. Cancelling download...")
				return multistep.ActionHalt
			}
		}
	}

	if err != nil {
		ui.Error(fmt.Sprintf("Error downloading ISO: %s", err))
		return multistep.ActionHalt
	}

	log.Printf("Path to ISO on disk: %s", cachePath)
	state["iso_path"] = cachePath

	return multistep.ActionContinue
}

func (stepDownloadISO) Cleanup(map[string]interface{}) {}

func (stepDownloadISO) downloadUrl(path string, url *url.URL, progress *uint) (string, error) {
	if url.Scheme == "file" {
		// If it is just a file URL, then we already have the ISO
		return path, nil
	}

	// Otherwise, it is an HTTP URL, and we must download it.
	f, err := os.Create(path)
	if err != nil {
		return "", err
	}
	defer f.Close()

	log.Printf("Beginning download of ISO: %s", url.String())
	resp, err := http.Get(url.String())
	if err != nil {
		return "", err
	}

	var buffer [4096]byte
	var totalRead int64
	for {
		n, err := resp.Body.Read(buffer[:])
		if err != nil && err != io.EOF {
			return "", err
		}

		totalRead += int64(n)
		*progress = uint((float64(totalRead) / float64(resp.ContentLength)) * 100)

		if _, werr := f.Write(buffer[:n]); werr != nil {
			return "", werr
		}

		if err == io.EOF {
			break
		}
	}

	return path, nil
}