Add upload and download functions

easyssh.go

@@ -16,6 +16,7 @@ import (
 
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
+	"github.com/pkg/sftp"
 )
 
 // Contains main authority information.
@@ -56,7 +57,7 @@ func getKeyFile(keypath string) (ssh.Signer, error) {
 }
 
 // connects to remote server using MakeConfig struct and returns *ssh.Session
-func (ssh_conf *MakeConfig) connect() (*ssh.Session, error) {
+func (ssh_conf *MakeConfig) connect() (*ssh.Client, *ssh.Session, error) {
 	// auths holds the detected ssh auth methods
 	auths := []ssh.AuthMethod{}
 
@@ -81,23 +82,23 @@ func (ssh_conf *MakeConfig) connect() (*ssh.Session, error) {
 
 	client, err := ssh.Dial("tcp", ssh_conf.Server+":"+ssh_conf.Port, config)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	session, err := client.NewSession()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return session, nil
+	return client, session, nil
 }
 
 // Stream returns one channel that combines the stdout and stderr of the command
 // as it is run on the remote machine, and another that sends true when the
 // command is done. The sessions and channels will then be closed.
 func (ssh_conf *MakeConfig) Stream(command string) (output chan string, done chan bool, err error) {
 	// connect to remote host
-	session, err := ssh_conf.connect()
+	_, session, err := ssh_conf.connect()
 	if err != nil {
 		return output, done, err
 	}
@@ -151,8 +152,8 @@ func (ssh_conf *MakeConfig) Run(command string) (outStr string, err error) {
 }
 
 // Scp uploads sourceFile to remote machine like native scp console app.
-func (ssh_conf *MakeConfig) Scp(sourceFile string) error {
-	session, err := ssh_conf.connect()
+func (ssh_conf *MakeConfig) Upload(sourceFile string, destinationDir string) error {
+	_, session, err := ssh_conf.connect()
 
 	if err != nil {
 		return err
@@ -188,9 +189,44 @@ func (ssh_conf *MakeConfig) Scp(sourceFile string) error {
 		}
 	}()
 
-	if err := session.Run(fmt.Sprintf("scp -t %s", targetFile)); err != nil {
+	if err := session.Run(fmt.Sprintf("cd %s && scp -t %s", destinationDir, targetFile)); err != nil {
 		return err
 	}
 
 	return nil
 }
+
+func (ssh_conf *MakeConfig) Download(sourceFile string, destinationFile string) error {
+	client, session, err := ssh_conf.connect()
+
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	// open an SFTP session over an existing ssh connection.
+	sftp, err := sftp.NewClient(client)
+	if err != nil {
+	    return err
+	}
+	defer sftp.Close()
+
+	// Open the source file
+	srcFile, err := sftp.Open(sourceFile)
+	if err != nil {
+	    return err
+	}
+	defer srcFile.Close()
+
+	// Create the destination file
+	dstFile, err := os.Create(destinationFile)
+	if err != nil {
+	    return err
+	}
+	defer dstFile.Close()
+
+	// Copy the file
+	srcFile.WriteTo(dstFile)
+
+	return err
+}