runtime/org.argeo.slc.support.simple/src/main/java/org/argeo/slc/jsch/RemoteExec.java

   1 /*
   2  * Copyright (C) 2007-2012 Mathieu Baudier
   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 package org.argeo.slc.jsch;
  17
  18 import java.io.BufferedReader;
  19 import java.io.BufferedWriter;
  20 import java.io.File;
  21 import java.io.FileOutputStream;
  22 import java.io.IOException;
  23 import java.io.InputStream;
  24 import java.io.InputStreamReader;
  25 import java.io.OutputStream;
  26 import java.io.OutputStreamWriter;
  27 import java.util.ArrayList;
  28 import java.util.HashMap;
  29 import java.util.Hashtable;
  30 import java.util.List;
  31 import java.util.Map;
  32 import java.util.StringTokenizer;
  33
  34 import org.apache.commons.exec.ExecuteStreamHandler;
  35 import org.apache.commons.io.IOUtils;
  36 import org.apache.commons.logging.Log;
  37 import org.apache.commons.logging.LogFactory;
  38 import org.argeo.slc.SlcException;
  39 import org.argeo.slc.core.execution.ExecutionResources;
  40 import org.argeo.slc.core.execution.tasks.SystemCall;
  41 import org.springframework.core.io.Resource;
  42 import org.springframework.util.StringUtils;
  43
  44 import com.jcraft.jsch.Channel;
  45 import com.jcraft.jsch.ChannelExec;
  46 import com.jcraft.jsch.ChannelShell;
  47 import com.jcraft.jsch.Session;
  48
  49 public class RemoteExec extends AbstractJschTask {
  50         private final static Log log = LogFactory.getLog(RemoteExec.class);
  51
  52         private Boolean failOnBadExitStatus = true;
  53
  54         private List<String> commands = new ArrayList<String>();
  55         private String command;
  56         private SystemCall systemCall;
  57         private List<SystemCall> systemCalls = new ArrayList<SystemCall>();
  58         private Resource script;
  59         private Boolean xForwarding = false;
  60         private Boolean agentForwarding = false;
  61         private Boolean forceShell = false;
  62         private Map<String, String> env = new HashMap<String, String>();
  63         private Resource stdIn = null;
  64         private Resource stdOut = null;
  65         private ExecutionResources executionResources;
  66
  67         private String user;
  68
  69         private ExecuteStreamHandler streamHandler = null;
  70
  71         private Integer lastExitStatus = null;
  72         /**
  73          * If set, stdout is written to it as a list of lines. Cleared before each
  74          * run.
  75          */
  76         private List<String> stdOutLines = null;
  77         private Boolean logEvenIfStdOutLines = false;
  78         private Boolean quiet = false;
  79
  80         public RemoteExec() {
  81         }
  82
  83         public RemoteExec(SshTarget sshTarget, String cmd) {
  84                 setSshTarget(sshTarget);
  85                 setCommand(cmd);
  86         }
  87
  88         public void run(Session session) {
  89                 List<String> commandsToUse = new ArrayList<String>(commands);
  90                 String commandToUse = command;
  91                 // convert system calls
  92                 if (systemCall != null) {
  93                         if (command != null)
  94                                 throw new SlcException("Cannot specify command AND systemCall");
  95                         commandToUse = convertSystemCall(systemCall);
  96                 }
  97
  98                 if (systemCalls.size() != 0) {
  99                         if (commandsToUse.size() != 0)
 100                                 throw new SlcException(
 101                                                 "Cannot specify commands AND systemCalls");
 102                         for (SystemCall systemCall : systemCalls)
 103                                 commandsToUse.add(convertSystemCall(systemCall));
 104                 }
 105
 106                 if (script != null) {
 107                         // TODO: simply pass the script as a string command
 108                         if (commandsToUse.size() != 0)
 109                                 throw new SlcException("Cannot specify commands and script");
 110                         BufferedReader reader = null;
 111                         try {
 112                                 reader = new BufferedReader(new InputStreamReader(
 113                                                 script.getInputStream()));
 114                                 String line = null;
 115                                 while ((line = reader.readLine()) != null) {
 116                                         if (!StringUtils.hasText(line))
 117                                                 continue;
 118                                         commandsToUse.add(line);
 119                                 }
 120                         } catch (IOException e) {
 121                                 throw new SlcException("Cannot read script " + script, e);
 122                         } finally {
 123                                 IOUtils.closeQuietly(reader);
 124                         }
 125                 }
 126
 127                 if (forceShell) {
 128                         // for the time being do not interpret both \n and ;
 129                         // priority to \n
 130                         // until we know how to parse ; within ""
 131                         if (commandToUse.indexOf('\n') >= 0) {
 132                                 StringTokenizer st = new StringTokenizer(commandToUse, "\n");
 133                                 while (st.hasMoreTokens()) {
 134                                         String cmd = st.nextToken();
 135                                         commandsToUse.add(cmd);
 136                                 }
 137                         } else if (commandToUse.indexOf(';') >= 0) {
 138                                 StringTokenizer st = new StringTokenizer(commandToUse, ";");
 139                                 while (st.hasMoreTokens()) {
 140                                         String cmd = st.nextToken();
 141                                         commandsToUse.add(cmd);
 142                                 }
 143                         } else {
 144                                 commandsToUse.add(commandToUse);
 145                         }
 146                         commandToUse = null;
 147                 }
 148
 149                 // run as user
 150                 if (user != null) {
 151                         if (commandsToUse.size() > 0) {
 152                                 commandsToUse.add(0, "su - " + user);
 153                                 commandsToUse.add("exit");
 154                         } else {
 155                                 if (command.indexOf('\"') >= 0)
 156                                         throw new SlcException(
 157                                                         "Don't know how to su a command with \", use shell instead.");
 158                                 commandToUse = "su - " + user + " -c \"" + command + "\"";
 159                         }
 160                 }
 161
 162                 // execute command(s)
 163                 if (commandToUse != null) {
 164                         if (commandsToUse.size() != 0)
 165                                 throw new SlcException(
 166                                                 "Specify either a single command or a list of commands.");
 167                         remoteExec(session, commandToUse);
 168                 } else {
 169                         if (commandsToUse.size() == 0)
 170                                 throw new SlcException(
 171                                                 "Neither a single command or a list of commands has been specified.");
 172
 173                         remoteExec(session, commandsToUse, script != null ? "script "
 174                                         + script.getFilename() : commandsToUse.size() + " commands");
 175                 }
 176         }
 177
 178         protected String convertSystemCall(SystemCall systemCall) {
 179                 // TODO: prepend environment variables
 180                 // TODO: deal with exec dir
 181                 return systemCall.asCommand();
 182         }
 183
 184         protected void remoteExec(Session session, final List<String> commands,
 185                         String description) {
 186                 try {
 187                         final ChannelShell channel = (ChannelShell) session
 188                                         .openChannel("shell");
 189                         channel.setInputStream(null);
 190                         channel.setXForwarding(xForwarding);
 191                         channel.setAgentForwarding(agentForwarding);
 192                         channel.setEnv(new Hashtable<String, String>(env));
 193
 194                         /*
 195                          * // Choose the pty-type "vt102".
 196                          * ((ChannelShell)channel).setPtyType("vt102");
 197                          */
 198                         // Writer thread
 199                         final BufferedWriter writer = new BufferedWriter(
 200                                         new OutputStreamWriter(channel.getOutputStream()));
 201
 202                         if (log.isDebugEnabled())
 203                                 log.debug("Run " + description + " on " + getSshTarget()
 204                                                 + "...");
 205                         channel.connect();
 206
 207                         // write commands to shell
 208                         Thread writerThread = new Thread("Shell writer " + getSshTarget()) {
 209                                 @Override
 210                                 public void run() {
 211                                         try {
 212                                                 for (String line : commands) {
 213                                                         if (!StringUtils.hasText(line))
 214                                                                 continue;
 215                                                         writer.write(line);
 216                                                         writer.newLine();
 217                                                 }
 218                                                 writer.append("exit");
 219                                                 writer.newLine();
 220                                                 writer.flush();
 221                                                 // channel.disconnect();
 222                                         } catch (IOException e) {
 223                                                 throw new SlcException("Cannot write to shell on "
 224                                                                 + getSshTarget(), e);
 225                                         } finally {
 226                                                 IOUtils.closeQuietly(writer);
 227                                         }
 228                                 }
 229                         };
 230                         writerThread.start();
 231
 232                         readStdOut(channel);
 233                         checkExitStatus(channel);
 234                         channel.disconnect();
 235
 236                 } catch (Exception e) {
 237                         throw new SlcException("Cannot use SSH shell on " + getSshTarget(),
 238                                         e);
 239                 }
 240
 241         }
 242
 243         protected void remoteExec(Session session, String command) {
 244                 try {
 245                         final ChannelExec channel = (ChannelExec) session
 246                                         .openChannel("exec");
 247                         channel.setCommand(command);
 248
 249                         channel.setInputStream(null);
 250                         channel.setXForwarding(xForwarding);
 251                         channel.setAgentForwarding(agentForwarding);
 252                         channel.setEnv(new Hashtable<String, String>(env));
 253                         channel.setErrStream(null);
 254
 255                         // Standard Error
 256                         readStdErr(channel);
 257
 258                         if (log.isDebugEnabled())
 259                                 log.debug("Run '" + command + "' on " + getSshTarget() + "...");
 260                         channel.connect();
 261
 262                         readStdIn(channel);
 263                         readStdOut(channel);
 264
 265                         if (streamHandler != null){
 266                                 streamHandler.start();
 267                                 while(!channel.isClosed()){
 268                                         try {
 269                                                 Thread.sleep(100);
 270                                         } catch (Exception e) {
 271                                                 break;
 272                                         }
 273                                 }
 274                         }
 275
 276                         checkExitStatus(channel);
 277                         channel.disconnect();
 278                 } catch (Exception e) {
 279                         throw new SlcException("Cannot execute remotely '" + command
 280                                         + "' on " + getSshTarget(), e);
 281                 }
 282         }
 283
 284         protected void readStdOut(Channel channel) {
 285                 try {
 286                         if (stdOut != null) {
 287                                 OutputStream localStdOut = createOutputStream(stdOut);
 288                                 try {
 289                                         IOUtils.copy(channel.getInputStream(), localStdOut);
 290                                 } finally {
 291                                         IOUtils.closeQuietly(localStdOut);
 292                                 }
 293                         } else if (streamHandler != null) {
 294                                 if (channel.getInputStream() != null)
 295                                         streamHandler.setProcessOutputStream(channel
 296                                                         .getInputStream());
 297                         } else {
 298                                 BufferedReader stdOut = null;
 299                                 try {
 300                                         InputStream in = channel.getInputStream();
 301                                         stdOut = new BufferedReader(new InputStreamReader(in));
 302                                         String line = null;
 303                                         while ((line = stdOut.readLine()) != null) {
 304                                                 if (!line.trim().equals("")) {
 305
 306                                                         if (stdOutLines != null) {
 307                                                                 stdOutLines.add(line);
 308                                                                 if (logEvenIfStdOutLines && !quiet)
 309                                                                         log.info(line);
 310                                                         } else {
 311                                                                 if (!quiet)
 312                                                                         log.info(line);
 313                                                         }
 314                                                 }
 315                                         }
 316                                 } finally {
 317                                         IOUtils.closeQuietly(stdOut);
 318                                 }
 319                         }
 320                 } catch (IOException e) {
 321                         throw new SlcException("Cannot redirect stdout from "
 322                                         + getSshTarget(), e);
 323                 }
 324         }
 325
 326         protected void readStdErr(final ChannelExec channel) {
 327                 if (streamHandler != null) {
 328                         try {
 329                                 streamHandler.setProcessOutputStream(channel.getErrStream());
 330                         } catch (IOException e) {
 331                                 throw new SlcException("Cannot read stderr from "
 332                                                 + getSshTarget(), e);
 333                         }
 334                 } else {
 335                         new Thread("stderr " + getSshTarget()) {
 336                                 public void run() {
 337                                         BufferedReader stdErr = null;
 338                                         try {
 339                                                 InputStream in = channel.getErrStream();
 340                                                 stdErr = new BufferedReader(new InputStreamReader(in));
 341                                                 String line = null;
 342                                                 while ((line = stdErr.readLine()) != null) {
 343                                                         if (!line.trim().equals(""))
 344                                                                 log.error(line);
 345                                                 }
 346                                         } catch (IOException e) {
 347                                                 if (log.isDebugEnabled())
 348                                                         log.error("Cannot read stderr from "
 349                                                                         + getSshTarget(), e);
 350                                         } finally {
 351                                                 IOUtils.closeQuietly(stdErr);
 352                                         }
 353                                 }
 354                         }.start();
 355                 }
 356         }
 357
 358         protected void readStdIn(final ChannelExec channel) {
 359                 if (stdIn != null) {
 360                         Thread stdInThread = new Thread("Stdin " + getSshTarget()) {
 361                                 @Override
 362                                 public void run() {
 363                                         OutputStream out = null;
 364                                         try {
 365                                                 out = channel.getOutputStream();
 366                                                 IOUtils.copy(stdIn.getInputStream(), out);
 367                                         } catch (IOException e) {
 368                                                 throw new SlcException("Cannot write stdin on "
 369                                                                 + getSshTarget(), e);
 370                                         } finally {
 371                                                 IOUtils.closeQuietly(out);
 372                                         }
 373                                 }
 374                         };
 375                         stdInThread.start();
 376                 } else if (streamHandler != null) {
 377                         try {
 378                                 streamHandler.setProcessInputStream(channel.getOutputStream());
 379                         } catch (IOException e) {
 380                                 throw new SlcException("Cannot write stdin on "
 381                                                 + getSshTarget(), e);
 382                         }
 383                 }
 384         }
 385
 386         protected void checkExitStatus(Channel channel) {
 387                 if (channel.isClosed()) {
 388                         lastExitStatus = channel.getExitStatus();
 389                         if (lastExitStatus == 0) {
 390                                 if (log.isTraceEnabled())
 391                                         log.trace("Remote execution exit status: " + lastExitStatus);
 392                         } else {
 393                                 String msg = "Remote execution failed with " + " exit status: "
 394                                                 + lastExitStatus;
 395                                 if (failOnBadExitStatus)
 396                                         throw new SlcException(msg);
 397                                 else
 398                                         log.error(msg);
 399                         }
 400                 }
 401
 402         }
 403
 404         protected OutputStream createOutputStream(Resource target) {
 405                 FileOutputStream out = null;
 406                 try {
 407
 408                         final File file;
 409                         if (executionResources != null)
 410                                 file = new File(executionResources.getAsOsPath(target, true));
 411                         else
 412                                 file = target.getFile();
 413                         out = new FileOutputStream(file, false);
 414                 } catch (IOException e) {
 415                         log.error("Cannot get file for " + target, e);
 416                         IOUtils.closeQuietly(out);
 417                 }
 418                 return out;
 419         }
 420
 421         public Integer getLastExitStatus() {
 422                 return lastExitStatus;
 423         }
 424
 425         public void setStreamHandler(ExecuteStreamHandler executeStreamHandler) {
 426                 this.streamHandler = executeStreamHandler;
 427         }
 428
 429         public void setCommand(String command) {
 430                 this.command = command;
 431         }
 432
 433         public void setCommands(List<String> commands) {
 434                 this.commands = commands;
 435         }
 436
 437         public void setFailOnBadExitStatus(Boolean failOnBadExitStatus) {
 438                 this.failOnBadExitStatus = failOnBadExitStatus;
 439         }
 440
 441         public void setSystemCall(SystemCall systemCall) {
 442                 this.systemCall = systemCall;
 443         }
 444
 445         public void setSystemCalls(List<SystemCall> systemCalls) {
 446                 this.systemCalls = systemCalls;
 447         }
 448
 449         public void setScript(Resource script) {
 450                 this.script = script;
 451         }
 452
 453         public void setxForwarding(Boolean xForwarding) {
 454                 this.xForwarding = xForwarding;
 455         }
 456
 457         public void setAgentForwarding(Boolean agentForwarding) {
 458                 this.agentForwarding = agentForwarding;
 459         }
 460
 461         public void setEnv(Map<String, String> env) {
 462                 this.env = env;
 463         }
 464
 465         public void setForceShell(Boolean forceShell) {
 466                 this.forceShell = forceShell;
 467         }
 468
 469         public List<String> getCommands() {
 470                 return commands;
 471         }
 472
 473         public void setStdOutLines(List<String> stdOutLines) {
 474                 this.stdOutLines = stdOutLines;
 475         }
 476
 477         public void setLogEvenIfStdOutLines(Boolean logEvenIfStdOutLines) {
 478                 this.logEvenIfStdOutLines = logEvenIfStdOutLines;
 479         }
 480
 481         public void setQuiet(Boolean quiet) {
 482                 this.quiet = quiet;
 483         }
 484
 485         public void setStdIn(Resource stdIn) {
 486                 this.stdIn = stdIn;
 487         }
 488
 489         public void setStdOut(Resource stdOut) {
 490                 this.stdOut = stdOut;
 491         }
 492
 493         public void setExecutionResources(ExecutionResources executionResources) {
 494                 this.executionResources = executionResources;
 495         }
 496
 497         public void setUser(String user) {
 498                 this.user = user;
 499         }
 500
 501 }