/* options_tcp.c
* Created March/2013
* By Ron Bowes
*
* See LICENSE.md
*/
#include <assert.h>
#include <errno.h>
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <sys/socket.h>
#include "buffer.h"
#include "log.h"
#include "memory.h"
#include "select_group.h"
#include "session.h"
#include "tcp.h"
#include "ui_stdin.h"
typedef struct
{
session_t *session;
int s;
char *host;
uint16_t port;
/* This is for buffering data until we get a full packet */
buffer_t *buffer;
select_group_t *group;
/* The UI */
ui_stdin_t *ui_stdin;
} options_t;
#define DEFAULT_HOST "localhost"
#define DEFAULT_PORT 4444
options_t *options = NULL;
static SELECT_RESPONSE_t timeout(void *group, void *param)
{
options_t *options = (options_t*) param;
session_do_actions(options->session);
return SELECT_OK;
}
static SELECT_RESPONSE_t recv_callback(void *group, int s, uint8_t *data, size_t length, char *addr, uint16_t port, void *param)
{
options_t *options = (options_t*)param;
/* Cleanup - if the buffer is empty, reset it */
if(buffer_get_remaining_bytes(options->buffer) == 0)
buffer_clear(options->buffer);
buffer_add_bytes(options->buffer, data, length);
/* If we have at least a length value */
if(buffer_get_remaining_bytes(options->buffer) >= 2)
{
/* Read the length. */
uint16_t expected_length = buffer_peek_next_int16(options->buffer);
/* Check if we have the full length. */
if(buffer_get_remaining_bytes(options->buffer) - 2 >= expected_length)
{
uint8_t *data;
size_t returned_length;
/* Consume the value we already know */
buffer_read_next_int16(options->buffer);
/* Read the rest of the buffer. */
data = buffer_read_remaining_bytes(options->buffer, &returned_length, expected_length, TRUE);
/* Sanity check. */
assert(expected_length == returned_length);
/* Do the callback. */
session_recv(options->session, data, returned_length);
/* Free it. */
safe_free(data);
/* Clear the buffer if it's empty. */
if(buffer_get_remaining_bytes(options->buffer) == 0)
buffer_clear(options->buffer);
}
}
return SELECT_OK;
}
void tcpcat_close(options_t *options)
{
LOG_INFO("Close()");
assert(options->s && options->s != -1); /* We can't close a closed socket */
/* Remove from the select_group */
select_group_remove_and_close_socket(options->group, options->s);
options->s = -1;
}
static SELECT_RESPONSE_t closed_callback(void *group, int s, void *param)
{
options_t *options = (options_t*)param;
LOG_ERROR("Connection closed");
tcpcat_close(options);
return SELECT_OK;
}
void tcpcat_send(uint8_t *data, size_t length, void *d)
{
options_t *options = (options_t*) d;
buffer_t *buffer;
uint8_t *encoded_data;
size_t encoded_length;
if(options->s == -1)
{
/* Attempt a TCP connection */
LOG_INFO("Connecting to %s:%d", options->host, options->port);
options->s = tcp_connect(options->host, options->port);
/* If it fails, just return (it will try again next send) */
if(options->s == -1)
{
LOG_FATAL("Connection failed!");
exit(1);
}
/* If it succeeds, add it to the select_group */
select_group_add_socket(options->group, options->s, SOCKET_TYPE_STREAM, options);
select_set_recv(options->group, options->s, recv_callback);
select_set_closed(options->group, options->s, closed_callback);
}
assert(options->s != -1); /* Make sure we have a valid socket. */
assert(data); /* Make sure they aren't trying to send NULL. */
assert(length > 0); /* Make sure they aren't trying to send 0 bytes. */
buffer = buffer_create(BO_BIG_ENDIAN);
buffer_add_int16(buffer, length);
buffer_add_bytes(buffer, data, length);
encoded_data = buffer_create_string_and_destroy(buffer, &encoded_length);
if(tcp_send(options->s, encoded_data, encoded_length) == -1)
{
LOG_ERROR("[[TCP]] send error, closing socket!");
tcpcat_close(options);
}
}
void cleanup()
{
LOG_INFO("[[tcpcat]] :: Terminating");
if(options)
{
session_destroy(options->session, options->group);
/* Ensure the options is closed */
if(options->s != -1)
tcpcat_close(options);
buffer_destroy(options->buffer);
safe_free(options);
options = NULL;
}
print_memory();
}
int main(int argc, char *argv[])
{
/* Define the options specific to the DNS protocol. */
struct option long_options[] =
{
{"host", required_argument, 0, 0},
{"port", required_argument, 0, 0},
{0, 0, 0, 0} /* End */
};
char c;
int option_index;
const char *option_name;
options = safe_malloc(sizeof(options_t));
srand(time(NULL));
/* Set up some default options. */
options->s = -1;
options->host = DEFAULT_HOST;
options->port = DEFAULT_PORT;
options->buffer = buffer_create(BO_BIG_ENDIAN);
options->group = select_group_create();
options->session = session_create(options->group, tcpcat_send, options, 65535);
/* Parse the command line options. */
opterr = 0;
while((c = getopt_long_only(argc, argv, "", long_options, &option_index)) != EOF)
{
switch(c)
{
case 0:
option_name = long_options[option_index].name;
if(!strcmp(option_name, "host"))
{
options->host = optarg;
}
else if(!strcmp(option_name, "port"))
{
options->port = atoi(optarg);
}
else
{
LOG_FATAL("Unknown option: %s\n", option_name);
exit(1);
/* TODO: Usage */
}
break;
case '?':
default:
/* Do nothing; we expect some unknown arguments. */
break;
}
}
/* Tell the user what's going on */
LOG_INFO("Host: %s\n", options->host);
LOG_INFO("Port: %d\n", options->port);
atexit(cleanup);
/* Add the timeout function */
select_set_timeout(options->group, timeout, (void*)options);
while(TRUE)
select_group_do_select(options->group, 1000);
return 0;
}