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