all: make the magic word "#put" able to copy encrypted file

When issuing "#put:" or "#put!" command in the script, if the input
file is not exist it will check for the encrypted file, the one with
".vault" extension.
If it exists, the encrypted file will be used as input for copy operation.

8 files changed, 236 insertions, 47 deletions

diff --git a/.gitignore b/.gitignore
index 9c0deff..85920b1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,6 +12,8 @@
 /awwan
 /cover.html
 /cover.out
+/testdata/decrypt-with-passphrase/.awwan.env
 /testdata/encrypt-with-passphrase/.awwan.env.vault
 /testdata/encrypt-without-passphrase/.awwan.env.vault
-/testdata/decrypt-with-passphrase/.awwan.env
+/testdata/encrypt/.cache
+/testdata/encrypt/file.txt.decrypted
diff --git a/awwan_test.go b/awwan_test.go
index ff67466..1d9a23c 100644
--- a/awwan_test.go
+++ b/awwan_test.go
@@ -1,3 +1,6 @@
+// SPDX-FileCopyrightText: 2023 M. Shulhan <ms@kilabit.info>
+// SPDX-License-Identifier: GPL-3.0-or-later
+
 package awwan
 
 import (
@@ -210,3 +213,109 @@ func TestAwwanLocal_withEncryption(t *testing.T) {
 		test.Assert(t, `stdout`, string(tdata.Output[c.tdataOut]), mockout.String())
 	}
 }
+
+func TestAwwanLocalPut_withEncryption(t *testing.T) {
+	type testCase struct {
+		desc       string
+		tdataOut   string
+		passphrase string
+		expError   string
+
+		// If true, the Awwan.privateKey will be set to nil before
+		// running Local.
+		resetPrivateKey bool
+	}
+
+	// Load the test data output.
+	var (
+		baseDir = filepath.Join(`testdata`, `encrypt`)
+
+		tdata *test.Data
+		err   error
+	)
+	tdata, err = test.LoadData(filepath.Join(baseDir, `test.data`))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create the Awwan instance.
+	var (
+		mockout = bytes.Buffer{}
+		mockerr = bytes.Buffer{}
+		mockrw  = mock.ReadWriter{}
+		aww     = Awwan{}
+	)
+
+	// Mock terminal to read passphrase for private key.
+	mockrw.BufRead.WriteString("s3cret\r")
+	aww.termrw = &mockrw
+
+	err = aww.init(baseDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var (
+		script    = filepath.Join(baseDir, `local.aww`)
+		lineRange = `3`
+		fileDest  = filepath.Join(baseDir, `file.txt.decrypted`)
+
+		cases = []testCase{{
+			desc:     `WithSuccess`,
+			tdataOut: `local.aww:3:exp_file_content`,
+		}, {
+			desc:            `WithEmptyPrivateKey`,
+			expError:        `Local: loadEnvFromPaths: private key is missing or not loaded`,
+			resetPrivateKey: true,
+		}, {
+			desc:            `WithInvalidPassphrase`,
+			passphrase:      "invalid\r",
+			expError:        `Local: loadEnvFromPaths: private key is missing or not loaded`,
+			resetPrivateKey: true,
+		}}
+
+		c          testCase
+		expContent string
+		gotContent []byte
+	)
+	for _, c = range cases {
+		t.Run(c.desc, func(tt *testing.T) {
+			_ = os.Remove(fileDest)
+
+			if c.resetPrivateKey {
+				aww.privateKey = nil
+
+				// Mock terminal to read passphrase for private key.
+				mockrw.BufRead.Reset()
+				mockrw.BufRead.WriteString(c.passphrase)
+			}
+
+			var req = NewRequest(CommandModeLocal, script, lineRange)
+
+			mockout.Reset()
+			mockerr.Reset()
+			req.stdout = &mockout
+			req.stderr = &mockerr
+
+			err = aww.Local(req)
+			if err != nil {
+				test.Assert(tt, c.desc, c.expError, err.Error())
+				return
+			}
+
+			// We cannot assert the stdout since its print dynamic
+			// paths.
+
+			test.Assert(tt, `stderr`, ``, mockerr.String())
+
+			gotContent, err = os.ReadFile(fileDest)
+			if err != nil {
+				tt.Fatal(err)
+			}
+
+			expContent = string(tdata.Output[c.tdataOut])
+
+			test.Assert(tt, `content`, expContent, string(gotContent))
+		})
+	}
+}
diff --git a/session.go b/session.go
index fb848b8..7470624 100644
--- a/session.go
+++ b/session.go
@@ -4,10 +4,12 @@
 package awwan
 
 import (
+	"bytes"
 	"crypto/rsa"
 	"errors"
 	"fmt"
 	"io/fs"
+	"log"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -94,12 +96,7 @@ func (ses *Session) Vals(keyPath string) []string {
 
 // Copy file in local system.
 func (ses *Session) Copy(stmt *Statement) (err error) {
-	var (
-		logp = "Copy"
-
-		src  string
-		dest string
-	)
+	var logp = `Copy`
 
 	if len(stmt.cmd) == 0 {
 		return fmt.Errorf("%s: missing source argument", logp)
@@ -111,16 +108,26 @@ func (ses *Session) Copy(stmt *Statement) (err error) {
 		return fmt.Errorf("%s: two or more destination arguments is given", logp)
 	}
 
-	src, err = ses.generateFileInput(stmt.cmd)
+	var (
+		src     string
+		isVault bool
+	)
+
+	src, isVault, err = ses.generateFileInput(stmt.cmd)
 	if err != nil {
 		return fmt.Errorf("%s: %w", logp, err)
 	}
 
-	dest = stmt.args[0]
-
-	err = libos.Copy(dest, src)
+	err = libos.Copy(stmt.args[0], src)
+	if isVault {
+		// Delete the decrypted file on exit.
+		var errRemove = os.Remove(src)
+		if errRemove != nil {
+			log.Printf(`%s: %s`, logp, errRemove)
+		}
+	}
 	if err != nil {
-		return fmt.Errorf("%s: %w", logp, err)
+		return fmt.Errorf(`%s: %w`, logp, err)
 	}
 	return nil
 }
@@ -160,12 +167,7 @@ func (ses *Session) Get(stmt *Statement) (err error) {
 
 // Put copy file from local to remote system.
 func (ses *Session) Put(stmt *Statement) (err error) {
-	var (
-		logp = "Put"
-
-		local  string
-		remote string
-	)
+	var logp = `Put`
 
 	if len(stmt.cmd) == 0 {
 		return fmt.Errorf("%s: missing source argument", logp)
@@ -177,18 +179,29 @@ func (ses *Session) Put(stmt *Statement) (err error) {
 		return fmt.Errorf("%s: two or more destination arguments is given", logp)
 	}
 
-	local, err = ses.generateFileInput(stmt.cmd)
+	var (
+		local   string
+		isVault bool
+	)
+
+	local, isVault, err = ses.generateFileInput(stmt.cmd)
 	if err != nil {
 		return fmt.Errorf("%s: %w", logp, err)
 	}
 
-	remote = stmt.args[0]
+	var remote = stmt.args[0]
 
 	if ses.sftpc == nil {
 		err = ses.sshClient.ScpPut(local, remote)
 	} else {
 		err = ses.sftpc.Put(local, remote)
 	}
+	if isVault {
+		var errRemove = os.Remove(local)
+		if errRemove != nil {
+			log.Printf(`%s: %s`, logp, errRemove)
+		}
+	}
 	if err != nil {
 		return fmt.Errorf("%s: %w", logp, err)
 	}
@@ -215,7 +228,15 @@ func (ses *Session) SudoCopy(req *Request, stmt *Statement, withParseInput bool)
 	}
 
 	if withParseInput {
-		src, err = ses.generateFileInput(stmt.cmd)
+		var isVault bool
+
+		src, isVault, err = ses.generateFileInput(stmt.cmd)
+		if isVault {
+			var errRemove = os.Remove(src)
+			if errRemove != nil {
+				log.Printf(`%s: %s`, logp, errRemove)
+			}
+		}
 		if err != nil {
 			return fmt.Errorf("%s: %w", logp, err)
 		}
@@ -297,9 +318,8 @@ func (ses *Session) SudoGet(stmt *Statement) (err error) {
 // SudoPut copy file from local to remote using sudo.
 func (ses *Session) SudoPut(stmt *Statement) (err error) {
 	var (
-		logp = "SudoPut"
+		logp = `SudoPut`
 
-		local    string
 		baseName string
 		tmp      string
 		remote   string
@@ -316,9 +336,19 @@ func (ses *Session) SudoPut(stmt *Statement) (err error) {
 		return fmt.Errorf("%s: two or more destination arguments is given", logp)
 	}
 
+	var (
+		local   string
+		isVault bool
+	)
 	// Apply the session variables into local file to be copied first, and
 	// save them into cache directory.
-	local, err = ses.generateFileInput(stmt.cmd)
+	local, isVault, err = ses.generateFileInput(stmt.cmd)
+	if isVault {
+		var errRemove = os.Remove(local)
+		if errRemove != nil {
+			log.Printf(`%s: %s`, logp, errRemove)
+		}
+	}
 	if err != nil {
 		return fmt.Errorf("%s: %w", logp, err)
 	}
@@ -474,50 +504,56 @@ func (ses *Session) executeScriptOnRemote(req *Request, pos linePosition) {
 //
 // For example, if the input file path is "{{.BaseDir}}/a/b/script" then the
 // output file path would be "{{.BaseDir}}/.cache/a/b/script".
-func (ses *Session) generateFileInput(in string) (out string, err error) {
+func (ses *Session) generateFileInput(in string) (out string, isVault bool, err error) {
+	// Check if the file is binary first, since binary file will not get
+	// encrypted.
+	if libos.IsBinary(in) {
+		return in, false, nil
+	}
+
 	var (
 		logp = `generateFileInput`
 
-		tmpl   *template.Template
-		f      *os.File
-		outDir string
-		base   string
+		contentInput []byte
 	)
 
-	if libos.IsBinary(in) {
-		return in, nil
+	contentInput, isVault, err = ses.loadFileInput(in)
+	if err != nil {
+		return ``, false, fmt.Errorf(`%s: %w`, logp, err)
 	}
 
-	outDir = filepath.Join(ses.BaseDir, defCacheDir, filepath.Dir(in))
-	base = filepath.Base(in)
-	out = filepath.Join(outDir, base)
+	var tmpl = template.New(in)
 
-	err = os.MkdirAll(outDir, 0700)
+	tmpl, err = tmpl.Parse(string(contentInput))
 	if err != nil {
-		return "", fmt.Errorf("%s %s: %w", logp, in, err)
+		return ``, false, fmt.Errorf(`%s: %w`, logp, err)
 	}
 
-	tmpl, err = template.ParseFiles(in)
-	if err != nil {
-		return "", fmt.Errorf("%s %s: %w", logp, in, err)
-	}
+	var contentOut bytes.Buffer
 
-	f, err = os.Create(out)
+	err = tmpl.Execute(&contentOut, ses)
 	if err != nil {
-		return "", fmt.Errorf("%s %s: %w", logp, in, err)
+		return ``, false, fmt.Errorf(`%s: %w`, logp, err)
 	}
 
-	err = tmpl.Execute(f, ses)
+	var (
+		outDir = filepath.Join(ses.BaseDir, defCacheDir, filepath.Dir(in))
+		base   = filepath.Base(in)
+	)
+
+	err = os.MkdirAll(outDir, 0700)
 	if err != nil {
-		return "", fmt.Errorf("%s %s: %w", logp, in, err)
+		return ``, false, fmt.Errorf(`%s: %s: %w`, logp, outDir, err)
 	}
 
-	err = f.Close()
+	out = filepath.Join(outDir, base)
+
+	err = os.WriteFile(out, contentOut.Bytes(), 0600)
 	if err != nil {
-		return "", fmt.Errorf("%s %s: %w", logp, in, err)
+		return ``, false, fmt.Errorf(`%s: %s: %w`, logp, out, err)
 	}
 
-	return out, nil
+	return out, isVault, nil
 }
 
 // generatePaths using baseDir return all paths from BaseDir to ScriptDir.
@@ -645,6 +681,36 @@ func (ses *Session) loadFileEnv(awwanEnv string, isVault bool) (err error) {
 	return nil
 }
 
+// loadFileInput read the input file for Copy or Put operation.
+// If the original input file does not exist, try loading the encrypted file
+// with ".vault" extension.
+//
+// On success, it will return the content of file and true if the file is
+// from encrypted file .vault.
+func (ses *Session) loadFileInput(path string) (content []byte, isVault bool, err error) {
+	content, err = os.ReadFile(path)
+	if err == nil {
+		return content, false, nil
+	}
+	if !errors.Is(err, fs.ErrNotExist) {
+		return nil, false, err
+	}
+
+	path = path + defEncryptExt
+
+	content, err = os.ReadFile(path)
+	if err != nil {
+		return nil, false, err
+	}
+
+	content, err = decrypt(ses.privateKey, content)
+	if err != nil {
+		return nil, false, err
+	}
+
+	return content, true, nil
+}
+
 func (ses *Session) loadRawEnv(content []byte) (err error) {
 	var in *ini.Ini
 
diff --git a/testdata/encrypt/awwan.env b/testdata/encrypt/awwan.env
new file mode 100644
index 0000000..81b55e7
--- /dev/null
+++ b/testdata/encrypt/awwan.env
@@ -0,0 +1,2 @@
+[host]
+name = encrypt
diff --git a/testdata/encrypt/file.txt.org b/testdata/encrypt/file.txt.org
new file mode 100644
index 0000000..de6797e
--- /dev/null
+++ b/testdata/encrypt/file.txt.org
@@ -0,0 +1,2 @@
+The host name is {{.Val "host::name"}}.
+The secret password is {{.Val "secret::pass"}}.
diff --git a/testdata/encrypt/file.txt.vault b/testdata/encrypt/file.txt.vault
new file mode 100644
index 0000000..43cb223
--- /dev/null
+++ b/testdata/encrypt/file.txt.vault
@@ -0,0 +1,2 @@
+��E���{���
+;�_��%u�ۂ�K��?C��5+B�clP�8"�b5m���D�� ���B!�P�A\`2T�W,
d<����.��$V-�{��S3UTָgjj@�*�qAX�Bb�42$TTn�X�@O��Kwa$����}��j|H�����ʹ#���Ma15�m�1M!�+�(��,���k:���c��t��s���B0����l7P��췧o�6J'<�T���qs�p�`��Q�����o�>�i;�_6zV�ij�S	��m�2�
�D?��&�7���3�'yf=��T&a�j�~�A���#�!�Q��u��@�J�c��?q�cp�����9 (&,ʑ�o�N1B�T��NN���|Y�G�����%Xt���
\ No newline at end of file
diff --git a/testdata/encrypt/local.aww b/testdata/encrypt/local.aww
index d9127db..31e737d 100644
--- a/testdata/encrypt/local.aww
+++ b/testdata/encrypt/local.aww
@@ -1 +1,3 @@
 echo {{.Val "secret::pass"}}
+
+#put: {{.ScriptDir}}/file.txt {{.ScriptDir}}/file.txt.decrypted
diff --git a/testdata/encrypt/test.data b/testdata/encrypt/test.data
index 9ecdb9f..40ad1de 100644
--- a/testdata/encrypt/test.data
+++ b/testdata/encrypt/test.data
@@ -7,3 +7,7 @@ this_is_a_secret
 
 --> local:   1: echo this_is_a_secret_in_sub
 this_is_a_secret_in_sub
+
+<<< local.aww:3:exp_file_content
+The host name is encrypt.
+The secret password is this_is_a_secret.