support WriteFile protocol

.gitignore

@@ -0,0 +1 @@
+.idea

go.mod

@@ -0,0 +1,9 @@
+module github.com/ImSingee/scp
+
+go 1.14
+
+require (
+	github.com/ImSingee/tt v1.0.2
+	github.com/alessio/shellescape v1.4.1
+	golang.org/x/crypto v0.0.0-20201208171446-5f87f3452ae9
+)

go.sum

@@ -0,0 +1,14 @@
+github.com/ImSingee/tt v1.0.2 h1:t6EWUfy6rB9fuBH3U4OaYcZ17hKa4TbYU211jCnD9fA=
+github.com/ImSingee/tt v1.0.2/go.mod h1:7O7v+cIBruYWGFObw85DDjH0gNLquen1cJqMR3GOgxw=
+github.com/alessio/shellescape v1.4.1 h1:V7yhSDDn8LP4lc4jS8pFkt0zCnzVJlG5JXy9BVKJUX0=
+github.com/alessio/shellescape v1.4.1/go.mod h1:PZAiSCk0LJaZkiCSkPv8qIobYglO3FPpyFjDCtHLS30=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20201208171446-5f87f3452ae9 h1:sYNJzB4J8toYPQTM6pAkcmBRgw9SnQKP9oXCHfgy604=
+golang.org/x/crypto v0.0.0-20201208171446-5f87f3452ae9/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/term v0.0.0-20201117132131-f5c789dd3221 h1:/ZHdbVpdR/jk3g30/d4yUL0JU9kksj8+F/bnQUVLGDM=
+golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

protocol.go

@@ -0,0 +1,125 @@
+package scp
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"github.com/alessio/shellescape"
+	"golang.org/x/crypto/ssh"
+	"io"
+	"strconv"
+	"strings"
+)
+
+type RemoteClient struct {
+	session *ssh.Session
+
+	stdout io.Reader
+	stdin  io.WriteCloser
+}
+
+func NewClient(session *ssh.Session) (*RemoteClient, error) {
+	stdout, err := session.StdoutPipe()
+
+	if err != nil {
+		return nil, fmt.Errorf("cannot catch stdout: %w", err)
+	}
+
+	stdin, err := session.StdinPipe()
+
+	if err != nil {
+		return nil, fmt.Errorf("cannot set stdin: %w", err)
+	}
+
+	return &RemoteClient{session: session, stdout: stdout, stdin: stdin}, nil
+}
+
+func (c *RemoteClient) Start(filename string, isDirectory bool) error {
+	commandBuilder := strings.Builder{}
+
+	commandBuilder.WriteString("scp -t")
+
+	if isDirectory {
+		commandBuilder.WriteString(" -d")
+	}
+
+	commandBuilder.WriteByte(' ')
+
+	commandBuilder.WriteString(shellescape.Quote(filename))
+
+	command := commandBuilder.String()
+
+	return c.session.Start(command)
+}
+
+func (c *RemoteClient) checkResponse() error {
+	buffer := make([]byte, 1)
+
+	_, err := c.stdout.Read(buffer)
+	if err != nil {
+		return fmt.Errorf("cannot read from remote stdout: %w", err)
+	}
+
+	if buffer[0] == 0 {
+		return nil
+	}
+
+	errorMessageBuilder := strings.Builder{}
+
+	if buffer[0] == 1 {
+		errorMessageBuilder.WriteString("Warning (1): ")
+	} else if buffer[0] == 2 {
+		errorMessageBuilder.WriteString("Error (2): ")
+	} else {
+		errorMessageBuilder.WriteString("UnknownError (")
+		errorMessageBuilder.WriteString(strconv.Itoa(int(buffer[0])))
+		errorMessageBuilder.WriteString("): ")
+	}
+
+	reader := bufio.NewReader(c.stdout)
+	all, err := reader.ReadString('\n')
+
+	if err != nil {
+		errorMessageBuilder.WriteString("[ERROR during get reason] (error: ")
+		errorMessageBuilder.WriteString(err.Error())
+		errorMessageBuilder.WriteString(")")
+	} else {
+		errorMessageBuilder.WriteString(all)
+	}
+
+	return errors.New(errorMessageBuilder.String())
+}
+
+func (c *RemoteClient) WriteFile(perm string, size int64, filename string, data io.Reader) error {
+	err := c.checkResponse()
+	if err != nil {
+		return err
+	}
+
+	_, err = fmt.Fprintln(c.stdin, "C"+perm, size, filename)
+	if err != nil {
+		return err
+	}
+
+	err = c.checkResponse()
+	if err != nil {
+		return err
+	}
+
+	_, err = io.Copy(c.stdin, data)
+	if err != nil {
+		return err
+	}
+
+	_, err = c.stdin.Write([]byte{0})
+	if err != nil {
+		return err
+	}
+
+	err = c.checkResponse()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

protocol_test.go

@@ -0,0 +1,125 @@
+package scp
+
+import (
+	"fmt"
+	"github.com/ImSingee/tt"
+	"github.com/alessio/shellescape"
+	"golang.org/x/crypto/ssh"
+	"os"
+	"strings"
+	"testing"
+)
+
+// Underlying client cannot be closed
+func getSshSession(t *testing.T) *ssh.Session {
+	client, err := ssh.Dial("tcp", "localhost:22", &ssh.ClientConfig{
+		User: "root",
+		Auth: []ssh.AuthMethod{
+			ssh.Password(os.Getenv("SSH_PASSWORD")),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		BannerCallback:  ssh.BannerDisplayStderr(),
+	})
+	tt.AssertIsNil(t, err)
+
+	session, err := client.NewSession()
+	tt.AssertIsNil(t, err)
+
+	return session
+}
+
+func reset(t *testing.T) {
+	session := getSshSession(t)
+	defer session.Close()
+
+	err := session.Run("rm -rf /test; mkdir /test")
+	tt.AssertIsNil(t, err)
+}
+
+func checkFileExist(t *testing.T, filename string) bool {
+	session := getSshSession(t)
+	defer session.Close()
+
+	output, err := session.Output("stat " + shellescape.Quote(filename))
+
+	if err != nil {
+		return false
+	}
+
+	return len(output) != 0
+}
+
+func readFile(t *testing.T, filename string) []byte {
+	session := getSshSession(t)
+	defer session.Close()
+
+	output, err := session.Output("cat " + shellescape.Quote(filename))
+	tt.AssertIsNil(t, err)
+
+	return output
+}
+
+func TestWriteFileToFile(t *testing.T) {
+	reset(t)
+
+	session := getSshSession(t)
+	defer session.Close()
+
+	client, err := NewClient(session)
+	tt.AssertIsNil(t, err)
+
+	err = client.Start("/test/some", false)
+	tt.AssertIsNil(t, err)
+
+	// filename (test01) can be difference from real (some)
+	err = client.WriteFile("0644", 6, "test01", strings.NewReader("hahaha"))
+	tt.AssertIsNil(t, err)
+
+	tt.AssertTrue(t, checkFileExist(t, "/test/some"))
+
+	file := readFile(t, "/test/some")
+
+	tt.AssertEqual(t, "hahaha", string(file))
+}
+
+func TestWriteFileToFolder(t *testing.T) {
+	reset(t)
+
+	session := getSshSession(t)
+	defer session.Close()
+
+	client, err := NewClient(session)
+	tt.AssertIsNil(t, err)
+
+	err = client.Start("/test", false)
+	tt.AssertIsNil(t, err)
+
+	err = client.WriteFile("0644", 6, "test01", strings.NewReader("hahaha"))
+	tt.AssertIsNil(t, err)
+
+	tt.AssertTrue(t, checkFileExist(t, "/test/test01"))
+
+	file := readFile(t, "/test/test01")
+
+	tt.AssertEqual(t, "hahaha", string(file))
+}
+
+func TestWriteFileToUnExistFolder(t *testing.T) {
+	reset(t)
+
+	session := getSshSession(t)
+	defer session.Close()
+
+	client, err := NewClient(session)
+	tt.AssertIsNil(t, err)
+
+	err = client.Start("/test/not/exist/folder", false)
+	tt.AssertIsNil(t, err)
+
+	err = client.WriteFile("0644", 6, "test01", strings.NewReader("hahaha"))
+	tt.AssertIsNotNil(t, err)
+
+	fmt.Println(err)
+
+	tt.AssertFalse(t, checkFileExist(t, "/test/not/exist/folder"))
+}