webdriver-downloader.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"path/filepath"
	"archive/zip"
	"io/ioutil"
	"regexp"
)

func unzip(sourceFilePath string,  message chan<- string, status chan<- int) {
	// set destination to current directory
	destination, err := os.Getwd()
	if err != nil {
		message <- "Error getting current directory"
		status <- 1
		return
	}

	message <- "Unzipping " + sourceFilePath + " to ./"

	// open a reader on source zip file
    r, err := zip.OpenReader(sourceFilePath)
	if err != nil {
		message <- "Error opening zip file " + sourceFilePath
		status <- 1
		return
	}
    defer r.Close()

    for _, f := range r.File {
		// get file descriptor for current iterated file
        readCloser, err := f.Open()
		if err != nil {
			message <- "Error opening file in archive " + f.Name
			status <- 1
			return
		}
        defer readCloser.Close()

		// create an output path from passed destination and current filename
        outputPath := filepath.Join(destination, f.Name)

        if f.FileInfo().IsDir() {
			// current file is an empty dir, create it
            os.MkdirAll(outputPath, os.ModePerm)
        } else {
			// find last path separator, split from beginning to that index to
			// get full path minus file name
            var fileDirPath string
            if lastIndex := strings.LastIndex(outputPath, string(os.PathSeparator)); lastIndex > -1 {
                fileDirPath = outputPath[:lastIndex]
            }

			// make all directories in current file path
            err = os.MkdirAll(fileDirPath, os.ModePerm)
			if err != nil {
				message <- "Error creating directory " + fileDirPath
				status <- 1
				return
			}

			// open file on disk for writing contents of zip file
            f, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
			if err != nil {
				message <- "Error opening file for writing" + outputPath
				status <- 1
				return
			}
            defer f.Close()

			// copy data from reader to output file
            _, err = io.Copy(f, readCloser)
			if err != nil {
				message <- "Error writing to file" + outputPath
				status <- 1
				return
			}
        }
    }

	// finished
	message <- "Finished unzipping " + sourceFilePath + ", removing zip file"

	// remove zip
	err = os.Remove(sourceFilePath)
	if err != nil {
		message <- "Could not remove zip file " + sourceFilePath
		status <- 1
		return
	}

	status <- 0
}

func dl(url string, message chan<- string, status chan<- int) {
	var fileName string
	// find last path separator, split from that index+1 to end of string to get filename
	if lastIndex := strings.LastIndex(url, "/"); lastIndex > -1 {
		fileName = url[lastIndex+1:]
	}

	// check if file already exists
	var _, err = os.Stat(fileName)
	// negate the IsNotExist result; return early if file DOES exist
	if !os.IsNotExist(err) {
		message <- "Skipping " + fileName + ", file already exists"
		status <- 0
		return
	}

	message <- "Downloading " + fileName

	// download data from url
	response, err := http.Get(url)
	if err != nil {
		message <- "Error contacting server for URL " + url
		status <- 1
	}
	defer response.Body.Close()

	// create file for downloaded data
	output, err := os.Create(fileName)
	if err != nil {
		message <- "Error creating file" + fileName
		return
	}
	defer output.Close()

	// copy downloaded data to file
	_, err = io.Copy(output, response.Body)
	if err != nil {
		message <- "Error copying data to file" + fileName
		status <- 1
	}

	message <- "Finished downloading " + fileName

	// unzip to cwd if zip file
	if strings.Contains(fileName, "zip") {
		unzip(fileName, message, status)
	} else {
		// respond with success
		status <- 0
	}
}

func main() {
	// match commented or invalid lines
	REGEX_IGNORE_LINE := regexp.MustCompile(`^[#\W\n]`)

	// read urls file
	data, err := ioutil.ReadFile("WEBDRIVER_DOWNLOAD_URLS.txt")
	if err != nil {
		fmt.Println("Error reading urls file")
		return
	}

	// replace all windows linebr chars with real ones
	strdata := strings.Replace(string(data), "\r\n", "\n", -1)
	urls := strings.Split(strdata, "\n")

	// channels
	message := make(chan string)
	status := make(chan int, len(urls))

	// keep track of the number of goroutines we spawn
	routineCount := 0

	// for each url
	for _, url := range urls {
		// ignore commented and empty lines
		if REGEX_IGNORE_LINE.MatchString(url) || len(url) < 1 {
			continue
		}

		// download from url
		// dl calls unzip on files when done
		go dl(url, message, status)
		routineCount++
	}

	// print anything that comes in through the message channel
	go func() {
		for { fmt.Println(<-message) }
	}()

	// block main returning until all urls respond on status channel
	for i := 0; i < routineCount; i++ {
		<-status
	}
}