From 2c2435a7a55e40eef0e7c6a8b849b1199b98b029 Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Sat, 13 May 2023 09:43:06 +0000
Subject: [PATCH] refactor: use cavaliergopher/grab to download files

---
 go.mod    |  1 +
 go.sum    |  2 ++
 images.go | 57 ++++++++++++++++---------------------------------------
 3 files changed, 19 insertions(+), 41 deletions(-)

diff --git a/go.mod b/go.mod
index 7112601..3f737dc 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ replace code.dumpstack.io/tools/out-of-tree/distro => ./distro
 require (
 	github.com/PuerkitoBio/goquery v1.8.1
 	github.com/alecthomas/kong v0.7.1
+	github.com/cavaliergopher/grab/v3 v3.0.1
 	github.com/go-git/go-git/v5 v5.6.1
 	github.com/mattn/go-sqlite3 v1.14.16
 	github.com/mitchellh/go-homedir v1.1.0
diff --git a/go.sum b/go.sum
index a1c2826..bf72659 100644
--- a/go.sum
+++ b/go.sum
@@ -23,6 +23,8 @@ github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkY
 github.com/boltdb/bolt v1.3.1 h1:JQmyP4ZBrce+ZQu0dY660FMfatumYDLun9hBCUVIkF4=
 github.com/boltdb/bolt v1.3.1/go.mod h1:clJnj/oiGkjum5o1McbSZDSLxVThjynRyGBgiAx27Ps=
 github.com/bwesterb/go-ristretto v1.2.0/go.mod h1:fUIoIZaG73pV5biE2Blr2xEzDoMj7NFEuV9ekS419A0=
+github.com/cavaliergopher/grab/v3 v3.0.1 h1:4z7TkBfmPjmLAAmkkAZNX/6QJ1nNFdv3SdIHXju0Fr4=
+github.com/cavaliergopher/grab/v3 v3.0.1/go.mod h1:1U/KNnD+Ft6JJiYoYBAimKH2XrYptb8Kl3DFGmsjpq4=
 github.com/cloudflare/circl v1.1.0 h1:bZgT/A+cikZnKIwn7xL2OBj012Bmvho/o6RpRvv3GKY=
 github.com/cloudflare/circl v1.1.0/go.mod h1:prBCrKB9DV4poKZY1l9zBXg2QJY7mvgRvtMxxK7fi4I=
 github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
diff --git a/images.go b/images.go
index b909af6..ac1630e 100644
--- a/images.go
+++ b/images.go
@@ -7,19 +7,20 @@ package main
 import (
 	"errors"
 	"fmt"
-	"io"
 	"io/ioutil"
-	"net/http"
+	"net/url"
 	"os"
 	"os/exec"
 	"os/user"
 	"strings"
 	"time"
 
+	"github.com/cavaliergopher/grab/v3"
+	"github.com/rs/zerolog/log"
+
 	"code.dumpstack.io/tools/out-of-tree/config"
 	"code.dumpstack.io/tools/out-of-tree/fs"
 	"code.dumpstack.io/tools/out-of-tree/qemu"
-	"github.com/rs/zerolog/log"
 )
 
 type ImageCmd struct {
@@ -126,39 +127,6 @@ func (cmd *ImageEditCmd) Run(g *Globals) (err error) {
 	return
 }
 
-// inspired by Edd Turtle code
-func downloadFile(filepath string, url string) (err error) {
-	out, err := os.Create(filepath)
-	if err != nil {
-		return
-	}
-	defer out.Close()
-
-	resp, err := http.Get(url)
-	if err != nil {
-		return
-	}
-	defer resp.Body.Close()
-
-	switch resp.StatusCode {
-	case http.StatusOK:
-		break
-	case http.StatusForbidden, http.StatusNotFound:
-		err = fmt.Errorf("Cannot download %s. It looks like you need "+
-			"to generate it manually and place it "+
-			"to ~/.out-of-tree/images/. "+
-			"Check documentation for additional information.", url)
-		return
-	default:
-		err = fmt.Errorf("Something weird happens while "+
-			"download file: %d", resp.StatusCode)
-		return
-	}
-
-	_, err = io.Copy(out, resp.Body)
-	return
-}
-
 func unpackTar(archive, destination string) (err error) {
 	// NOTE: If you're change anything in tar command please check also
 	// BSD tar (or if you're using macOS, do not forget to check GNU Tar)
@@ -184,14 +152,21 @@ func downloadImage(path, file string) (err error) {
 	}
 	defer os.RemoveAll(tmp)
 
-	archive := tmp + "/" + file + ".tar.gz"
-	url := imagesBaseURL + file + ".tar.gz"
-
-	err = downloadFile(archive, url)
+	fileurl, err := url.JoinPath(imagesBaseURL, file+".tar.gz")
 	if err != nil {
 		return
 	}
 
-	err = unpackTar(archive, path)
+	resp, err := grab.Get(tmp, fileurl)
+	if err != nil {
+		err = fmt.Errorf("Cannot download %s. It looks like you need "+
+			"to generate it manually and place it "+
+			"to ~/.out-of-tree/images/. "+
+			"Check documentation for additional information.",
+			fileurl)
+		return
+	}
+
+	err = unpackTar(resp.Filename, path)
 	return
 }