/*
 * JobServClient
 *
 * v1.0
 *
 * May 18, 2019
 */

package JobServ;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;
import io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.NettyChannelBuilder;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslContextBuilder;

import javax.net.ssl.SSLException;
import java.io.File;
import java.util.InputMismatchException;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.Scanner;

/*
 * The JobServClient class extends the gRPC stub code
 * Additionally, it plugs a command line interface into the API code.
 */
public class JobServClient {

    /*
     * The client should not use the same logging module as the server.
     * In a more robust product the server logging module will take advantage of system level
     * log aggregators such as journalctl, which the client should not be writing to on the users system
     */
    private static final Logger logger = Logger.getLogger(JobServClient.class.getName());

    private final ManagedChannel channel;

    /*
     * blockingStub is used when the client needs to block until the server responds
     * the client doesnt nessesarily need to support asynchronously firing off commands
     * in this shell-like interface it would be disconcerting to get multiple returns out of order
     */
    private final ShellServerGrpc.ShellServerBlockingStub blockingStub;

    /*
     * Constructor
     * Creates an SslContext from cert, key, and trust store
     * Creates a ManagedChannel object from SSL Parameters
     * Spawns a new blockingStub for network operations with the server
     */
    public JobServClient(String host,
                         int port,
                         String trustStore,
                         String clientCert,
                         String clientPrivateKey) throws SSLException {
        SslContextBuilder builder = GrpcSslContexts.forClient();
        builder.trustManager(new File(trustStore));
        builder.keyManager(new File(clientCert), new File(clientPrivateKey));

        this.channel = NettyChannelBuilder.forAddress(host, port)
            .sslContext(builder.build())
            .build();

        blockingStub = ShellServerGrpc.newBlockingStub(this.channel);
    }

    /*
     * shutdown()
     * Gets called when you press cntrl+c
     * takes at most 5 seconds to close its connection
     */
    public void shutdown() throws InterruptedException {
        channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
    }

    /*
     * getProcessInfo()
     * sends the server a request for output from the process identified by 'pid'
     * returns process output as string
     */
    public String getProcessOutput(int pid) {
        logger.info("[+] requesting output");

        PIDMessage request = PIDMessage.newBuilder()
            .setPid(pid)
            .build();
        OutputMessage response;

        try {
            // blocking network operation
            response = blockingStub.getOutput(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "(API Failure) Request for output failed: " + e.getStatus());
            return "<Error connecting to API>";
        }

        return response.getOutput();
    }

    /*
     * sendNewJobMessage()
     * sends a shell command to the api server
     * returns new pid of job
     */
    public int sendNewJobMessage(String command) {
        // thought of escaping this, but the vulnerability is only client side, from client user input.
        logger.info("[+] Sending command to server");

        NewJobMessage request = NewJobMessage.newBuilder()
            .setCommand(command)
            .build();
        PIDMessage response;

        try {
            // blocking network operation
            response = blockingStub.newJob(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "(API Failure) Request for new job failed: " + e.getStatus());
            return -1;
        }

        if(response.getPid() == -1) {
            logger.log(Level.WARNING, "New job creation failed server side!");
        }

        return response.getPid();
    }

    /*
     * getProcessStatus()
     * requests running status of process pid
     * returns true if process still running else false
     */
    public Boolean getProcessStatus(int pid) {
        logger.info("[+] Requesting status of a job");

        PIDMessage request = PIDMessage.newBuilder()
            .setPid(pid)
            .build();
        StatusMessage response;

        try {
            // blocking network operation
            response = blockingStub.getStatus(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "(API Failure) Request for status failed: " + e.getStatus());
            return false;
        }

        return response.getIsRunning();
    }

    /*
     * sends PID to server
     * returns process exit code
     * returns a 0-255 return code or 277 if still running
     * or 278 if error in API
     */
    public int getProcessReturn(int pid) {
        logger.info("[+] Requesting return code of a job");

        PIDMessage request = PIDMessage.newBuilder()
            .setPid(pid)
            .build();
        ReturnMessage response;

        try {
            // blocking network operation
            response = blockingStub.getReturn(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "(API Failure) Failed to get return code: " + e.getStatus());
            return 278;
        }

        return response.getProcessReturnCode();
    }

    /*
     * killProcess()
     * send a PID to be killed, function returns nothing
     * logs warning if job status comes back still running
     */
    public void killProcess(int pid) {
        logger.info("[+] Killing a job");

        PIDMessage request = PIDMessage.newBuilder()
            .setPid(pid)
            .build();
        StatusMessage response;

        try {
            // blocking network operation
            response = blockingStub.killJob(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "(API Failure) Failed to send request: " + e.getStatus());
            return;
        }

        if (response.getIsRunning()) {
            logger.log(Level.WARNING, "[-] Server failed to kill job!");
        }
    }

    /*
     * main()
     * Client entrypoint
     * Parses arguments and calls the correct function
     */
    public static void main(String[] args) throws Exception {
        if (args.length == 1 && args[0] == "help"){
            outputHelp();
        }

        // check args
        if (args.length < 7) {
            System.out.println("Usage: $ ./jobserv-client privatekey, cert, truststore, host, port, command, args");
            System.out.println("Or try $ ./jobserv-client help");
            return;
        }

        // start client
        // fails if port is improperly formatted or if an ssl exception occurs
        JobServClient client;
        try {
            client = new JobServClient(args[3], Integer.parseInt(args[4]), args[2], args[1], args[0]);

        } catch (NumberFormatException e) {
            System.out.println("Invalid Port");
            return;
        } catch (SSLException e) {
            System.out.println(e.getMessage());
            return;
        }

        // declare pid up here so that multiple switch cases can use it
        int candidatePid;
        // parse remaining args
        switch (args[5]) {
            case "new":
                if (args.length < 7) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

                String command = "";
                for (int token = 6; token < args.length; token++) {
                    command += " " + args[token];
                }

                int newProcess = client.sendNewJobMessage(command);
                System.out.printf("Process started, assigned pid is %d\n", newProcess);
                break;

            case "output":
                if (args.length != 7) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

                try {
                    candidatePid = Integer.parseInt(args[6]);
                } catch (InputMismatchException e) {
                    System.out.println(args[6] + " is not a valid int, much less a valid pid");
                    break;
                }

                String processOutput = client.getProcessOutput(candidatePid);
                System.out.println(processOutput);
                break;

            case "status":
                if (args.length != 7) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

                try {
                    candidatePid = Integer.parseInt(args[6]);

                } catch (InputMismatchException e) {
                    System.out.println(args[6] + " is not a valid int, much less a valid pid");
                    break;
                }

                Boolean processStatus = client.getProcessStatus(candidatePid);
                System.out.printf("Process is currently running? %b\n", processStatus);
                break;

            case "kill":
                if (args.length != 7) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

                try {
                    candidatePid = Integer.parseInt(args[6]);

                } catch (InputMismatchException e) {
                    System.out.println(args[6] + " is not a valid int, much less a valid pid");
                    break;
                }

                client.killProcess(candidatePid);
                System.out.println("End process request recieved!");
                break;

            case "return":
                if (args.length != 7) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

                try {
                    candidatePid = Integer.parseInt(args[6]);

                } catch (InputMismatchException e) {
                    System.out.println(args[6] + " is not a valid int, much less a valid pid");
                    break;
                }

                int returnCode = client.getProcessReturn(candidatePid);

                if (returnCode == 277) {
                    System.out.println("Process is still running");
                    break;

                } else if (returnCode == 278) {
                    System.out.println("RPC Call error!");
                    break;

                } else {
                    System.out.printf("Process Exit Code: %d\n", returnCode);
                }

            default:
                System.out.println("Improper command, try 'help'");
                break;
            }
        }

    public static void outputHelp() {
        System.out.println("... new (command)");
        System.out.println("Starts a new process on the server");
        System.out.println("... output (pid)");
        System.out.println("Garners output from process on server");
        System.out.println("... status (pid)");
        System.out.println("Returns whether process on server is running");
        System.out.println("... return (pid)");
        System.out.println("Collects return code from remote process");
        System.out.println("... kill (pid)");
        System.out.println("Immediately destroys remote process");
    }
}