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

package org.openslx.taskmanager.tasks;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import org.apache.logging.log4j.util.Strings;
import org.openslx.satserver.util.MessageSink;
import org.openslx.satserver.util.Util;
import org.openslx.satserver.util.WakeOnLanExecutor;
import org.openslx.taskmanager.api.AbstractTask;
import org.openslx.taskmanager.tasks.RemoteExec.Output;

import com.google.gson.annotations.Expose;

public class WakeOnLan extends AbstractTask
{

	/**
	 * List of clients to wake up
	 */
	@Expose
	private List<Machine> clients;

	/**
	 * SHH data (key, command) for waking via another host
	 */
	@Expose
	private Map<String, SshData> ssh;

	/**
	 * UDP port to send WOL packets to
	 */
	@Expose
	private int port;

	private StatusObject status;

	@Override
	protected boolean initTask()
	{
		status = new StatusObject();
		this.setStatusObject( status );
		if ( clients == null ) {
			status.addMsg( "Clients null" );
		} else {
			// Clean
			for ( Iterator<Machine> it = this.clients.iterator(); it.hasNext(); ) {
				Machine client = it.next();
				if ( client == null || Util.isEmpty( client.mac ) || Util.isEmpty( client.ip )
						|| client.methods == null || client.methods.isEmpty() ) {
					it.remove();
				}
			}
			if ( clients.isEmpty() ) {
				status.addMsg( "Clients empty" );
			}
		}
		if ( !Util.isEmpty( status.messages ) )
			return false;
		return true;
	}

	@Override
	protected boolean execute()
	{
		// Loop over clients until they all were handled (i.e. methods is empty)
		ExecutorService tp = Executors.newFixedThreadPool( ssh.size() > 4 ? 4 : ssh.size() );
		Map<String, ArrayList<String>> byMethod;
		do {
			byMethod = new HashMap<>();
			// Fetch next method for all clients
			for ( Machine client : clients ) {
				if ( client.methods.isEmpty() )
					continue;
				// Group by method and destination IP address
				String method = client.methods.remove( 0 ) + "//" + client.ip;
				ArrayList<String> list = byMethod.get( method );
				if ( list == null ) {
					list = new ArrayList<>();
					byMethod.put( method, list );
				}
				list.add( client.mac );
			}
			// Execute
			List<Future<?>> waitList = new ArrayList<>();
			for ( Entry<String, ArrayList<String>> it : byMethod.entrySet() ) {
				String[] parts = it.getKey().split( "//" );
				String method = parts[0];
				String ip = parts[1];
				if ( method.equalsIgnoreCase( "DIRECT" ) ) {
					// Directly from server
					waitList.add( tp.submit( () -> {
						WakeOnLanExecutor wol = new WakeOnLanExecutor( null );
						String[] success = wol.execute( status, it.getValue(), ip, this.port );
						markSuccess( Arrays.asList( success ) );
					} ) );
				} else if ( this.ssh.containsKey( method ) ) {
					// Via SSH
					waitList.add( tp.submit( () -> {
						SshData sshData = this.ssh.get( method );
						String macs = Strings.join( it.getValue().iterator(), ' ' );
						String command = sshData.command.replace( "%MACS%", macs ).replace( "%IP%", ip );
						RemoteExec.Client c = new RemoteExec.Client( "x", sshData.ip, sshData.port, sshData.username );
						RemoteExec task = new RemoteExec( new RemoteExec.Client[] { c }, sshData.sshkey, sshData.port, command, 5 );
						task.execute();
						Output s = task.getStatusObject();
						if ( s != null ) {
							if ( s.result != null && s.result.containsKey( "x" ) && s.result.get( "x" ).exitCode == 0 ) {
								markSuccess( Arrays.asList( macs ) );
							}
							if ( !Util.isEmpty( s.error ) ) {
								status.addMsg( s.error );
							}
						}
					} ) );
				} else {
					status.addMsg( "Ignoring unknown SSH method '" + method + "'" );
				}
			}
			// Wait for all jobs
			for ( Future<?> f : waitList ) {
				try {
					f.get();
				} catch ( InterruptedException e ) {
					status.addMsg( "Threadpool got interrupted" );
					Thread.currentThread().interrupt();
					return false;
				} catch ( ExecutionException e ) {
					status.addMsg( "A treadpool job threw an exception" );
					e.printStackTrace();
				}
			}
		} while ( !byMethod.isEmpty() );
		if ( this.clients.isEmpty() )
			return true;
		status.addMsg( "Failed clients:" );
		for ( Machine c : clients ) {
			status.addMsg( c.ip );
		}
		return false;
	}

	private void markSuccess( Collection<String> macs )
	{
		synchronized ( this.clients ) {
			for ( Iterator<Machine> it = this.clients.iterator(); it.hasNext(); ) {
				Machine client = it.next();
				if ( macs.contains( client.mac ) ) {
					it.remove();
				}
			}
		}
	}

	static class Machine
	{
		/**
		 * Destination IP (for directed broadcast)
		 */
		@Expose
		String ip;
		/**
		 * Destination MAC
		 */
		@Expose
		String mac;
		/**
		 * Desired WOL modes, from most to least preferred. either a key to the top level ssh map,
		 * or "DIRECT" for sending a WOL packet directly from the server
		 */
		@Expose
		List<String> methods;
		/**
		 * WOL password (currently unused by slx-admin)
		 */
		@Expose
		private String password;
	}

	static class SshData
	{
		@Expose
		String username;
		@Expose
		String sshkey;
		@Expose
		String ip;
		@Expose
		int port;
		@Expose
		String command;
	}

	/*
	 * 
	 */

	/**
	 * Return any messages about progress and status
	 */
	static class StatusObject implements MessageSink
	{
		private String messages = "";

		public synchronized void addMsg( String str )
		{
			messages = messages + "\n" + str;
		}
	}

}