s3zipper.go

package main

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"net/http"

	"github.com/AdRoll/goamz/aws"
	"github.com/AdRoll/goamz/s3"
	redigo "github.com/garyburd/redigo/redis"
)

type Configuration struct {
	AccessKey          string
	SecretKey          string
	Bucket             string
	Region             string
	RedisServerAndPort string
	Port               int
}

var config = Configuration{}
var aws_bucket *s3.Bucket
var redisPool *redigo.Pool

type RedisFile struct {
	FileName  string
	Folder    string
	S3Path    string
	DataValue string
	//// Optional - we use are Teamwork.com but feel free to rmove
	//FileId       int64 `json:",string"`
	//ProjectId    int64 `json:",string"`
	//ProjectName  string
	//Modified     string
	//ModifiedTime time.Time
}

func main() {
	if 1 == 0 {
		test()
		return
	}

	configFile, _ := os.Open("conf.json")
	decoder := json.NewDecoder(configFile)
	err := decoder.Decode(&config)
	if err != nil {
		panic("Error reading conf")
	}

	initAwsBucket()
	InitRedis()

	fmt.Println("Running on port", config.Port)
	http.HandleFunc("/", handler)
	http.ListenAndServe(":" + strconv.Itoa(config.Port), nil)
}

func test() {
	var err error
	var files []*RedisFile
	jsonData := "[{\"S3Path\":\"1\\/p23216.tf_A89A5199-F04D-A2DE-5824E635AC398956.Avis_Rent_A_Car_Print_Reservation.pdf\",\"FileVersionId\":\"4164\",\"FileName\":\"Avis Rent A Car_ Print Reservation.pdf\",\"ProjectName\":\"Superman\",\"ProjectId\":\"23216\",\"Folder\":\"\",\"FileId\":\"4169\"},{\"modified\":\"2015-07-18T02:05:04Z\",\"S3Path\":\"1\\/p23216.tf_351310E0-DF49-701F-60601109C2792187.a1.jpg\",\"FileVersionId\":\"4165\",\"FileName\":\"a1.jpg\",\"ProjectName\":\"Superman\",\"ProjectId\":\"23216\",\"Folder\":\"Level 1\\/Level 2 x\\/Level 3\",\"FileId\":\"4170\"}]"

	resultByte := []byte(jsonData)

	err = json.Unmarshal(resultByte, &files)
	if err != nil {
		err = errors.New("Error decoding json: " + jsonData)
	}

	//parseFileDates(files)
}

//func parseFileDates(files []*RedisFile) {
//	layout := "2006-01-02T15:04:05Z"
//	for _, file := range files {
//		t, err := time.Parse(layout, file.Modified)
//		if err != nil {
//			fmt.Println(err)
//			continue
//		}
//		file.ModifiedTime = t
//	}
//}

func initAwsBucket() {
	expiration := time.Now().Add(time.Hour * 1)
	auth, err := aws.GetAuth(config.AccessKey, config.SecretKey, "", expiration) //"" = token which isn't needed
	if err != nil {
		panic(err)
	}

	aws_bucket = s3.New(auth, aws.GetRegion(config.Region)).Bucket(config.Bucket)
}

func InitRedis() {
	redisPool = &redigo.Pool{
		MaxIdle:     10,
		IdleTimeout: 1 * time.Second,
		Dial: func() (redigo.Conn, error) {
			return redigo.Dial("tcp", config.RedisServerAndPort)
		},
		TestOnBorrow: func(c redigo.Conn, t time.Time) (err error) {
			_, err = c.Do("PING")
			if err != nil {
				panic("Error connecting to redis")
			}
			return
		},
	}
}

// Remove all other unrecognised characters apart from
var makeSafeFileName = regexp.MustCompile(`[#<>:"/\|?*\\]`)

func getFilesFromRedis(ref string) (files []*RedisFile, err error) {

	// Testing - enable to test. Remove later.
	if 1 == 0 && ref == "test" {
		files = append(files, &RedisFile{FileName: "test.zip", Folder: "", S3Path: "test/test.zip"}) // Edit and dplicate line to test
		return
	}

	redis := redisPool.Get()
	defer redis.Close()

	// Get the value from Redis
	result, err := redis.Do("GET", "zip:" + ref)
	if err != nil || result == nil {
		err = errors.New("Access Denied (sorry your link has timed out)")
	}

	// Convert to bytes
	var resultByte []byte
	var ok bool
	if resultByte, ok = result.([]byte); !ok {
		err = errors.New("Error converting data stream to bytes")
		return
	}

	// Decode JSON
	err = json.Unmarshal(resultByte, &files)
	if err != nil {
		err = errors.New("Error decoding json: " + string(resultByte))
	}

	// Convert mofified date strings to time objects
	//parseFileDates(files)

	return
}

func getCnfFilesFromRedis(cnf string) (CnfDatas []*RedisFile, err error) {

	redis := redisPool.Get()
	defer redis.Close()

	// Get the value from Redis
	result, err := redis.Do("GET", "cnf:" + cnf)
	if err != nil || result == nil {
		err = errors.New("Conf File Access Denied (sorry your link has timed out)")
		return
	}

	// Convert to bytes
	var resultByte []byte
	var ok bool
	if resultByte, ok = result.([]byte); !ok {
		err = errors.New("Conf File Error converting data stream to bytes")
		return
	}

	// Decode JSON
	err = json.Unmarshal(resultByte, &CnfDatas)
	if err != nil {
		err = errors.New("Conf File Error decoding json: " + string(resultByte))
	}

	return
}

func slicer(a []*RedisFile, b int) [][]*RedisFile {
	c := make([][]*RedisFile, 0)
	for i := 0; i < len(a); i += b {
		if i + b > len(a) {
			c = append(c, a[i:])
		} else {
			c = append(c, a[i:i + b])
		}
	}
	return c
}

func handler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	// Get "ref" URL params
	refs, ok := r.URL.Query()["ref"]
	if !ok || len(refs) < 1 {
		http.Error(w, "S3 File Zipper. Pass ?ref= to use.", 500)
		return
	}
	ref := refs[0]

	// Get "downloadas" URL params
	downloadas, ok := r.URL.Query()["downloadas"]
	if !ok && len(downloadas) > 0 {
		downloadas[0] = makeSafeFileName.ReplaceAllString(downloadas[0], "")
		if downloadas[0] == "" {
			downloadas[0] = "download.zip"
		}
	} else {
		downloadas = append(downloadas, "download.zip")
	}

	files, err := getFilesFromRedis(ref)

	if err != nil {
		http.Error(w, err.Error(), 403)
		log.Printf("%s\t%s\t%s", r.Method, r.RequestURI, err.Error())
		return
	}

	// Get "conf" URL parms
	cnfs, ok := r.URL.Query()["cnf"]
	if !ok || len(cnfs) < 1 {
		http.Error(w, "S3 File Zipper. Pass ?cnf= to use.", 500)
		return
	}

	cnf := cnfs[0]
	CnfDatas, err := getCnfFilesFromRedis(cnf)
	if err != nil {
		http.Error(w, err.Error(), 403)
		log.Printf("%s\t%s\t%s", r.Method, r.RequestURI, err.Error())
		return
	}

	// Start processing the response
	w.Header().Add("Content-Disposition", "attachment; filename=\"" + downloadas[0] + "\"")
	w.Header().Add("Content-Type", "application/zip")

	// Loop over files, add them to the
	zipWriter := zip.NewWriter(w)

	// We have to set a special flag so zip files recognize utf file names
	// See http://stackoverflow.com/questions/30026083/creating-a-zip-archive-with-unicode-filenames-using-gos-archive-zip

	for _, CnfData := range CnfDatas {
		// Build safe file file name
		safeFileName := makeSafeFileName.ReplaceAllString(CnfData.FileName, "")

		// Build a good path for the file within the zip
		zipPath := ""

		// Prefix folder name, if any
		if CnfData.Folder != "" {
			zipPath += CnfData.Folder
			if !strings.HasSuffix(zipPath, "/") {
				zipPath += "/"
			}
		}
		zipPath += safeFileName

		h := &zip.FileHeader{
			Name:   zipPath,
			Method: zip.Deflate,
			Flags:  0x800,
		}

		f, _ := zipWriter.CreateHeader(h)

		r := strings.NewReader(CnfData.DataValue)
		io.Copy(f, r)
	}

	sz := len(files);
	fmt.Println("Total Patches:", sz)

	chunks := slicer(files, 1000)

	fmt.Printf("Total Chunks %v\n", len(chunks))

	for _, f1 := range chunks {
		initAwsBucket();
		for _, file := range f1 {

			if file.S3Path == "" {
				log.Printf("Missing path for file: %v", file)
				continue
			}

			// Build safe file file name
			safeFileName := makeSafeFileName.ReplaceAllString(file.FileName, "")
			if safeFileName == "" {
				// Unlikely but just in case
				safeFileName = "file"
			}

			// Read file from S3, log any errors
			rdr, err := aws_bucket.GetReader(file.S3Path)
			if err != nil {
				switch t := err.(type) {
				case *s3.Error:
					if t.StatusCode == 404 {
						log.Printf("File not found. %s", file.S3Path)
					}
				default:
					log.Printf("Error downloading \"%s\" - %s", file.S3Path, err.Error())
				}
				continue
			}

			// Build a good path for the file within the zip
			zipPath := ""
			// Prefix folder name, if any
			if file.Folder != "" {
				zipPath += file.Folder
				if !strings.HasSuffix(zipPath, "/") {
					zipPath += "/"
				}
			}

			zipPath += safeFileName

			// We have to set a special flag so zip files recognize utf file names
			// See http://stackoverflow.com/questions/30026083/creating-a-zip-archive-with-unicode-filenames-using-gos-archive-zip
			h := &zip.FileHeader{
				Name:   zipPath,
				Method: zip.Deflate,
				Flags:  0x800,
			}

			f, _ := zipWriter.CreateHeader(h)

			io.Copy(f, rdr)
			rdr.Close()
		}
	}
	//for _, file := range files {
	//
	//	if file.S3Path == "" {
	//		log.Printf("Missing path for file: %v", file)
	//		continue
	//	}
	//
	//	// Build safe file file name
	//	safeFileName := makeSafeFileName.ReplaceAllString(file.FileName, "")
	//	if safeFileName == "" {
	//		// Unlikely but just in case
	//		safeFileName = "file"
	//	}
	//
	//	// Read file from S3, log any errors
	//	rdr, err := aws_bucket.GetReader(file.S3Path)
	//	if err != nil {
	//		switch t := err.(type) {
	//		case *s3.Error:
	//			if t.StatusCode == 404 {
	//				log.Printf("File not found. %s", file.S3Path)
	//			}
	//		default:
	//			log.Printf("Error downloading \"%s\" - %s", file.S3Path, err.Error())
	//		}
	//		continue
	//	}
	//
	//	// Build a good path for the file within the zip
	//	zipPath := ""
	//	// Prefix folder name, if any
	//	if file.Folder != "" {
	//		zipPath += file.Folder
	//		if !strings.HasSuffix(zipPath, "/") {
	//			zipPath += "/"
	//		}
	//	}
	//
	//	zipPath += safeFileName
	//
	//	// We have to set a special flag so zip files recognize utf file names
	//	// See http://stackoverflow.com/questions/30026083/creating-a-zip-archive-with-unicode-filenames-using-gos-archive-zip
	//	h := &zip.FileHeader{
	//		Name:   zipPath,
	//		Method: zip.Deflate,
	//		Flags:  0x800,
	//	}
	//
	//	f, _ := zipWriter.CreateHeader(h)
	//
	//	io.Copy(f, rdr)
	//	rdr.Close()
	//}

	zipWriter.Close()

	log.Printf("%s\t%s\t%s", r.Method, r.RequestURI, time.Since(start))
}