path: root/src/main/java/org/openslx/taskmanager/tasks/RemoteReboot.java

package org.openslx.taskmanager.tasks;

import java.io.IOException;
import java.net.Socket;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import org.openslx.taskmanager.api.AbstractTask;

import com.jcabi.ssh.Shell;
import com.google.gson.annotations.Expose;
import com.jcabi.ssh.SSH;

public class RemoteReboot extends AbstractTask
{
	@Expose
	private HashMap<String, String>[] clients;
	
	@Expose
	private boolean shutdown = false;
	
	@Expose
	private int minutes = 0;
	
	@Expose
	private String locationId;
	
	@Expose
	private String locationName;
	
	@Expose
	private String sshkey;
	
	@Expose
	private int port = 22;
		
	private Output status = new Output();
	
	@Override
	protected boolean initTask()
	{
		this.setStatusObject( this.status );

		status.clients = clients;
		Date shutdownTime = new Date(System.currentTimeMillis()+minutes*60*1000);
		SimpleDateFormat sdf = new SimpleDateFormat("HH:mm");
		status.time = sdf.format(shutdownTime);
		status.locationId = locationId;
		status.locationName = locationName;
		return true;
	}

	@Override
	protected boolean execute()
	{
		// try to connect to every client and start the reboot/shutdown process
		ExecutorService tp = Executors.newFixedThreadPool( clients.length > 4 ? 4 : clients.length );
		for (HashMap<String, String> client : clients) {
			final String ip = client.get("ip");
			status.clientStatus.put(ip, "connecting");
			tp.submit(new Runnable() {
			    public void run() {
					try {
						Shell shell = new SSH(ip, port, "root", sshkey);
						if (shutdown) {
							new Shell.Empty(shell).exec("/sbin/shutdown +" + minutes);
							status.clientStatus.put(ip, minutes == 0 ? "shutdown" : "shutdownat");
						} else {
							new Shell.Empty(shell).exec("/sbin/reboot");
							status.clientStatus.put(ip, "rebooting");
						}
					} catch (IOException e) {
						status.clientStatus.put(ip, "error");
					}
			    }
			});
		}
		tp.shutdown();
		
		try {
			tp.awaitTermination( clients.length * 5, TimeUnit.SECONDS );
		} catch ( InterruptedException e ) {
			// ...
		}
		
		// wait for rebooting clients to finish rebooting
		ArrayList<String> rebootingClients = new ArrayList<String>();
		for (Map.Entry<String, String> entry : status.clientStatus.entrySet()) {
			if (entry.getValue() == "rebooting") {
				rebootingClients.add(entry.getKey());
			}
		}
		
		if (rebootingClients.size() > 0) {
			ExecutorService statusTP = Executors.newFixedThreadPool( rebootingClients.size() > 4 ? 4 : rebootingClients.size() );
			for (final String ip : rebootingClients) {
				statusTP.submit(new Runnable() {
				    public void run() {
				    	while (!isOnline(ip)) {
				    		try {
								Thread.sleep(3000);
							} catch (InterruptedException e) {
							}
				    	}
				    	status.clientStatus.put(ip, "online");
				    }
				});
			}
			statusTP.shutdown();
			
			try {
				statusTP.awaitTermination( 180, TimeUnit.SECONDS );
			} catch ( InterruptedException e ) {
				// ...
			}
		}
		
		// change status of clients that got stuck because of timeouts
		for (Map.Entry<String, String> entry : status.clientStatus.entrySet()) {
			String value = entry.getValue();
		    if (value == "connecting" || value == "rebooting") {
		    	entry.setValue("error");
		    }
		}
		
		return true;
	}
	
	
	private boolean isOnline(String address)
	{ 
	    try (Socket s = new Socket(address, port)) {
	        return true;
	    } catch (IOException ex) {
	    }
	    return false;
	}
	
	
	/**
	 * Output - contains additional status data of this task
	 */
	class Output
	{
		private ConcurrentHashMap<String, String> clientStatus = new ConcurrentHashMap<String, String>();
		private HashMap<String, String>[] clients;
		private String time;
		private String locationId;
		private String locationName;
	}
}