1 /* 2 * Copyright (C) 2021 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.bedstead.nene.utils; 18 19 import androidx.annotation.CheckResult; 20 import androidx.annotation.Nullable; 21 22 import com.android.bedstead.nene.exceptions.AdbException; 23 import com.android.bedstead.nene.exceptions.NeneException; 24 import com.android.bedstead.nene.users.UserReference; 25 26 import java.util.function.Function; 27 28 /** 29 * A tool for progressively building and then executing a shell command. 30 */ 31 public final class ShellCommand { 32 33 // 10 seconds 34 private static final int MAX_WAIT_UNTIL_ATTEMPTS = 100; 35 private static final long WAIT_UNTIL_DELAY_MILLIS = 100; 36 37 /** 38 * Begin building a new {@link ShellCommand}. 39 */ 40 @CheckResult builder(String command)41 public static Builder builder(String command) { 42 if (command == null) { 43 throw new NullPointerException(); 44 } 45 return new Builder(command); 46 } 47 48 /** 49 * Create a builder and if {@code userReference} is not {@code null}, add "--user <userId>". 50 */ 51 @CheckResult builderForUser(@ullable UserReference userReference, String command)52 public static Builder builderForUser(@Nullable UserReference userReference, String command) { 53 Builder builder = builder(command); 54 if (userReference != null) { 55 builder.addOption("--user", userReference.id()); 56 } 57 58 return builder; 59 } 60 61 public static final class Builder { 62 private String mLinuxUser; 63 private final StringBuilder commandBuilder; 64 @Nullable 65 private byte[] mStdInBytes = null; 66 @Nullable 67 private boolean mAllowEmptyOutput = false; 68 @Nullable 69 private Function<String, Boolean> mOutputSuccessChecker = null; 70 Builder(String command)71 private Builder(String command) { 72 commandBuilder = new StringBuilder(command); 73 } 74 75 /** 76 * Add an option to the command. 77 * 78 * <p>e.g. --user 10 79 */ 80 @CheckResult addOption(String key, Object value)81 public Builder addOption(String key, Object value) { 82 // TODO: Deal with spaces/etc. 83 commandBuilder.append(" ").append(key).append(" ").append(value); 84 return this; 85 } 86 87 /** 88 * Add an operand to the command. 89 */ 90 @CheckResult addOperand(Object value)91 public Builder addOperand(Object value) { 92 // TODO: Deal with spaces/etc. 93 commandBuilder.append(" ").append(value); 94 return this; 95 } 96 97 /** 98 * If {@code false} an error will be thrown if the command has no output. 99 * 100 * <p>Defaults to {@code false} 101 */ 102 @CheckResult allowEmptyOutput(boolean allowEmptyOutput)103 public Builder allowEmptyOutput(boolean allowEmptyOutput) { 104 mAllowEmptyOutput = allowEmptyOutput; 105 return this; 106 } 107 108 /** 109 * Write the given {@code stdIn} to standard in. 110 */ 111 @CheckResult writeToStdIn(byte[] stdIn)112 public Builder writeToStdIn(byte[] stdIn) { 113 mStdInBytes = stdIn; 114 return this; 115 } 116 117 /** 118 * Validate the output when executing. 119 * 120 * <p>{@code outputSuccessChecker} should return {@code true} if the output is valid. 121 */ 122 @CheckResult validate(Function<String, Boolean> outputSuccessChecker)123 public Builder validate(Function<String, Boolean> outputSuccessChecker) { 124 mOutputSuccessChecker = outputSuccessChecker; 125 return this; 126 } 127 128 /** 129 * Run the command as a given linux user. 130 */ 131 @CheckResult asLinuxUser(String user)132 public Builder asLinuxUser(String user) { 133 mLinuxUser = user; 134 return this; 135 } 136 137 /** 138 * Run the command as the root linux user. 139 */ 140 @CheckResult asRoot()141 public Builder asRoot() { 142 return asLinuxUser("root"); 143 } 144 145 /** 146 * Build the full command including all options and operands. 147 */ build()148 public String build() { 149 if (mLinuxUser != null) { 150 return "su " + mLinuxUser + " " + commandBuilder.toString(); 151 } 152 return commandBuilder.toString(); 153 } 154 155 /** 156 * See {@link #execute()} except that any {@link AdbException} is wrapped in a 157 * {@link NeneException} with the message {@code errorMessage}. 158 */ executeOrThrowNeneException(String errorMessage)159 public String executeOrThrowNeneException(String errorMessage) throws NeneException { 160 try { 161 return execute(); 162 } catch (AdbException e) { 163 throw new NeneException(errorMessage, e); 164 } 165 } 166 167 /** See {@link ShellCommandUtils#executeCommand(java.lang.String)}. */ execute()168 public String execute() throws AdbException { 169 if (mOutputSuccessChecker != null) { 170 return ShellCommandUtils.executeCommandAndValidateOutput( 171 build(), 172 /* allowEmptyOutput= */ mAllowEmptyOutput, 173 mStdInBytes, 174 mOutputSuccessChecker); 175 } 176 177 return ShellCommandUtils.executeCommand( 178 build(), 179 /* allowEmptyOutput= */ mAllowEmptyOutput, 180 mStdInBytes); 181 } 182 183 /** 184 * See {@link #execute} and then extract information from the output using 185 * {@code outputParser}. 186 * 187 * <p>If any {@link Exception} is thrown by {@code outputParser}, and {@link AdbException} 188 * will be thrown. 189 */ executeAndParseOutput(Function<String, E> outputParser)190 public <E> E executeAndParseOutput(Function<String, E> outputParser) throws AdbException { 191 String output = execute(); 192 193 try { 194 return outputParser.apply(output); 195 } catch (RuntimeException e) { 196 throw new AdbException( 197 "Could not parse output", commandBuilder.toString(), output, e); 198 } 199 } 200 201 /** 202 * Execute the command and check that the output meets a given criteria. Run the 203 * command repeatedly until the output meets the criteria. 204 * 205 * <p>{@code outputSuccessChecker} should return {@code true} if the output indicates the 206 * command executed successfully. 207 */ executeUntilValid()208 public String executeUntilValid() throws InterruptedException, AdbException { 209 int attempts = 0; 210 while (attempts++ < MAX_WAIT_UNTIL_ATTEMPTS) { 211 try { 212 return execute(); 213 } catch (AdbException e) { 214 // ignore, will retry 215 Thread.sleep(WAIT_UNTIL_DELAY_MILLIS); 216 } 217 } 218 return execute(); 219 } 220 forBytes()221 public BytesBuilder forBytes() { 222 if (mOutputSuccessChecker != null) { 223 throw new IllegalStateException("Cannot call .forBytes after .validate"); 224 } 225 226 return new BytesBuilder(this); 227 } 228 229 @Override toString()230 public String toString() { 231 return "ShellCommand$Builder{cmd=" + build() + "}"; 232 } 233 } 234 235 public static final class BytesBuilder { 236 237 private final Builder mBuilder; 238 BytesBuilder(Builder builder)239 private BytesBuilder(Builder builder) { 240 mBuilder = builder; 241 } 242 243 /** See {@link ShellCommandUtils#executeCommandForBytes(java.lang.String)}. */ execute()244 public byte[] execute() throws AdbException { 245 return ShellCommandUtils.executeCommandForBytes( 246 mBuilder.build(), 247 mBuilder.mStdInBytes); 248 } 249 } 250 } 251