diff --git a/main.go b/main.go index 95fa49d..c67bdd6 100644 --- a/main.go +++ b/main.go @@ -51,6 +51,7 @@ func runMain() error { cacert := flags.String("cacert", "", "ca certificate to validate against") opTimeout := flags.Duration("op-timeout", time.Second*60, "operation timeout") maxOpsPerShell := flags.Int("max-ops-per-shell", 15, "max operations per shell") + maxShells := flags.Int("max-shells", 10, "max shells per user") flags.Parse(os.Args[1:]) var certBytes []byte @@ -72,6 +73,7 @@ func runMain() error { CACertBytes: certBytes, OperationTimeout: *opTimeout, MaxOperationsPerShell: *maxOpsPerShell, + MaxShells: *maxShells, }) if err != nil { return err diff --git a/winrmcp/cp.go b/winrmcp/cp.go index 5599b6d..5429d2c 100644 --- a/winrmcp/cp.go +++ b/winrmcp/cp.go @@ -21,7 +21,7 @@ func doCopy(client *winrm.Client, config *Config, in io.Reader, toPath string) e log.Printf("Copying file to %s\n", tempPath) } - err := uploadContent(client, config.MaxOperationsPerShell, "%TEMP%\\"+tempFile, in) + err := uploadContent(client, config.MaxShells, config.MaxOperationsPerShell, "%TEMP%\\"+tempFile, in) if err != nil { return errors.New(fmt.Sprintf("Error uploading file to %s: %v", tempPath, err)) } @@ -47,18 +47,21 @@ func doCopy(client *winrm.Client, config *Config, in io.Reader, toPath string) e return nil } -func uploadContent(client *winrm.Client, maxChunks int, filePath string, reader io.Reader) error { +func uploadContent(client *winrm.Client, maxShell int, maxChunks int, filePath string, reader io.Reader) error { var err error var piece = 0 var wg sync.WaitGroup - parallel := 4 if maxChunks == 0 { maxChunks = 1 } + if maxChunks == 0 { + maxChunks = 10 + } + // Create 4 Parallel workers - for p := 0; p < parallel; p++ { + for p := 0; p < maxShell; p++ { done := make(chan bool, 1) // Add worker to the WaitGroup wg.Add(1) diff --git a/winrmcp/winrmcp.go b/winrmcp/winrmcp.go index 844c5a2..1d5b088 100644 --- a/winrmcp/winrmcp.go +++ b/winrmcp/winrmcp.go @@ -24,6 +24,7 @@ type Config struct { CACertBytes []byte OperationTimeout time.Duration MaxOperationsPerShell int + MaxShells int } type Auth struct {