github.com/jasonkeene/cli@v6.14.1-0.20160816203908-ca5715166dfb+incompatible/utils/downloader/file_download.go (about)

     1  package downloader
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net/http"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  )
    11  
    12  type Downloader interface {
    13  	DownloadFile(string) (int64, string, error)
    14  	RemoveFile() error
    15  	SavePath() string
    16  }
    17  
    18  type downloader struct {
    19  	saveDir    string
    20  	filename   string
    21  	downloaded bool
    22  }
    23  
    24  func NewDownloader(saveDir string) Downloader {
    25  	return &downloader{
    26  		saveDir:    saveDir,
    27  		downloaded: false,
    28  	}
    29  }
    30  
    31  //this func returns byte written, filename and error
    32  func (d *downloader) DownloadFile(url string) (int64, string, error) {
    33  	c := http.Client{
    34  		CheckRedirect: func(r *http.Request, via []*http.Request) error {
    35  			r.URL.Opaque = r.URL.Path
    36  
    37  			//some redirect return '/' as url
    38  			if strings.Trim(r.URL.Opaque, "/") != "" {
    39  				url = r.URL.Opaque
    40  			}
    41  
    42  			return nil
    43  		},
    44  	}
    45  
    46  	r, err := c.Get(url)
    47  
    48  	if err != nil {
    49  		return 0, "", err
    50  	}
    51  	defer r.Body.Close()
    52  
    53  	if r.StatusCode == 200 {
    54  		d.filename = getFilenameFromHeader(r.Header.Get("Content-Disposition"))
    55  
    56  		if d.filename == "" {
    57  			d.filename = getFilenameFromURL(url)
    58  		}
    59  
    60  		f, err := os.Create(filepath.Join(d.saveDir, d.filename))
    61  		if err != nil {
    62  			return 0, "", err
    63  		}
    64  		defer f.Close()
    65  
    66  		size, err := io.Copy(f, r.Body)
    67  		if err != nil {
    68  			return 0, "", err
    69  		}
    70  
    71  		d.downloaded = true
    72  		return size, d.filename, nil
    73  
    74  	}
    75  	return 0, "", fmt.Errorf("Error downloading file from %s", url)
    76  }
    77  
    78  func (d *downloader) RemoveFile() error {
    79  	if !d.downloaded {
    80  		return nil
    81  	}
    82  	d.downloaded = false
    83  	return os.Remove(filepath.Join(d.saveDir, d.filename))
    84  }
    85  
    86  func getFilenameFromHeader(h string) string {
    87  	if h == "" {
    88  		return ""
    89  	}
    90  
    91  	contents := strings.Split(h, ";")
    92  	for _, content := range contents {
    93  		if strings.Contains(content, "filename=") {
    94  			content = strings.TrimSpace(content)
    95  			name := strings.TrimLeft(content, "filename=")
    96  			return strings.Trim(name, `"`)
    97  		}
    98  	}
    99  
   100  	return ""
   101  }
   102  
   103  func getFilenameFromURL(url string) string {
   104  	tmp := strings.Split(url, "/")
   105  	token := tmp[len(tmp)-1]
   106  
   107  	if i := strings.LastIndex(token, "?"); i != -1 {
   108  		token = token[i+1:]
   109  	}
   110  
   111  	if i := strings.LastIndex(token, "&"); i != -1 {
   112  		token = token[i+1:]
   113  	}
   114  
   115  	if i := strings.LastIndex(token, "="); i != -1 {
   116  		return token[i+1:]
   117  	}
   118  
   119  	return token
   120  }
   121  
   122  func (d *downloader) SavePath() string {
   123  	return d.saveDir
   124  }