From 9c0c5c210d815f43f8350d678c271d4bc4374e16 Mon Sep 17 00:00:00 2001
From: Lucas Bremgartner <lucas.bremgartner@swisscom.com>
Date: Sun, 2 Jul 2017 21:26:48 +0200
Subject: [PATCH] Allow to set env variables based on commands

---
 README.md      |  6 ++++--
 command.go     | 36 ++++++++++++++++++++++++++++++++++++
 environment.go | 12 +++++++++---
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 5142a36..b3c03aa 100644
--- a/README.md
+++ b/README.md
@@ -298,8 +298,9 @@ on your system, Homemaker defines a couple of extra ones for ease of use:
 
     Variant used for task and macro execution.
 
-Environment variables can also be set within tasks block by assigning them to the `envs` variable. The example below
-demonstrates the setting and clearing of environment variables:
+Environment variables can also be set within tasks block by assigning them to the `envs` variable. The `!` prefix for,
+the first value allows to assign environment variables with the result (output of stdout) of an arbitrary command.
+The example below demonstrates the setting and clearing of environment variables:
 
 ```
 [tasks.default]
@@ -307,6 +308,7 @@ demonstrates the setting and clearing of environment variables:
         ["MYENV1", "foo"],        # set MYENV1 to foo
         ["MYENV2", "foo", "bar"], # set MYENV2 to foo,bar
         ["MYENV3"],               # clear MYENV3
+        ["MYENV4", "!hostname", "-s"], # set MYENV4 to the output of `hostname -s`
     ]
 ```
 
diff --git a/command.go b/command.go
index a128937..7a242f7 100644
--- a/command.go
+++ b/command.go
@@ -23,6 +23,7 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"log"
 	"os"
@@ -116,3 +117,38 @@ func processCmd(params []string, interact bool, conf *config) error {
 
 	return exec()
 }
+
+func processCmdWithReturn(params []string, conf *config) (string, error) {
+	args := appendExpEnv(nil, params)
+	if len(args) == 0 {
+		return "", fmt.Errorf("invalid command statement")
+	}
+
+	cmdName := args[0]
+	var cmdArgs []string
+	if len(args) > 1 {
+		cmdArgs = args[1:]
+	}
+
+	if strings.HasPrefix(cmdName, "@") {
+		return "", processCmdMacro(cmdName, cmdArgs, false, conf)
+	}
+
+	if conf.flags&flagVerbose != 0 {
+		log.Printf("executing command (with return): %s %s", cmdName, strings.Join(cmdArgs, " "))
+	}
+
+	exec := func() (string, error) {
+		var stdout bytes.Buffer
+		cmd := exec.Command(cmdName, cmdArgs...)
+		cmd.Dir = conf.dstDir
+		cmd.Stderr = os.Stderr
+		cmd.Stdout = &stdout
+		cmd.Stdin = os.Stdin
+
+		err := cmd.Run()
+		return strings.Trim(stdout.String(), "\r\n"), err
+	}
+
+	return exec()
+}
diff --git a/environment.go b/environment.go
index 335c519..c8d32e8 100644
--- a/environment.go
+++ b/environment.go
@@ -42,10 +42,16 @@ func processEnv(env []string, conf *config) error {
 		}
 		os.Unsetenv(args[0])
 		return nil
-	case len(args) == 2:
-		value = args[1]
 	default:
-		value = strings.Join(args[1:], ",")
+		if strings.HasPrefix(args[1], "!") {
+			var err error
+			args[1] = strings.TrimLeft(args[1], "!")
+			if value, err = processCmdWithReturn(args[1:], conf); err != nil {
+				return err
+			}
+		} else {
+			value = strings.Join(args[1:], ",")
+		}
 	}
 
 	if conf.flags&flagVerbose != 0 {