Last active
September 26, 2016 11:15
-
-
Save cmeon/8fd10b48575b22b526c57e055f272641 to your computer and use it in GitHub Desktop.
Replace an environment variable in config ( derived from mozilla/heka )
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| package main | |
| import ( | |
| "bufio" | |
| "bytes" | |
| "errors" | |
| "flag" | |
| "io" | |
| "io/ioutil" | |
| "log" | |
| "os" | |
| ) | |
| const ( | |
| HEKA_DAEMON = "hekad" | |
| invalidEnvChars = "\n\r\t " | |
| ) | |
| func ReplaceEnvsFile(path string) (string, error) { | |
| file, err := os.Open(path) | |
| if err != nil { | |
| return "", err | |
| } | |
| r, err := EnvSub(file) | |
| if err != nil { | |
| return "", err | |
| } | |
| contents, err := ioutil.ReadAll(r) | |
| if err != nil { | |
| return "", err | |
| } | |
| return string(contents), nil | |
| } | |
| func EnvSub(r io.Reader) (io.Reader, error) { | |
| var ( | |
| invalidEnvPrefix = []byte("${") | |
| ErrMissingCloseDelim = errors.New("Missing closing delimiter") | |
| ErrInvalidChars = errors.New("Invalid characters in environmental variable") | |
| ) | |
| bufIn := bufio.NewReader(r) | |
| bufOut := new(bytes.Buffer) | |
| for { | |
| chunk, err := bufIn.ReadBytes(byte('$')) | |
| if err != nil { | |
| if err == io.EOF { | |
| // We're done. | |
| bufOut.Write(chunk) | |
| break | |
| } | |
| return nil, err | |
| } | |
| bufOut.Write(chunk[:len(chunk)-1]) | |
| tmp := make([]byte, 1) | |
| tmp, err = bufIn.Peek(1) | |
| if err != nil { | |
| if err == io.EOF { | |
| // End of file, write the last few bytes out and exit. | |
| bufOut.WriteRune('$') | |
| bufOut.Write(tmp) | |
| break | |
| } | |
| return nil, err | |
| } | |
| if string(tmp) == "{" { | |
| chunk, err = bufIn.ReadBytes(byte('}')) | |
| if err != nil { | |
| if err == io.EOF { | |
| // No closing delimiter, return an error | |
| return nil, ErrMissingCloseDelim | |
| } | |
| return nil, err | |
| } | |
| // `chunk` is now holding var name + closing delimiter. | |
| // var name contains invalid characters, return an error | |
| if bytes.IndexAny(chunk, invalidEnvChars) != -1 || | |
| bytes.Index(chunk, invalidEnvPrefix) != -1 { | |
| return nil, ErrInvalidChars | |
| } | |
| varName := string(chunk[1 : len(chunk)-1]) | |
| varVal := os.Getenv(varName) | |
| println(varName) | |
| bufOut.WriteString(varVal) | |
| } else { | |
| // Just a random '%', not an opening delimiter, write it out and | |
| // keep going. | |
| bufOut.WriteRune('$') | |
| } | |
| } | |
| return bufOut, nil | |
| } | |
| func main() { | |
| var ( | |
| i string | |
| o string | |
| ) | |
| flag.StringVar(&i, "infile", "", "Input file name") | |
| flag.StringVar(&o, "outfile", "", "Output file name") | |
| flag.Parse() | |
| iFile, err := os.Open(i) | |
| out, err := EnvSub(iFile) | |
| if err != nil { | |
| log.Fatalf("[ERR] %v", err) | |
| } | |
| oFile, err := os.Create(o) | |
| if err != nil { | |
| log.Fatalf("[ERR] %v", err) | |
| } | |
| _, err = io.Copy(oFile, out) | |
| if err != nil { | |
| log.Fatalf("[ERR] %v", err) | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment