terraform/vendor/github.com/packer-community/winrmcp/winrmcp/cp.go

251 lines
5.4 KiB
Go

package winrmcp
import (
"encoding/base64"
"fmt"
"io"
"log"
"os"
"sync"
"github.com/masterzen/winrm"
"github.com/nu7hatch/gouuid"
)
func doCopy(client *winrm.Client, config *Config, in io.Reader, toPath string) error {
tempFile, err := tempFileName()
if err != nil {
return fmt.Errorf("Error generating unique filename: %v", err)
}
tempPath := "$env:TEMP\\" + tempFile
if os.Getenv("WINRMCP_DEBUG") != "" {
log.Printf("Copying file to %s\n", tempPath)
}
err = uploadContent(client, config.MaxOperationsPerShell, "%TEMP%\\"+tempFile, in)
if err != nil {
return fmt.Errorf("Error uploading file to %s: %v", tempPath, err)
}
if os.Getenv("WINRMCP_DEBUG") != "" {
log.Printf("Moving file from %s to %s", tempPath, toPath)
}
err = restoreContent(client, tempPath, toPath)
if err != nil {
return fmt.Errorf("Error restoring file from %s to %s: %v", tempPath, toPath, err)
}
if os.Getenv("WINRMCP_DEBUG") != "" {
log.Printf("Removing temporary file %s", tempPath)
}
err = cleanupContent(client, tempPath)
if err != nil {
return fmt.Errorf("Error removing temporary file %s: %v", tempPath, err)
}
return nil
}
func uploadContent(client *winrm.Client, maxChunks int, filePath string, reader io.Reader) error {
var err error
done := false
for !done {
done, err = uploadChunks(client, filePath, maxChunks, reader)
if err != nil {
return err
}
}
return nil
}
func uploadChunks(client *winrm.Client, filePath string, maxChunks int, reader io.Reader) (bool, error) {
shell, err := client.CreateShell()
if err != nil {
return false, fmt.Errorf("Couldn't create shell: %v", err)
}
defer shell.Close()
// Upload the file in chunks to get around the Windows command line size limit.
// Base64 encodes each set of three bytes into four bytes. In addition the output
// is padded to always be a multiple of four.
//
// ceil(n / 3) * 4 = m1 - m2
//
// where:
// n = bytes
// m1 = max (8192 character command limit.)
// m2 = len(filePath)
chunkSize := ((8000 - len(filePath)) / 4) * 3
chunk := make([]byte, chunkSize)
if maxChunks == 0 {
maxChunks = 1
}
for i := 0; i < maxChunks; i++ {
n, err := reader.Read(chunk)
if err != nil && err != io.EOF {
return false, err
}
if n == 0 {
return true, nil
}
content := base64.StdEncoding.EncodeToString(chunk[:n])
if err = appendContent(shell, filePath, content); err != nil {
return false, err
}
}
return false, nil
}
func restoreContent(client *winrm.Client, fromPath, toPath string) error {
shell, err := client.CreateShell()
if err != nil {
return err
}
defer shell.Close()
script := fmt.Sprintf(`
$tmp_file_path = [System.IO.Path]::GetFullPath("%s")
$dest_file_path = [System.IO.Path]::GetFullPath("%s".Trim("'"))
if (Test-Path $dest_file_path) {
if (Test-Path -Path $dest_file_path -PathType container) {
Exit 1
} else {
rm $dest_file_path
}
}
else {
$dest_dir = ([System.IO.Path]::GetDirectoryName($dest_file_path))
New-Item -ItemType directory -Force -ErrorAction SilentlyContinue -Path $dest_dir | Out-Null
}
if (Test-Path $tmp_file_path) {
$reader = [System.IO.File]::OpenText($tmp_file_path)
$writer = [System.IO.File]::OpenWrite($dest_file_path)
try {
for(;;) {
$base64_line = $reader.ReadLine()
if ($base64_line -eq $null) { break }
$bytes = [System.Convert]::FromBase64String($base64_line)
$writer.write($bytes, 0, $bytes.Length)
}
}
finally {
$reader.Close()
$writer.Close()
}
} else {
echo $null > $dest_file_path
}
`, fromPath, toPath)
cmd, err := shell.Execute(winrm.Powershell(script))
if err != nil {
return err
}
defer cmd.Close()
var wg sync.WaitGroup
copyFunc := func(w io.Writer, r io.Reader) {
defer wg.Done()
io.Copy(w, r)
}
wg.Add(2)
go copyFunc(os.Stdout, cmd.Stdout)
go copyFunc(os.Stderr, cmd.Stderr)
cmd.Wait()
wg.Wait()
if cmd.ExitCode() != 0 {
return fmt.Errorf("restore operation returned code=%d", cmd.ExitCode())
}
return nil
}
func cleanupContent(client *winrm.Client, filePath string) error {
shell, err := client.CreateShell()
if err != nil {
return err
}
defer shell.Close()
script := fmt.Sprintf(`
$tmp_file_path = [System.IO.Path]::GetFullPath("%s")
if (Test-Path $tmp_file_path) {
Remove-Item $tmp_file_path -ErrorAction SilentlyContinue
}
`, filePath)
cmd, err := shell.Execute(winrm.Powershell(script))
if err != nil {
return err
}
defer cmd.Close()
var wg sync.WaitGroup
copyFunc := func(w io.Writer, r io.Reader) {
defer wg.Done()
io.Copy(w, r)
}
wg.Add(2)
go copyFunc(os.Stdout, cmd.Stdout)
go copyFunc(os.Stderr, cmd.Stderr)
cmd.Wait()
wg.Wait()
if cmd.ExitCode() != 0 {
return fmt.Errorf("cleanup operation returned code=%d", cmd.ExitCode())
}
return nil
}
func appendContent(shell *winrm.Shell, filePath, content string) error {
cmd, err := shell.Execute(fmt.Sprintf("echo %s >> \"%s\"", content, filePath))
if err != nil {
return err
}
defer cmd.Close()
var wg sync.WaitGroup
copyFunc := func(w io.Writer, r io.Reader) {
defer wg.Done()
io.Copy(w, r)
}
wg.Add(2)
go copyFunc(os.Stdout, cmd.Stdout)
go copyFunc(os.Stderr, cmd.Stderr)
cmd.Wait()
wg.Wait()
if cmd.ExitCode() != 0 {
return fmt.Errorf("upload operation returned code=%d", cmd.ExitCode())
}
return nil
}
func tempFileName() (string, error) {
uniquePart, err := uuid.NewV4()
if err != nil {
return "", err
}
return fmt.Sprintf("winrmcp-%s.tmp", uniquePart), nil
}