main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

//定义命令行参数方式1
var blobFileName string
var embedFolder string

func Init() {
	flag.StringVar(&blobFileName, "o", "box_blob.go", "out file name")
	flag.StringVar(&embedFolder, "d", "static", "directory to pack")

	flag.Parse()
	pwd, _ := os.Getwd()
	fmt.Println("pwd: ", pwd)
	fmt.Println("out file name: " + blobFileName + ", directory to pack: " + embedFolder)
}

var conv = map[string]interface{}{"conv": fmtByteSlice}
var tmplBlob = template.Must(template.New("").Funcs(conv).Parse(`
// Code generated by go generate; DO NOT EDIT.

// +build !codeanalysis

package {{.pkg}}

func init() {
	{{- range $name, $file := .configs }}
    	box.Add("{{ $name }}", []byte{ {{ conv $file }} })
	{{- end }}
}`),
)

var tmpl = template.Must(template.New("").Funcs(conv).Parse(`
// Code generated by go generate; DO NOT EDIT.

package {{.pkg}}

type embedBox struct {
	storage map[string][]byte
}

// Create new box for embed files
func newEmbedBox() *embedBox {
	return &embedBox{storage: make(map[string][]byte)}
}

// Add a file to box
func (e *embedBox) Add(file string, content []byte) {
	e.storage[file] = content
}

// Get file's content
// Always use / for looking up
// For example: /init/README.md is actually configs/init/README.md
func (e *embedBox) Get(file string) []byte {
	if f, ok := e.storage[file]; ok {
		return f
	}
	return nil
}

// Find for a file
func (e *embedBox) Has(file string) bool {
	if _, ok := e.storage[file]; ok {
		return true
	}
	return false
}

// Return inner storage
func (e *embedBox) Map() map[string][]byte {
	return e.storage
}

// Embed box expose
var box = newEmbedBox()
`),
)

func fmtByteSlice(s []byte) string {
	builder := strings.Builder{}

	for _, v := range s {
		builder.WriteString(fmt.Sprintf("%d,", int(v)))
	}

	return builder.String()
}

func main() {
	Init()
	// Checking directory with files
	if _, err := os.Stat(embedFolder); os.IsNotExist(err) {
		log.Fatalf("Static directory:%s does not exists!", embedFolder)
	}

	// Create map for filenames
	configs := make(map[string][]byte)
	infos := make(map[string]interface{})
	infos["pkg"] = os.Getenv("GOPACKAGE")
	if infos["pkg"] == "" {
		infos["pkg"] = "main"
	}
	infos["configs"] = configs

	// Walking through embed directory
	err := filepath.Walk(embedFolder, func(path string, info os.FileInfo, err error) error {
		relativePath := filepath.ToSlash(strings.TrimPrefix(path, embedFolder))

		if info.IsDir() {
			// Skip directories
			log.Println(path, "is a directory, skipping...")
			return nil
		} else {
			// If element is a simple file, embed
			log.Println(path, "is a file, packing in...")

			b, err := ioutil.ReadFile(path)
			if err != nil {
				// If file not reading
				log.Printf("Error reading %s: %s", path, err)
				return err
			}

			// Add file name to map
			configs[relativePath] = b
		}

		return nil
	})
	if err != nil {
		log.Fatal("Error walking through embed directory:", err)
	}

	// Create buffer
	builderBox := &bytes.Buffer{}
	builderBoxBlob := &bytes.Buffer{}

	// Execute template
	if err = tmpl.Execute(builderBox, infos); err != nil {
		log.Fatal("Error executing template", err)
	}
	if err = tmplBlob.Execute(builderBoxBlob, infos); err != nil {
		log.Fatal("Error executing template", err)
	}

	// Formatting generated code
	data, err := format.Source(builderBox.Bytes())
	if err != nil {
		log.Fatal("Error formatting generated code", err)
	}
	// Writing blob file
	if err = ioutil.WriteFile(blobFileName, data, os.ModePerm); err != nil {
		log.Fatal("Error writing blob file", err)
	}

	dataBlob, err := format.Source(builderBoxBlob.Bytes())
	if err != nil {
		log.Fatal("Error formatting generated code", err)
	}
	if err = ioutil.WriteFile(strings.TrimSuffix(filepath.Base(blobFileName), filepath.Ext(blobFileName))+"_blob.go", dataBlob, os.ModePerm); err != nil {
		log.Fatal("Error writing blob file", err)
	}
}