jobserv/src/main/java/JobServ/JobServClient.java

package JobServ;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;

import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.Scanner;

// GRPC Client Class
public class JobServClient {

    // RequestMessage types
    // TODO: refactor to enum?
    public static final int OUTPUT = 0;
    public static final int RETURN = 1;
    public static final int PAUSE  = 2;
    public static final int RESUME = 3;
    public static final int STATUS = 4;

    /* 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 connects to server
    public JobServClient(String host, int port) {
        this(ManagedChannelBuilder.forAddress(host, port)
             // TODO: MTLS
             .usePlaintext()
             .build());
    }

    // private overload of constructor, used in the above constructor
    JobServClient(ManagedChannel channel) {
        this.channel = channel;
        blockingStub = ShellServerGrpc.newBlockingStub(channel);
    }

    public void shutdown() throws InterruptedException {
        channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
    }

    // sends the server a RequestMessage, blocks until response
    // returns an integer representing process status
    private int makeProcessRequest(int pid, int requestType) {
        logger.info("[+] Sending request");

        RequestMessage request = RequestMessage.newBuilder()
            .setPID(pid)
            .setRequestType(requestType)
            .build();
        JobStatusMessage response;

        try {
            // blocking network operation
            response = blockingStub.getStatusDetail(request);
        } catch (StatusRuntimeException e) {
            /* if this was using an async stub it might be
             * worthwhile to include the PID and request type 
             */
            logger.log(Level.WARNING, "[-]Status Request Failed: %s", e.getStatus());
            return -1;
        }

        // return process status
        return response.getProcessStatus();
    }

    // sends the server a request for output from PID
    // different from getProcessStatus in output expected from the server
    // returns process output as string
    public String getProcessOutput(int pid) {
        logger.info("[+] requesting output");

        RequestMessage request = RequestMessage.newBuilder()
            .setPID(pid)
            .setRequestType(OUTPUT)
            .build();
        OutputMessage response;

        try {
            // blocking network operation
            response = blockingStub.getJobOutput(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "[-] Request for output failed: %s", e.getStatus());
            // TODO: refactor this out by throwing here and catching in the shell
            return "";
        }

        return response.getProcessOutput();
    }

    // sends the server a command for a new job, blocks until response
    // 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();
        JobStatusMessage response;

        try {
            // blocking network operation
            response = blockingStub.makeNewJob(request);
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "[-] Request for new job failed!");
            return -1;
        }

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

        return response.getPID();
    }

    // Client entrypoint
    public static void main(String[] args) throws Exception {
        // check args
        if(args.length != 3) {
            System.out.println("Usage: $ jobservclient host port");
            return;
        }

        // start client (or fail if port is improperly formatted)
        JobServClient client;
        try {
            client = new JobServClient(args[1], Integer.parseInt(args[2]));
        } catch (NumberFormatException e) {
            System.out.println("Invalid Port");
            return;
        }

        Scanner reader = new Scanner(System.in);
        while(true) {
            System.out.print("> ");
            int pid, status;
            String input = reader.next();
            switch (input) {

            case "pause":
                pid = getPid();
                status = client.makeProcessRequest(pid, PAUSE);
                // TODO: parse status return
                break;

            case "resume":
                pid = getPid();
                status = client.makeProcessRequest(pid, RESUME);
                // TODO: parse status return
                break;

            case "new":
                System.out.println("Enter a command");
                String command = reader.next();
                pid = client.sendNewJobMessage(command);
                System.out.println(String.format("New process on server: %d", pid));
                break;

            case "output":
                pid = getPid();
                String out = client.getProcessOutput(pid);
                System.out.println(out);
                break;

            case "status":
                pid = getPid();
                status = client.makeProcessRequest(pid, STATUS);
                System.out.println(String.format("Current status of program is: %d", status));
                break;

            case "return":
                pid = getPid();
                status = client.makeProcessRequest(pid, RETURN);
                System.out.println(String.format("Exit code of process is: %d", status));
                break;

            case "quit":
                reader.close();
                return;

            case "help":
                System.out.println("pause:  pauses a process on the server");
                System.out.println("resume: resumes a process on the server");
                System.out.println("new:    starts a new process on the server");
                System.out.println("output: garners (new) output from a process on the server");
                System.out.println("status: outputs the current status of a program on the server");
                System.out.println("return: outputs exit code of a process on the server");
                System.out.println("quit:   exits this server");
                break;

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

        private static int getPid(Scanner reader) {
            System.out.println("Enter a PID");
            int pid;

            while(true) {
                try {
                    pid = reader.nextInt();

                } catch (InputMismatchException e) {
                    System.out.println("That was not a valid integer");
                    continue;
                }

                break;
            }
        }
    }
}