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