From 10370adbabbbccbd87fd5eca45f757629f33b397 Mon Sep 17 00:00:00 2001
From: Krzysztof Wilczynski <kw@linux.com>
Date: Tue, 21 Nov 2017 21:09:01 +0100
Subject: [PATCH] ansible-local: Add ability to clean staging directory.

Signed-off-by: Krzysztof Wilczynski <kw@linux.com>
---
 provisioner/ansible-local/provisioner.go      | 32 +++++++++++++++++--
 provisioner/ansible-local/provisioner_test.go | 23 +++++++++++++
 .../docs/provisioners/ansible-local.html.md   |  4 +++
 3 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/provisioner/ansible-local/provisioner.go b/provisioner/ansible-local/provisioner.go
index 2f6a873f9..ec1eec041 100644
--- a/provisioner/ansible-local/provisioner.go
+++ b/provisioner/ansible-local/provisioner.go
@@ -48,6 +48,9 @@ type Config struct {
 	// permissions in this directory.
 	StagingDir string `mapstructure:"staging_directory"`
 
+	// If true, staging directory is removed after executing ansible.
+	CleanStagingDir bool `mapstructure:"clean_staging_directory"`
+
 	// The optional inventory file
 	InventoryFile string `mapstructure:"inventory_file"`
 
@@ -260,6 +263,13 @@ func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
 	if err := p.executeAnsible(ui, comm); err != nil {
 		return fmt.Errorf("Error executing Ansible: %s", err)
 	}
+
+	if p.config.CleanStagingDir {
+		ui.Message("Removing staging directory...")
+		if err := p.removeDir(ui, comm, p.config.StagingDir); err != nil {
+			return fmt.Errorf("Error removing staging directory: %s", err)
+		}
+	}
 	return nil
 }
 
@@ -367,15 +377,33 @@ func (p *Provisioner) uploadFile(ui packer.Ui, comm packer.Communicator, dst, sr
 }
 
 func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
-	ui.Message(fmt.Sprintf("Creating directory: %s", dir))
 	cmd := &packer.RemoteCmd{
 		Command: fmt.Sprintf("mkdir -p '%s'", dir),
 	}
+
+	ui.Message(fmt.Sprintf("Creating directory: %s", dir))
 	if err := cmd.StartWithUi(comm, ui); err != nil {
 		return err
 	}
+
 	if cmd.ExitStatus != 0 {
-		return fmt.Errorf("Non-zero exit status.")
+		return fmt.Errorf("Non-zero exit status. See output above for more information.")
+	}
+	return nil
+}
+
+func (p *Provisioner) removeDir(ui packer.Ui, comm packer.Communicator, dir string) error {
+	cmd := &packer.RemoteCmd{
+		Command: fmt.Sprintf("rm -rf '%s'", dir),
+	}
+
+	ui.Message(fmt.Sprintf("Removing directory: %s", dir))
+	if err := cmd.StartWithUi(comm, ui); err != nil {
+		return err
+	}
+
+	if cmd.ExitStatus != 0 {
+		return fmt.Errorf("Non-zero exit status. See output above for more information.")
 	}
 	return nil
 }
diff --git a/provisioner/ansible-local/provisioner_test.go b/provisioner/ansible-local/provisioner_test.go
index 2195b7107..d96ed113f 100644
--- a/provisioner/ansible-local/provisioner_test.go
+++ b/provisioner/ansible-local/provisioner_test.go
@@ -188,3 +188,26 @@ func TestProvisionerPrepare_Dirs(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 }
+
+func TestProvisionerPrepare_CleanStagingDir(t *testing.T) {
+	var p Provisioner
+	config := testConfig()
+
+	playbook_file, err := ioutil.TempFile("", "playbook")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	defer os.Remove(playbook_file.Name())
+
+	config["playbook_file"] = playbook_file.Name()
+	config["clean_staging_directory"] = true
+
+	err = p.Prepare(config)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if !p.config.CleanStagingDir {
+		t.Fatalf("expected clean_staging_directory to be set")
+	}
+}
diff --git a/website/source/docs/provisioners/ansible-local.html.md b/website/source/docs/provisioners/ansible-local.html.md
index 92d46fb3d..5277c8a89 100644
--- a/website/source/docs/provisioners/ansible-local.html.md
+++ b/website/source/docs/provisioners/ansible-local.html.md
@@ -140,6 +140,10 @@ chi-appservers
     are not correct, use a shell provisioner prior to this to configure it
     properly.
 
+-   `clean_staging_directory` (boolean) - If set to `true`, the content of
+     the `staging_directory` will be removed after executing ansible. By
+     default, this is set to `false`.
+
 ## Default Extra Variables
 
 In addition to being able to specify extra arguments using the