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

package org.openslx.taskmanager.tasks;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
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.google.gson.annotations.Expose;
import com.jcabi.ssh.SSH;
import com.jcabi.ssh.Shell;

public class RemoteReboot extends AbstractTask
{
	@Expose
	private Client[] clients;
	
	@Expose
	private boolean shutdown;
	
	@Expose
	private int minutes;
	
	@Expose
	private String locationId;
	
	@Expose
	private String locationName;
	
	@Expose
	private String sshkey;
	
	@Expose
	private int port;
		
	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()
	{
		if ( clients.length == 0 )
			return true;
		// try to connect to every client and start the reboot/shutdown process
		ExecutorService tp = Executors.newFixedThreadPool( clients.length > 4 ? 4 : clients.length );
		for (final Client client : clients) {
			if ( client == null || client.clientip == null )
				continue;
			status.clientStatus.put(client.clientip, ClientStatus.CONNECTING);
			tp.submit(new Runnable() {
			    public void run() {
					try {
						Shell shell = new SSH(client.clientip, port, "root", sshkey);
						if (shutdown) {
							new Shell.Empty(shell).exec("/sbin/shutdown +" + minutes);
							status.clientStatus.put(client.clientip, minutes == 0 ? ClientStatus.SHUTDOWN : ClientStatus.SHUTDOWN_AT);
						} else {
							new Shell.Empty(shell).exec("/sbin/reboot");
							status.clientStatus.put(client.clientip, ClientStatus.REBOOTING);
						}
					} catch (IOException e) {
						status.clientStatus.put(client.clientip, ClientStatus.ERROR);
					}
			    }
			});
		}
		tp.shutdown();
		
		try {
			tp.awaitTermination( clients.length * 5, TimeUnit.SECONDS );
		} catch ( InterruptedException e ) {
			Thread.currentThread().interrupt();
			return false;
		}
		
		// wait for rebooting clients to finish rebooting
		List<String> rebootingClients = new ArrayList<>();
		for (Entry<String, ClientStatus> entry : status.clientStatus.entrySet()) {
			if (entry.getValue() == ClientStatus.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) {
								Thread.currentThread().interrupt();
								return;
							}
				    	}
				    	status.clientStatus.put(ip, ClientStatus.ONLINE);
				    }
				});
			}
			statusTP.shutdown();
			
			try {
				statusTP.awaitTermination( 180, TimeUnit.SECONDS );
			} catch ( InterruptedException e ) {
				Thread.currentThread().interrupt();
				return false;
			}
		}
		
		// change status of clients that got stuck because of timeouts
		for (Map.Entry<String, ClientStatus> entry : status.clientStatus.entrySet()) {
			ClientStatus value = entry.getValue();
		    if (value == ClientStatus.CONNECTING || value == ClientStatus.REBOOTING) {
		    	entry.setValue(ClientStatus.ERROR);
		    }
		}
		
		return true;
	}
	
	
	private boolean isOnline(String address)
	{ 
		try ( Socket s = new Socket() ) {
			s.connect( new InetSocketAddress( address, port ), 3000 );
			return true;
		} catch ( Exception ex ) {
		}
		return false;
	}
	
	
	/**
	 * Output - contains additional status data of this task
	 */
	@SuppressWarnings( "unused" )
	static class Output
	{
		private Map<String, ClientStatus> clientStatus = new ConcurrentHashMap<>();
		private Client[] clients;
		private String time;
		private String locationId;
		private String locationName;
	}
	
	@SuppressWarnings( "unused" )
	static class Client
	{
		@Expose
		private String machineuuid;
		@Expose
		private String clientip;
	}
	
	static enum ClientStatus {
		CONNECTING, REBOOTING, SHUTDOWN, SHUTDOWN_AT, ONLINE, ERROR;
	}
	
}