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.Iterator;
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 );

		if ( minutes < 0 ) {
			status.addError( "Delay cannot be negative" );
		}
		if ( sshkey == null || sshkey.length() == 0 ) {
			status.addError( "No SSH key given" );
		}
		if ( port < 1 || port > 65535 ) {
			status.addError( "Invalid port number" );
		}

		if ( status.error != null )
			return false;

		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;
		final List<Client> rebootingClients = new ArrayList<>();
		// 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 || client.machineuuid == null ) {
				status.addError( "null Client or missing ip/uuid in list, ignoring." );
				continue;
			}
			status.clientStatus.put(client.machineuuid, 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.machineuuid, minutes == 0 ? ClientStatus.SHUTDOWN : ClientStatus.SHUTDOWN_AT);
						} else {
							new Shell.Empty(shell).exec("/sbin/reboot");
							status.clientStatus.put(client.machineuuid, ClientStatus.REBOOTING);
							rebootingClients.add( client );
						}
					} catch (IOException e) {
						status.clientStatus.put(client.machineuuid, ClientStatus.ERROR);
					}
				}
			});
		}
		tp.shutdown();
		
		try {
			tp.awaitTermination( clients.length * 5, TimeUnit.SECONDS );
		} catch ( InterruptedException e ) {
			Thread.currentThread().interrupt();
			return false;
		}
		
		if (rebootingClients.size() > 0) {
			// Give about 3 minutes for reboot, should be plenty
			final int[] ports;
			if ( this.port == 22 ) {
				ports = new int[] { 22, 445 };
			} else {
				ports = new int[] { this.port, 22, 445 };
			}
			try {
				Thread.sleep(30000);
			} catch (InterruptedException e) {
				Thread.currentThread().interrupt();
				return false;
			}
			long lastcheck = System.currentTimeMillis();
			long deadline = lastcheck + 120 * 1000;
			while ( rebootingClients.size() > 0 && System.currentTimeMillis() < deadline ) {
				long delay = 10000 - (System.currentTimeMillis() - lastcheck);
				if ( delay > 0 ) {
					try {
						Thread.sleep(delay);
					} catch (InterruptedException e) {
						Thread.currentThread().interrupt();
						return false;
					}
				}
				lastcheck = System.currentTimeMillis();
				Iterator<Client> it = rebootingClients.iterator();
				while (it.hasNext()) {
					Client client = it.next();
					if (isOnline(client.clientip, ports)) {
						it.remove();
						status.clientStatus.put(client.machineuuid, ClientStatus.ONLINE);
					}
				}
			}
		}
		
		// 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, int... ports)
	{
		for ( int port : ports ) {
			try ( Socket s = new Socket() ) {
				s.connect( new InetSocketAddress( address, port ), 1000 );
				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;
		private String error;
		private void addError( String e )
		{
			if ( error == null ) {
				error = e + "\n";
			} else {
				error += e + "\n";
			}
		}
	}
	
	@SuppressWarnings( "unused" )
	static class Client
	{
		@Expose
		private String machineuuid;
		@Expose
		private String clientip;
	}
	
	static enum ClientStatus {
		CONNECTING, REBOOTING, SHUTDOWN, SHUTDOWN_AT, ONLINE, ERROR;
	}
	
}