Implement download for WinRM communicator
Signed-off-by: Mark DeLillo <mdelillo@pivotal.io>
This commit is contained in:
parent
9f992b8f80
commit
0614fc234f
|
@ -1,8 +1,10 @@
|
||||||
package winrm
|
package winrm
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"encoding/base64"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"io/ioutil"
|
||||||
"log"
|
"log"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
@ -143,7 +145,20 @@ func (c *Communicator) UploadDir(dst string, src string, exclude []string) error
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Communicator) Download(src string, dst io.Writer) error {
|
func (c *Communicator) Download(src string, dst io.Writer) error {
|
||||||
return fmt.Errorf("WinRM doesn't support download.")
|
endpoint := winrm.NewEndpoint(c.endpoint.Host, c.endpoint.Port, c.config.Https, c.config.Insecure, nil, nil, nil, c.config.Timeout)
|
||||||
|
client, err := winrm.NewClient(endpoint, c.config.Username, c.config.Password)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
encodeScript := `$file=[System.IO.File]::ReadAllBytes("%s"); Write-Output $([System.Convert]::ToBase64String($file))`
|
||||||
|
|
||||||
|
base64DecodePipe := &Base64Pipe{w: dst}
|
||||||
|
|
||||||
|
cmd := winrm.Powershell(fmt.Sprintf(encodeScript, src))
|
||||||
|
_, err = client.Run(cmd, base64DecodePipe, ioutil.Discard)
|
||||||
|
|
||||||
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Communicator) DownloadDir(src string, dst string, exclude []string) error {
|
func (c *Communicator) DownloadDir(src string, dst string, exclude []string) error {
|
||||||
|
@ -164,3 +179,33 @@ func (c *Communicator) newCopyClient() (*winrmcp.Winrmcp, error) {
|
||||||
TransportDecorator: c.config.TransportDecorator,
|
TransportDecorator: c.config.TransportDecorator,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type Base64Pipe struct {
|
||||||
|
w io.Writer // underlying writer (file, buffer)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *Base64Pipe) ReadFrom(r io.Reader) (int64, error) {
|
||||||
|
b, err := ioutil.ReadAll(r)
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
var i int
|
||||||
|
i, err = d.Write(b)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return int64(i), err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *Base64Pipe) Write(p []byte) (int, error) {
|
||||||
|
dst := make([]byte, base64.StdEncoding.DecodedLen(len(p)))
|
||||||
|
|
||||||
|
if _, err := base64.StdEncoding.Decode(dst, p); err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return d.w.Write(dst)
|
||||||
|
}
|
||||||
|
|
|
@ -3,6 +3,7 @@ package winrm
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"io"
|
"io"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
@ -10,6 +11,9 @@ import (
|
||||||
"github.com/hashicorp/packer/packer"
|
"github.com/hashicorp/packer/packer"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const PAYLOAD = "stuff"
|
||||||
|
const BASE64_ENCODED_PAYLOAD = "c3R1ZmY="
|
||||||
|
|
||||||
func newMockWinRMServer(t *testing.T) *winrmtest.Remote {
|
func newMockWinRMServer(t *testing.T) *winrmtest.Remote {
|
||||||
wrm := winrmtest.NewRemote()
|
wrm := winrmtest.NewRemote()
|
||||||
|
|
||||||
|
@ -26,9 +30,16 @@ func newMockWinRMServer(t *testing.T) *winrmtest.Remote {
|
||||||
return 0
|
return 0
|
||||||
})
|
})
|
||||||
|
|
||||||
|
wrm.CommandFunc(
|
||||||
|
winrmtest.MatchPattern(`^echo `+BASE64_ENCODED_PAYLOAD+` >> ".*"$`),
|
||||||
|
func(out, err io.Writer) int {
|
||||||
|
return 0
|
||||||
|
})
|
||||||
|
|
||||||
wrm.CommandFunc(
|
wrm.CommandFunc(
|
||||||
winrmtest.MatchPattern(`^powershell.exe -EncodedCommand .*$`),
|
winrmtest.MatchPattern(`^powershell.exe -EncodedCommand .*$`),
|
||||||
func(out, err io.Writer) int {
|
func(out, err io.Writer) int {
|
||||||
|
out.Write([]byte(BASE64_ENCODED_PAYLOAD))
|
||||||
return 0
|
return 0
|
||||||
})
|
})
|
||||||
|
|
||||||
|
@ -86,9 +97,21 @@ func TestUpload(t *testing.T) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("error creating communicator: %s", err)
|
t.Fatalf("error creating communicator: %s", err)
|
||||||
}
|
}
|
||||||
|
file := "C:/Temp/packer.cmd"
|
||||||
err = c.Upload("C:/Temp/packer.cmd", bytes.NewReader([]byte("something")), nil)
|
err = c.Upload(file, strings.NewReader(PAYLOAD), nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("error uploading file: %s", err)
|
t.Fatalf("error uploading file: %s", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
dest := new(bytes.Buffer)
|
||||||
|
err = c.Download(file, dest)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("error downloading file: %s", err)
|
||||||
|
}
|
||||||
|
downloadedPayload := strings.TrimRight(dest.String(), "\x00")
|
||||||
|
|
||||||
|
if downloadedPayload != PAYLOAD {
|
||||||
|
t.Fatalf("files are not equal: expected [%s] length: %v, got [%s] length %v", PAYLOAD, len(PAYLOAD), downloadedPayload, len(downloadedPayload))
|
||||||
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
|
@ -5,6 +5,7 @@ import (
|
||||||
"crypto/x509"
|
"crypto/x509"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"sync"
|
||||||
|
|
||||||
"github.com/masterzen/winrm/soap"
|
"github.com/masterzen/winrm/soap"
|
||||||
)
|
)
|
||||||
|
@ -114,10 +115,21 @@ func (c *Client) Run(command string, stdout io.Writer, stderr io.Writer) (int, e
|
||||||
return 1, err
|
return 1, err
|
||||||
}
|
}
|
||||||
|
|
||||||
go io.Copy(stdout, cmd.Stdout)
|
var wg sync.WaitGroup
|
||||||
go io.Copy(stderr, cmd.Stderr)
|
wg.Add(2)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
io.Copy(stdout, cmd.Stdout)
|
||||||
|
}()
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
io.Copy(stderr, cmd.Stderr)
|
||||||
|
}()
|
||||||
|
|
||||||
cmd.Wait()
|
cmd.Wait()
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
return cmd.ExitCode(), cmd.err
|
return cmd.ExitCode(), cmd.err
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue