Skip to content

Commit

Permalink
WinRM Script tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Valentin Aitken committed Oct 15, 2015
1 parent 5e07983 commit 36d8ed0
Show file tree
Hide file tree
Showing 9 changed files with 355 additions and 40 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,10 @@
import org.apache.brooklyn.core.location.AbstractLocation;
import org.apache.brooklyn.core.location.access.PortForwardManager;
import org.apache.brooklyn.core.mgmt.BrooklynTaskTags;
import org.apache.brooklyn.util.core.internal.winrm.NativeWindowsScriptRunner;
import org.apache.brooklyn.util.collections.MutableMap;
import org.apache.brooklyn.util.core.ResourceUtils;
import org.apache.brooklyn.util.core.internal.winrm.NaiveWindowsScriptRunner;
import org.apache.brooklyn.util.core.internal.winrm.WinRmScriptTool;
import org.apache.brooklyn.util.core.task.Tasks;
import org.apache.brooklyn.util.text.Strings;
import org.apache.commons.codec.binary.Base64;
Expand All @@ -59,13 +62,15 @@
import com.google.common.net.HostAndPort;
import com.google.common.reflect.TypeToken;

import java.util.List;

import org.apache.brooklyn.util.exceptions.Exceptions;
import org.apache.brooklyn.util.stream.Streams;

import io.cloudsoft.winrm4j.winrm.WinRmTool;
import io.cloudsoft.winrm4j.winrm.WinRmToolResponse;

public class WinRmMachineLocation extends AbstractLocation implements MachineLocation, NativeWindowsScriptRunner {
public class WinRmMachineLocation extends AbstractLocation implements MachineLocation, NaiveWindowsScriptRunner {

private static final Logger LOG = LoggerFactory.getLogger(WinRmMachineLocation.class);

Expand Down Expand Up @@ -177,6 +182,13 @@ public WinRmToolResponse executeCommand(Function<WinRmTool, WinRmToolResponse> c

@Override
public Integer executeNativeOrPsCommand(Map flags, String regularCommand, final String powerShellCommand, String summaryForLogging, Boolean allowNoOp) {
if (Strings.isBlank(regularCommand) && Strings.isBlank(powerShellCommand)) {
if (allowNoOp) {
return new WinRmToolResponse("", "", 0).getStatusCode();
} else {
throw new IllegalStateException(String.format("Exactly one of cmd or psCmd must be set for %s", summaryForLogging));
}
}
ByteArrayOutputStream stdIn = new ByteArrayOutputStream();
ByteArrayOutputStream stdOut = flags.get("out") != null ? (ByteArrayOutputStream)flags.get("out") : new ByteArrayOutputStream();
ByteArrayOutputStream stdErr = flags.get("err") != null ? (ByteArrayOutputStream)flags.get("err") : new ByteArrayOutputStream();
Expand Down Expand Up @@ -212,6 +224,59 @@ public Integer executeNativeOrPsCommand(Map flags, String regularCommand, final
return response.getStatusCode();
}

public int executeNativeOrPsScript(Map flags, List<String> regularCommands, final List<String> powerShellCommands, String summaryForLogging, Boolean allowNoOp) {
if ((regularCommands == null || regularCommands.size() == 0) && (powerShellCommands == null || powerShellCommands.size() == 0)) {
if (allowNoOp) {
return new WinRmToolResponse("", "", 0).getStatusCode();
} else {
throw new IllegalStateException(String.format("Exactly one of cmd or psCmd must be set for %s", summaryForLogging));
}
}

MutableMap<String, Object> scriptProps = MutableMap.<String, Object>of(WinRmScriptTool.PROP_SUMMARY.getName(), summaryForLogging);
WinRmScriptTool scriptTool = new WinRmScriptTool(scriptProps, this);

if (regularCommands == null || regularCommands.size() == 0) {
return scriptTool.execPsScript(flags, powerShellCommands);
} else {
return scriptTool.execNativeScript(flags, regularCommands);
}
}

public int copyResource(Map<Object, Object> flags, InputStream source, String target, boolean createParentDir) {
if (createParentDir) {
createDirectory(getDirectory(target), "Creating resource directory");
}
return copyTo(source, target);
}

public int copyResource(Map<Object, Object> flags, String source, String target, boolean createParentDir, ResourceUtils resource) {
if (createParentDir) {
createDirectory(getDirectory(target), "Creating resource directory");
}

InputStream stream = null;
try {
Tasks.setBlockingDetails("retrieving resource "+source+" for copying across");
stream = resource.getResourceFromUrl(source);
Tasks.setBlockingDetails("copying resource "+source+" to server");
return copyTo(stream, target);
} catch (Exception e) {
throw Exceptions.propagate(e);
} finally {
Tasks.setBlockingDetails(null);
if (stream != null) Streams.closeQuietly(stream);
}
}

public int createDirectory(String directoryName, String summaryForLogging) {
return executePsCommand("New-Item -path \"" + directoryName + "\" -type directory -ErrorAction SilentlyContinue").getStatusCode();
}

private String getDirectory(String fileName) {
return fileName.substring(0, fileName.lastIndexOf("\\"));
}

private void writeToStream(ByteArrayOutputStream stream, String string) {
try {
stream.write(string.getBytes());
Expand Down Expand Up @@ -274,6 +339,7 @@ public int copyTo(File source, String destination) {
}
}

@Override
public int copyTo(InputStream source, String destination) {
executePsCommand("rm -ErrorAction SilentlyContinue " + destination);
try {
Expand All @@ -288,9 +354,9 @@ public int copyTo(InputStream source, String destination) {
} else {
chunk = Arrays.copyOf(inputData, bytesRead);
}
executePsCommand(Joiner.on("\r\n").join(ImmutableList.of("If ((!(Test-Path " + destination + ")) -or ((Get-Item '" + destination + "').length -eq " +
expectedFileSize + ")) {Add-Content -Encoding Byte -path " + destination +
" -value ([System.Convert]::FromBase64String(\"" + new String(Base64.encodeBase64(chunk)) + "\"))}"), "copyFile"));
executePsCommand("If ((!(Test-Path " + destination + ")) -or ((Get-Item '" + destination + "').length -eq " +
expectedFileSize + ")) {Add-Content -Encoding Byte -path \"" + destination +
"\" -value ([System.Convert]::FromBase64String(\"" + new String(Base64.encodeBase64(chunk)) + "\"))}");
expectedFileSize += bytesRead;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ protected static Boolean hasVal(Map<String,?> map, ConfigKey<?> keyC) {
return map.containsKey(key);
}

protected static <T> T getMandatoryVal(Map<String,?> map, ConfigKey<T> keyC) {
public static <T> T getMandatoryVal(Map<String,?> map, ConfigKey<T> keyC) {
String key = keyC.getName();
checkArgument(map.containsKey(key), "must contain key '"+keyC+"'");
return TypeCoercions.coerce(map.get(key), keyC.getTypeToken());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,13 @@
*/
package org.apache.brooklyn.util.core.internal.winrm;

import java.io.InputStream;
import java.util.Map;

public interface NativeWindowsScriptRunner {
public interface NaiveWindowsScriptRunner {

/** Runs a command and returns the result code */
Integer executeNativeOrPsCommand(Map flags, String regularCommand, String powershellCommand, String summaryForLogging, Boolean allowNoOp);

int copyTo(InputStream source, String destination);
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.brooklyn.util.core.internal.winrm;

import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableMap;
import com.sun.xml.internal.messaging.saaj.util.ByteInputStream;
import org.apache.brooklyn.config.ConfigKey;
import org.apache.brooklyn.util.collections.MutableMap;
import org.apache.brooklyn.util.core.internal.ssh.ShellTool;
import org.apache.brooklyn.util.os.Os;
import org.apache.brooklyn.util.text.Identifiers;
import org.apache.brooklyn.util.text.Strings;
import org.apache.brooklyn.util.time.Time;

import java.util.List;
import java.util.Map;

import static org.apache.brooklyn.core.config.ConfigKeys.newConfigKeyWithDefault;
import static org.apache.brooklyn.util.core.internal.ssh.ShellAbstractTool.getOptionalVal;

/**
* Implemented like SshCliTool
*/
// TODO Use a general interface for SshTool and WinRmScriptTool with
public class WinRmScriptTool {
public static final ConfigKey<String> PROP_SCRIPT_DIR = newConfigKeyWithDefault(ShellTool.PROP_SCRIPT_DIR, "The directory where the script should be uploaded. It is an env variable.", "temp");
public static final ConfigKey<String> PROP_SUMMARY = ShellTool.PROP_SUMMARY;

private static final String SCRIPT_TEMP_DIR_VARIABLE = "BROOKLYN_TEMP_SCRIPT_DIR";

private String scriptNameWithoutExtension;
private String scriptDir;
private String summary;
private NaiveWindowsScriptRunner scriptRunner;

public WinRmScriptTool(Map<String, ?> props, NaiveWindowsScriptRunner scriptRunner) {
this.scriptDir = getOptionalVal(props, PROP_SCRIPT_DIR);

String summary = getOptionalVal(props, PROP_SUMMARY);
if (summary != null) {
summary = Strings.makeValidFilename(summary);
if (summary.length() > 30)
summary = summary.substring(0, 30);
}

this.summary = summary;
this.scriptNameWithoutExtension = "brooklyn-" +
Time.makeDateStampString() + "-" + Identifiers.makeRandomId(4) +
(Strings.isBlank(summary) ? "" : "-" + summary);

this.scriptRunner = scriptRunner;
}

public static String psStringExpressionToBatchVariable(String psString, String batchVariable) {
return "for /f \"delims=\" %%i in ('powershell -noprofile \\'Write-Host \""+ psString +"\"\\') do @set "+batchVariable+ "=%i";
}

private String psExtension(String filename) {
return filename + ".ps1";
}

private String batchExtension(String filename) {
return filename + ".bat";
}

public int execPsScript(List<String> psCommands) {
return execPsScript(ImmutableMap.<String, Object>of(), psCommands);
}

public int execPsScript(Map<String, Object> flags, List<String> commands) {
byte[] scriptBytes = toScript(commands).getBytes();
String scriptFilename = psExtension(scriptNameWithoutExtension);
if (flags.containsKey("scriptFilename"))
flags.put("scriptFilename", scriptFilename);
String scriptPath = Os.mergePathsWin("$env:"+scriptDir, scriptFilename);
scriptRunner.copyTo(new ByteInputStream(scriptBytes, scriptBytes.length), scriptPath);
return scriptRunner.executeNativeOrPsCommand(MutableMap.of(), null, "& " + scriptPath, summary, false);
}

public int execNativeScript(List<String> commands) {
return execNativeScript(ImmutableMap.<String, Object>of(), commands);
}

public int execNativeScript(Map<String, Object> flags, List<String> commands) {
byte[] scriptBytes = toScript(commands).getBytes();
String scriptFilename = batchExtension(scriptNameWithoutExtension);
if (flags.containsKey("scriptFilename"))
flags.put("scriptFilename", scriptFilename);
scriptRunner.copyTo(new ByteInputStream(scriptBytes, scriptBytes.length), Os.mergePathsWin("$env:"+scriptDir, scriptFilename));
String scriptPath = Os.mergePathsWin("%"+scriptDir+"%", scriptFilename);
return scriptRunner.executeNativeOrPsCommand(MutableMap.of(), scriptPath, null, summary, false);
}

private String toScript(List<String> commands) {
return Joiner.on("\r\n").join(commands);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@
import org.apache.brooklyn.core.sensor.Sensors;
import org.apache.brooklyn.entity.software.base.lifecycle.WinRmExecuteHelper;
import org.apache.brooklyn.location.winrm.WinRmMachineLocation;
import org.apache.brooklyn.util.core.internal.winrm.NativeWindowsScriptRunner;
import org.apache.brooklyn.util.core.ResourceUtils;
import org.apache.brooklyn.util.core.internal.winrm.NaiveWindowsScriptRunner;
import org.apache.brooklyn.util.core.task.Tasks;
import org.apache.brooklyn.util.exceptions.Exceptions;
import org.apache.brooklyn.util.exceptions.ReferenceWithError;
Expand All @@ -47,7 +48,7 @@

import static org.apache.brooklyn.util.JavaGroovyEquivalents.elvis;

public abstract class AbstractSoftwareProcessWinRmDriver extends AbstractSoftwareProcessDriver implements NativeWindowsScriptRunner {
public abstract class AbstractSoftwareProcessWinRmDriver extends AbstractSoftwareProcessDriver implements NaiveWindowsScriptRunner {
private static final Logger LOG = LoggerFactory.getLogger(AbstractSoftwareProcessWinRmDriver.class);

AttributeSensor<String> WINDOWS_USERNAME = Sensors.newStringSensor("windows.username",
Expand Down Expand Up @@ -153,36 +154,18 @@ public String getInstallDir() {
}

@Override
public int copyResource(Map<Object, Object> sshFlags, String source, String target, boolean createParentDir) {
if (createParentDir) {
createDirectory(getDirectory(target), "Creating resource directory");
}

InputStream stream = null;
try {
Tasks.setBlockingDetails("retrieving resource "+source+" for copying across");
stream = resource.getResourceFromUrl(source);
Tasks.setBlockingDetails("copying resource "+source+" to server");
return copyTo(stream, target);
} catch (Exception e) {
throw Exceptions.propagate(e);
} finally {
Tasks.setBlockingDetails(null);
if (stream != null) Streams.closeQuietly(stream);
}
public int copyResource(Map<Object, Object> flags, String source, String target, boolean createParentDir) {
return getLocation().copyResource(flags, source, target, createParentDir, resource);
}

@Override
public int copyResource(Map<Object, Object> sshFlags, InputStream source, String target, boolean createParentDir) {
if (createParentDir) {
createDirectory(getDirectory(target), "Creating resource directory");
}
return copyTo(source, target);
public int copyResource(Map<Object, Object> flags, InputStream source, String target, boolean createParentDir) {
return getLocation().copyResource(flags, source, target, createParentDir);
}

@Override
protected void createDirectory(String directoryName, String summaryForLogging) {
getLocation().executePsCommand("New-Item -path \"" + directoryName + "\" -type directory -ErrorAction SilentlyContinue");
getLocation().createDirectory(directoryName, summaryForLogging);
}

@Override
Expand Down Expand Up @@ -234,10 +217,6 @@ public void rebootAndWait() {
waitForWinRmStatus(true, entity.getConfig(VanillaWindowsProcess.REBOOT_COMPLETED_TIMEOUT)).getWithError();
}

private String getDirectory(String fileName) {
return fileName.substring(0, fileName.lastIndexOf("\\"));
}

private ReferenceWithError<Boolean> waitForWinRmStatus(final boolean requiredStatus, Duration timeout) {
// TODO: Reduce / remove duplication between this and JcloudsLocation.waitForWinRmAvailable
Callable<Boolean> checker = new Callable<Boolean>() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
import org.apache.brooklyn.api.mgmt.Task;
import org.apache.brooklyn.api.mgmt.TaskQueueingContext;
import org.apache.brooklyn.core.mgmt.BrooklynTaskTags;
import org.apache.brooklyn.util.core.internal.winrm.NativeWindowsScriptRunner;
import org.apache.brooklyn.util.core.internal.winrm.NaiveWindowsScriptRunner;
import org.apache.brooklyn.util.core.task.DynamicTasks;
import org.apache.brooklyn.util.core.task.TaskBuilder;
import org.apache.brooklyn.util.core.task.Tasks;
Expand All @@ -52,7 +52,7 @@ public class WinRmExecuteHelper {

private Task<Integer> task;

protected final NativeWindowsScriptRunner runner;
protected final NaiveWindowsScriptRunner runner;
public final String summary;

private String command;
Expand All @@ -63,7 +63,7 @@ public class WinRmExecuteHelper {
protected Predicate<? super Integer> resultCodeCheck = Predicates.alwaysTrue();
protected ByteArrayOutputStream stdout, stderr, stdin;

public WinRmExecuteHelper(NativeWindowsScriptRunner runner, String summary) {
public WinRmExecuteHelper(NaiveWindowsScriptRunner runner, String summary) {
this.runner = runner;
this.summary = summary;
}
Expand Down
Loading

0 comments on commit 36d8ed0

Please sign in to comment.