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

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