all: implement remote "#get!" and "#put!" with owner and mode

When script with magic command "#get!" or "#put!" executed using "play"
command, one can changes the owner and/or permission mode by setting
the user/group and permission bits after the magic command, for example,

  #get!user:group+0600 src dst

Will changes the owner of dst in local into "user:group" with permission
"0600", while

  #put!user:group+0600 src dst

Will changes the owner of dst in remote into "user:group" with permission
"0600".

6 files changed, 303 insertions, 7 deletions

diff --git a/awwan_play_test.go b/awwan_play_test.go
index 7bcc827..0894c3e 100644
--- a/awwan_play_test.go
+++ b/awwan_play_test.go
@@ -16,9 +16,11 @@ import (
 
 type testCaseGetPut struct {
 	desc      string
-	fileDest  string
 	lineRange string
 
+	sudoPass string
+	fileDest string
+
 	expContent string
 	expError   string
 
@@ -190,3 +192,181 @@ func TestAwwan_Play_Put(t *testing.T) {
 		test.Assert(t, `mode`, c.expMode, fi.Mode().Perm())
 	}
 }
+
+func TestAwwan_Play_SudoGet(t *testing.T) {
+	var (
+		baseDir    = `testdata/play`
+		scriptDir  = filepath.Join(baseDir, `awwanssh.test`)
+		scriptFile = filepath.Join(scriptDir, `get.aww`)
+
+		tdata *test.Data
+		aww   *Awwan
+		err   error
+	)
+
+	tdata, err = test.LoadData(filepath.Join(scriptDir, `get_test.data`))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	aww, err = New(baseDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var cases = []testCaseGetPut{{
+		desc:       `WithMode`,
+		lineRange:  `14`,
+		sudoPass:   "awwan\n",
+		fileDest:   filepath.Join(scriptDir, `tmp`, `sudoget_with_mode.txt`),
+		expContent: string(tdata.Output[`/etc/crypttab`]),
+		expMode:    0601,
+	}, {
+		desc:       `WithOwner`,
+		lineRange:  `16`,
+		sudoPass:   "awwan\n",
+		fileDest:   filepath.Join(scriptDir, `tmp`, `sudoget_with_owner.txt`),
+		expContent: string(tdata.Output[`/etc/crypttab`]),
+		expMode:    420,
+	}, {
+		desc:       `WithOwnerAndMode`,
+		lineRange:  `18`,
+		sudoPass:   "awwan\nawwan\n",
+		fileDest:   filepath.Join(scriptDir, `tmp`, `sudoget_with_owner_mode.txt`),
+		expContent: string(tdata.Output[`/etc/crypttab`]),
+		expMode:    0602,
+	}}
+
+	var (
+		mockin = &mockStdin{}
+
+		c          testCaseGetPut
+		fi         os.FileInfo
+		gotContent []byte
+	)
+
+	for _, c = range cases {
+		t.Log(c.desc)
+
+		if len(c.fileDest) != 0 {
+			_ = os.Remove(c.fileDest)
+		}
+
+		var req = NewRequest(CommandModePlay, scriptFile, c.lineRange)
+
+		// Mock the request stdin to read password from buffer.
+		mockin.buf.Reset()
+		mockin.buf.WriteString(c.sudoPass)
+		req.stdin = mockin
+
+		err = aww.Play(req)
+		if err != nil {
+			test.Assert(t, `play error`, c.expError, err.Error())
+		}
+
+		if len(c.fileDest) == 0 {
+			continue
+		}
+
+		// File successfully copied but maybe error when setting
+		// owner or permission.
+
+		gotContent, err = os.ReadFile(c.fileDest)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, `content`, c.expContent, string(gotContent))
+
+		fi, err = os.Stat(c.fileDest)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, `mode`, c.expMode, fi.Mode().Perm())
+	}
+}
+
+func TestAwwan_Play_SudoPut(t *testing.T) {
+	var (
+		baseDir    = `testdata/play`
+		scriptDir  = filepath.Join(baseDir, `awwanssh.test`)
+		scriptFile = filepath.Join(scriptDir, `put.aww`)
+
+		tdata *test.Data
+		aww   *Awwan
+		err   error
+	)
+
+	tdata, err = test.LoadData(filepath.Join(scriptDir, `put_test.data`))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	aww, err = New(baseDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var cases = []testCaseGetPut{{
+		desc:       `WithMode`,
+		lineRange:  `14`,
+		fileDest:   `/home/awwanssh/sudoput_with_mode.txt`,
+		expContent: string(tdata.Output[`plain.txt`]),
+		expMode:    0604,
+	}, {
+		desc:       `WithOwner`,
+		lineRange:  `16`,
+		fileDest:   `/home/awwanssh/sudoput_with_owner.txt`,
+		expContent: string(tdata.Output[`plain.txt`]),
+		expMode:    0600,
+	}, {
+		desc:       `WithOwnerAndMode`,
+		lineRange:  `18`,
+		fileDest:   `/home/awwanssh/sudoput_with_owner_mode.txt`,
+		expContent: string(tdata.Output[`plain.txt`]),
+		expMode:    0602,
+	}}
+
+	var (
+		c          testCaseGetPut
+		fi         os.FileInfo
+		gotContent []byte
+	)
+
+	for _, c = range cases {
+		t.Log(c.desc)
+
+		if len(c.fileDest) != 0 {
+			_ = os.Remove(c.fileDest)
+		}
+
+		var req = NewRequest(CommandModePlay, scriptFile, c.lineRange)
+
+		err = aww.Play(req)
+		if err != nil {
+			test.Assert(t, `play error`, c.expError, err.Error())
+		}
+
+		if len(c.fileDest) == 0 {
+			continue
+		}
+
+		// File successfully copied but maybe error when setting
+		// owner or permission.
+
+		gotContent, err = os.ReadFile(c.fileDest)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, `content`, c.expContent, string(gotContent))
+
+		fi, err = os.Stat(c.fileDest)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, `mode`, c.expMode, fi.Mode().Perm())
+	}
+}
diff --git a/session.go b/session.go
index bab3607..90c951d 100644
--- a/session.go
+++ b/session.go
@@ -306,18 +306,37 @@ func (ses *Session) SudoCopy(req *Request, stmt *Statement) (err error) {
 
 // SudoGet copy file from remote that can be accessed by root on remote, to
 // local.
-func (ses *Session) SudoGet(stmt *Statement) (err error) {
-	var logp = `SudoGet`
+// If the owner and mode is set, it will also changes using sudo.
+func (ses *Session) SudoGet(req *Request, stmt *Statement) (err error) {
+	var (
+		logp = `SudoGet`
+		src  = stmt.args[0]
+		dst  = stmt.args[1]
+	)
 
-	err = ses.sshc.sudoGet(stmt.args[0], stmt.args[1])
+	err = ses.sshc.sudoGet(src, dst)
 	if err != nil {
 		return fmt.Errorf("%s: %w", logp, err)
 	}
+
+	if stmt.mode != 0 {
+		err = ses.localSudoChmod(req, dst, stmt.mode)
+		if err != nil {
+			return fmt.Errorf(`%s: %w`, logp, err)
+		}
+	}
+	if len(stmt.owner) != 0 {
+		err = ses.localSudoChown(req, dst, stmt.owner)
+		if err != nil {
+			return fmt.Errorf(`%s: %w`, logp, err)
+		}
+	}
+
 	return nil
 }
 
 // SudoPut copy file from local to remote using sudo.
-func (ses *Session) SudoPut(stmt *Statement) (err error) {
+func (ses *Session) SudoPut(req *Request, stmt *Statement) (err error) {
 	var (
 		logp = `SudoPut`
 		src  = stmt.args[0]
@@ -342,6 +361,19 @@ func (ses *Session) SudoPut(stmt *Statement) (err error) {
 		return fmt.Errorf("%s: %w", logp, err)
 	}
 
+	if stmt.mode != 0 {
+		err = ses.sshc.sudoChmod(dst, stmt.mode)
+		if err != nil {
+			return fmt.Errorf(`%s: %w`, logp, err)
+		}
+	}
+	if len(stmt.owner) != 0 {
+		err = ses.sshc.sudoChown(dst, stmt.owner)
+		if err != nil {
+			return fmt.Errorf(`%s: %w`, logp, err)
+		}
+	}
+
 	return nil
 }
 
@@ -475,9 +507,9 @@ func (ses *Session) executeScriptOnRemote(req *Request, pos linePosition) (err e
 		case statementKindPut:
 			err = ses.Put(stmt)
 		case statementKindSudoGet:
-			err = ses.SudoGet(stmt)
+			err = ses.SudoGet(req, stmt)
 		case statementKindSudoPut:
-			err = ses.SudoPut(stmt)
+			err = ses.SudoPut(req, stmt)
 		}
 		if err != nil {
 			fmt.Fprintf(req.stderr, "!!! %s\n", err)
@@ -705,6 +737,40 @@ func (ses *Session) loadRawEnv(content []byte) (err error) {
 	return nil
 }
 
+// localSudoChmod change the file permission in local environment using
+// sudo.
+func (ses *Session) localSudoChmod(req *Request, file string, mode fs.FileMode) (err error) {
+	var (
+		fsmode    = strconv.FormatUint(uint64(mode), 8)
+		sudoChmod = &Statement{
+			kind: statementKindDefault,
+			cmd:  `sudo`,
+			args: []string{`chmod`, fsmode, file},
+			raw:  []byte(fmt.Sprintf(`sudo chmod %o %q`, mode, file)),
+		}
+	)
+	err = ExecLocal(req, sudoChmod)
+	if err != nil {
+		return fmt.Errorf(`%s: %w`, sudoChmod.raw, err)
+	}
+	return nil
+}
+
+// localSudoChown change the file owner in local environment using sudo.
+func (ses *Session) localSudoChown(req *Request, file, owner string) (err error) {
+	var sudoChown = &Statement{
+		kind: statementKindDefault,
+		cmd:  `sudo`,
+		args: []string{`chown`, owner, file},
+		raw:  []byte(fmt.Sprintf(`sudo chown %s %q`, owner, file)),
+	}
+	err = ExecLocal(req, sudoChown)
+	if err != nil {
+		return fmt.Errorf(`%s: %w`, sudoChown.raw, err)
+	}
+	return nil
+}
+
 // render apply the session and environment variables into input stream `in`
 // and return the result.
 // It will return an error if the input cannot be parsed or one variable
diff --git a/ssh_client.go b/ssh_client.go
index 692f050..786417c 100644
--- a/ssh_client.go
+++ b/ssh_client.go
@@ -155,6 +155,29 @@ func (sshc *sshClient) rmdirAll(dir string) {
 	}
 }
 
+// sudoChmod change the permission of remoteFile using sudo.
+func (sshc *sshClient) sudoChmod(remoteFile string, mode fs.FileMode) (err error) {
+	var cmd = fmt.Sprintf(`sudo chmod %o %q`, mode, remoteFile)
+
+	err = sshc.conn.Execute(cmd)
+	if err != nil {
+		return err
+	}
+	return nil
+
+}
+
+// sudoChown change the owner of remoteFile using sudo.
+func (sshc *sshClient) sudoChown(remoteFile, owner string) (err error) {
+	var cmd = fmt.Sprintf(`sudo chown %s %q`, owner, remoteFile)
+
+	err = sshc.conn.Execute(cmd)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 // sudoGet copy the remote file using sudo to local.
 // The remote file is copied to temporary directory first, chmod-ed to
 // current SSH user so it can be read.
diff --git a/testdata/play/awwanssh.test/get.aww b/testdata/play/awwanssh.test/get.aww
index c10ae27..4b5d604 100644
--- a/testdata/play/awwanssh.test/get.aww
+++ b/testdata/play/awwanssh.test/get.aww
@@ -10,3 +10,9 @@
 ## privileged.
 
 #get:awwan:bin /etc/os-release {{.ScriptDir}}/tmp/get_with_owner.txt
+
+#get!+0601 /etc/crypttab {{.ScriptDir}}/tmp/sudoget_with_mode.txt
+
+#get!awwan:bin /etc/crypttab {{.ScriptDir}}/tmp/sudoget_with_owner.txt
+
+#get!awwan:bin+0602 /etc/crypttab {{.ScriptDir}}/tmp/sudoget_with_owner_mode.txt
diff --git a/testdata/play/awwanssh.test/get_test.data b/testdata/play/awwanssh.test/get_test.data
index ad4577d..aec7acb 100644
--- a/testdata/play/awwanssh.test/get_test.data
+++ b/testdata/play/awwanssh.test/get_test.data
@@ -18,3 +18,18 @@ LOGO=archlinux-logo
 
 <<< WithOwner:error
 Play: Get: chown awwan:bin "/home/awwan/src/testdata/play/awwanssh.test/tmp/get_with_owner.txt": exit status 1
+
+<<< /etc/crypttab
+# Configuration for encrypted block devices.
+# See crypttab(5) for details.
+
+# NOTE: Do not list your root (/) partition here, it must be set up
+#       beforehand by the initramfs (/etc/mkinitcpio.conf).
+
+# <name>       <device>                                     <password>              <options>
+# home         UUID=b8ad5c18-f445-495d-9095-c9ec4f9d2f37    /etc/mypassword1
+# data1        /dev/sda3                                    /etc/mypassword2
+# data2        /dev/sda5                                    /etc/cryptfs.key
+# swap         /dev/sdx4                                    /dev/urandom            swap,cipher=aes-cbc-essiv:sha256,size=256
+# vol          /dev/sdb7                                    none
+
diff --git a/testdata/play/awwanssh.test/put.aww b/testdata/play/awwanssh.test/put.aww
index 9779ed2..e419a4f 100644
--- a/testdata/play/awwanssh.test/put.aww
+++ b/testdata/play/awwanssh.test/put.aww
@@ -10,3 +10,9 @@
 ## privileged.
 
 #put:awwan:bin+666 {{.ScriptDir}}/plain.txt put_with_owner.txt
+
+#put!+604 {{.ScriptDir}}/plain.txt sudoput_with_mode.txt
+
+#put!awwan:bin {{.ScriptDir}}/plain.txt sudoput_with_owner.txt
+
+#put!awwan:bin+602 {{.ScriptDir}}/plain.txt sudoput_with_owner_mode.txt