content: /**
   * 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 backtype.storm.utils;
  
  import java.io.BufferedReader;
  import java.io.File;
  import java.io.IOException;
  import java.io.InputStreamReader;
  import java.io.InputStream;
  import java.util.Map;
  import java.util.Timer;
  import java.util.TimerTask;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  abstract public class ShellUtils {
      public static final Logger LOG = LoggerFactory.getLogger(ShellUtils.class);
  
      // OSType detection
      public enum OSType {
          OS_TYPE_LINUX,
          OS_TYPE_WIN,
          OS_TYPE_SOLARIS,
          OS_TYPE_MAC,
          OS_TYPE_FREEBSD,
          OS_TYPE_OTHER
      }
  
      public static final OSType osType = getOSType();
  
      static private OSType getOSType() {
          String osName = System.getProperty("os.name");
          if (osName.startsWith("Windows")) {
              return OSType.OS_TYPE_WIN;
          } else if (osName.contains("SunOS") || osName.contains("Solaris")) {
              return OSType.OS_TYPE_SOLARIS;
          } else if (osName.contains("Mac")) {
              return OSType.OS_TYPE_MAC;
          } else if (osName.contains("FreeBSD")) {
              return OSType.OS_TYPE_FREEBSD;
          } else if (osName.startsWith("Linux")) {
              return OSType.OS_TYPE_LINUX;
          } else {
              // Some other form of Unix
              return OSType.OS_TYPE_OTHER;
          }
      }
  
      // Helper static vars for each platform
      public static final boolean WINDOWS = (osType == OSType.OS_TYPE_WIN);
      public static final boolean SOLARIS = (osType == OSType.OS_TYPE_SOLARIS);
      public static final boolean MAC = (osType == OSType.OS_TYPE_MAC);
      public static final boolean FREEBSD = (osType == OSType.OS_TYPE_FREEBSD);
      public static final boolean LINUX = (osType == OSType.OS_TYPE_LINUX);
      public static final boolean OTHER = (osType == OSType.OS_TYPE_OTHER);
  
      /** Token separator regex used to parse Shell tool outputs */
      public static final String TOKEN_SEPARATOR_REGEX
          = WINDOWS ? "[|\n\r]" : "[ \t\n\r\f]";
  
      private long interval; // refresh interval in msec
      private long lastTime; // last time the command was performed
      final private boolean redirectErrorStream; // merge stdout and stderr
      private Map<String, String> environment; // env for the command execution
      private File dir;
      private Process process; // sub process used to execute the command
      private int exitCode;
      /** Time after which the executing script would be timedout */
      protected long timeOutInterval = 0L;
      /** If or not script timed out */
      private AtomicBoolean timedOut;
  
      /** If or not script finished executing */
      private volatile AtomicBoolean completed;
  
      public ShellUtils() {
          this(0L);
      }
  
      public ShellUtils(long interval) {
          this(interval, false);
      }
  
      public int getExitCode() {
        return exitCode;
      }
      /**
       * @param interval the minimum duration to wait before re-executing the command.
       */
      public ShellUtils(long interval, boolean redirectErrorStream) {
          this.interval = interval;
          this.lastTime = (interval < 0) ? 0 : -interval;
          this.redirectErrorStream = redirectErrorStream;
      }
  
      /**
       * set the environment for the command
       * 
       * @param env Mapping of environment variables
       */
      protected void setEnvironment(Map<String, String> env) {
          this.environment = env;
      }
  
      /**
       * set the working directory
       * 
       * @param dir The directory where the command would be executed
       */
      protected void setWorkingDirectory(File dir) {
          this.dir = dir;
      }
  
      /** a Unix command to get the current user's groups list */
      public static String[] getGroupsCommand() {
          return (WINDOWS)? new String[]{"cmd", "/c", "groups"}
          : new String[]{"bash", "-c", "groups"};
      }
  
      /**
       * a Unix command to get a given user's groups list. If the OS is not WINDOWS,
       * the command will get the user's primary group first and finally get the
       * groups list which includes the primary group. i.e. the user's primary group will be included twice.
       */
      public static String[] getGroupsForUserCommand(final String user) {
          // 'groups username' command return is non-consistent across different unixes
          return new String [] {"bash", "-c", "id -gn " + user
                           + "&& id -Gn " + user};
      }
  
      /** check to see if a command needs to be executed and execute if needed */
      protected void run() throws IOException {
          if (lastTime + interval > System.currentTimeMillis())
              return;
          exitCode = 0; // reset for next run
          runCommand();
      }
  
      /** Run a command */
      private void runCommand() throws IOException {
          ProcessBuilder builder = new ProcessBuilder(getExecString());
          Timer timeOutTimer = null;
          ShellTimeoutTimerTask timeoutTimerTask;
          timedOut = new AtomicBoolean(false);
          completed = new AtomicBoolean(false);
  
          if (environment != null) {
              builder.environment().putAll(this.environment);
          }
          if (dir != null) {
              builder.directory(this.dir);
          }
  
          builder.redirectErrorStream(redirectErrorStream);
          process = builder.start();
  
          if (timeOutInterval > 0) {
              timeOutTimer = new Timer("Shell command timeout");
              timeoutTimerTask = new ShellTimeoutTimerTask(this);
              // One time scheduling.
              timeOutTimer.schedule(timeoutTimerTask, timeOutInterval);
          }
          final BufferedReader errReader =
              new BufferedReader(new InputStreamReader(process
                                                       .getErrorStream()));
          BufferedReader inReader =
              new BufferedReader(new InputStreamReader(process
                                                       .getInputStream()));
          final StringBuffer errMsg = new StringBuffer();
  
          // read error and input streams as this would free up the buffers
          // free the error stream buffer
          Thread errThread = new Thread() {
              @Override
              public void run() {
                  try {
                      String line = errReader.readLine();
                      while ((line != null) && !isInterrupted()) {
                          errMsg.append(line);
                          errMsg.append(System.getProperty("line.separator"));
                          line = errReader.readLine();
                      }
                  } catch (IOException ioe) {
                      LOG.warn("Error reading the error stream", ioe);
                  }
              }
          };
          try {
              errThread.start();
          } catch (IllegalStateException ignored) {
          }
          try {
              parseExecResult(inReader); // parse the output
              // clear the input stream buffer
              String line = inReader.readLine();
              while (line != null) {
                  line = inReader.readLine();
              }
              // wait for the process to finish and check the exit code
              exitCode = process.waitFor();
              // make sure that the error thread exits
              joinThread(errThread);
              completed.set(true);
              // the timeout thread handling
              // taken care in finally block
              if (exitCode != 0) {
                  throw new ExitCodeException(exitCode, errMsg.toString());
              }
          } catch (InterruptedException ie) {
              throw new IOException(ie.toString());
          } finally {
              if (timeOutTimer != null) {
                  timeOutTimer.cancel();
              }
              // close the input stream
              try {
                  // JDK 7 tries to automatically drain the input streams for us
                  // when the process exits, but since close is not synchronized,
                  // it creates a race if we close the stream first and the same
                  // fd is recycled. the stream draining thread will attempt to
                  // drain that fd!! it may block, OOM, or cause bizarre behavior
                  // see: https://bugs.openjdk.java.net/browse/JDK-8024521
                  // issue is fixed in build 7u60
                  InputStream stdout = process.getInputStream();
                  synchronized (stdout) {
                      inReader.close();
                  }
              } catch (IOException ioe) {
                  LOG.warn("Error while closing the input stream", ioe);
              }
              if (!completed.get()) {
                  errThread.interrupt();
                  joinThread(errThread);
              }
              try {
                  InputStream stderr = process.getErrorStream();
                  synchronized (stderr) {
                      errReader.close();
                  }
              } catch (IOException ioe) {
                  LOG.warn("Error while closing the error stream", ioe);
              }
              process.destroy();
              lastTime = System.currentTimeMillis();
          }
      }
  
      private static void joinThread(Thread t) {
          while (t.isAlive()) {
              try {
                  t.join();
              } catch (InterruptedException ie) {
                  if (LOG.isWarnEnabled()) {
                      LOG.warn("Interrupted while joining on: " + t, ie);
                  }
                  t.interrupt(); // propagate interrupt
              }
          }
      }
  
      /** return an array containing the command name & its parameters */
      protected abstract String[] getExecString();
  
      /** Parse the execution result */
      protected abstract void parseExecResult(BufferedReader lines)
          throws IOException;
      /**
       * get the current sub-process executing the given command
       * 
       * @return process executing the command
       */
      public Process getProcess() {
          return process;
      }
  
      /**
       * This is an IOException with exit code added.
       */
      public static class ExitCodeException extends IOException {
          int exitCode;
  
          public ExitCodeException(int exitCode, String message) {
              super(message);
              this.exitCode = exitCode;
          }
  
          public int getExitCode() {
              return exitCode;
          }
      }
  
      /**
       * A simple shell command executor.
       * 
       * <code>ShellCommandExecutor</code>should be used in cases where the output of the command needs no explicit parsing and where the command, working
       * directory and the environment remains unchanged. The output of the command is stored as-is and is expected to be small.
       */
      public static class ShellCommandExecutor extends ShellUtils {
  
          private String[] command;
          private StringBuffer output;
  
          public ShellCommandExecutor(String[] execString) {
              this(execString, null);
          }
  
          public ShellCommandExecutor(String[] execString, File dir) {
              this(execString, dir, null);
          }
  
          public ShellCommandExecutor(String[] execString, File dir,
                                      Map<String, String> env) {
              this(execString, dir, env, 0L);
          }
  
          /**
           * Create a new instance of the ShellCommandExecutor to execute a command.
           * 
           * @param execString The command to execute with arguments
           * @param dir If not-null, specifies the directory which should be set as the current working directory for the command. If null, the current working
           *            directory is not modified.
           * @param env If not-null, environment of the command will include the key-value pairs specified in the map. If null, the current environment is not
           *            modified.
           * @param timeout Specifies the time in milliseconds, after which the command will be killed and the status marked as timedout. If 0, the command will
           *            not be timed out.
           */
          public ShellCommandExecutor(String[] execString, File dir,
                                      Map<String, String> env, long timeout) {
              command = execString.clone();
              if (dir != null) {
                  setWorkingDirectory(dir);
              }
              if (env != null) {
                  setEnvironment(env);
              }
              timeOutInterval = timeout;
          }
  
          /** Execute the shell command. */
          public void execute() throws IOException {
              this.run();
          }
  
          @Override
          public String[] getExecString() {
              return command;
          }
  
          @Override
          protected void parseExecResult(BufferedReader lines) throws IOException {
              output = new StringBuffer();
              char[] buf = new char[512];
              int nRead;
              while ((nRead = lines.read(buf, 0, buf.length)) > 0) {
                  output.append(buf, 0, nRead);
              }
          }
  
          /** Get the output of the shell command. */
          public String getOutput() {
              return (output == null) ? "" : output.toString();
          }
  
          /**
           * Returns the commands of this instance. Arguments with spaces in are presented with quotes round; other arguments are presented raw
           * 
           * @return a string representation of the object.
           */
          @Override
          public String toString() {
              StringBuilder builder = new StringBuilder();
              String[] args = getExecString();
              for (String s : args) {
                  if (s.indexOf(' ') >= 0) {
                      builder.append('"').append(s).append('"');
                  } else {
                      builder.append(s);
                  }
                  builder.append(' ');
              }
              return builder.toString();
          }
      }
  
      /**
       * To check if the passed script to shell command executor timed out or not.
       * 
       * @return if the script timed out.
       */
      public boolean isTimedOut() {
          return timedOut.get();
      }
  
      /**
       * Set if the command has timed out.
       * 
       */
      private void setTimedOut() {
          this.timedOut.set(true);
      }
  
      /**
       * Static method to execute a shell command. Covers most of the simple cases without requiring the user to implement the <code>Shell</code> interface.
       * 
       * @param cmd shell command to execute.
       * @return the output of the executed command.
       */
      public static String execCommand(String... cmd) throws IOException {
          return execCommand(null, cmd, 0L);
      }
  
      /**
       * Static method to execute a shell command. Covers most of the simple cases without requiring the user to implement the <code>Shell</code> interface.
       * 
       * @param env the map of environment key=value
       * @param cmd shell command to execute.
       * @param timeout time in milliseconds after which script should be marked timeout
       * @return the output of the executed command.o
       */
  
      public static String execCommand(Map<String, String> env, String[] cmd,
                                       long timeout) throws IOException {
          ShellCommandExecutor exec = new ShellCommandExecutor(cmd, null, env,
                                                               timeout);
          exec.execute();
          return exec.getOutput();
      }
  
      /**
       * Static method to execute a shell command. Covers most of the simple cases without requiring the user to implement the <code>Shell</code> interface.
       * 
       * @param env the map of environment key=value
       * @param cmd shell command to execute.
       * @return the output of the executed command.
       */
      public static String execCommand(Map<String,String> env, String ... cmd)
          throws IOException {
          return execCommand(env, cmd, 0L);
      }
  
      /**
       * Timer which is used to timeout scripts spawned off by shell.
       */
      private static class ShellTimeoutTimerTask extends TimerTask {
  
          private ShellUtils shell;
  
          public ShellTimeoutTimerTask(ShellUtils shell) {
              this.shell = shell;
          }
  
          @Override
          public void run() {
              Process p = shell.getProcess();
              try {
                  p.exitValue();
              } catch (Exception e) {
                  // Process has not terminated.
                  // So check if it has completed
                  // if not just destroy it.
                  if (p != null && !shell.completed.get()) {
                      shell.setTimedOut();
                      p.destroy();
                  }
              }
          }
      }
  
  }
  