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

package JobServ;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;
import java.util.InputMismatchException;
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 {

    /* 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 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");

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

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

        return response.getOutput();
    }

    // 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();
        PIDMessage response;

        try {
            // blocking network operation
            response = blockingStub.newJob(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();
    }

    // requests running status of job
    // returns true if job 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, "[-] Request for status failed!");
            return false;
        }

        return response.getIsRunning();
    }

    // sends PID to server expecting the return cod eof a process
    // function 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, "[-] Failed to get return code!");
            return 278;
        }

        return response.getProcessReturnCode();
    }

    // send a PID to be killed, 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, "[-] Failed to send request!");
            return;
        }

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

    // Client entrypoint
    public static void main(String[] args) throws Exception {
        if (args.length == 1 && args[0] == "--help"){
            outputHelp();
        }

        // check args
        if (args.length < 3) {
            System.out.println("Usage: $ jobservclient host port command");
            System.out.println("Or try client --help");
            return;
        }

        // start client (or fail if port is improperly formatted)
        JobServClient client;
        try {
            client = new JobServClient(args[0], Integer.parseInt(args[1]));

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

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

                String command = "";
                for (int token = 3; 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 != 4) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

                try {
                    candidatePid = Integer.parseInt(args[3]);
                } catch (InputMismatchException e) {
                    System.out.println(args[3] + " 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 != 4) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

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

                } catch (InputMismatchException e) {
                    System.out.println(args[3] + " 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 != 4) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

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

                } catch (InputMismatchException e) {
                    System.out.println(args[3] + " 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 != 4) {
                    System.out.println("Improper formatting, try client --help");
                    break;
                }

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

                } catch (InputMismatchException e) {
                    System.out.println(args[3] + " 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");
    }
}