/* Sock v0.3
 * By Spockie / Brigade (spockie@brigade.dhs.org)
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>

int main(int argc, char *argv[])
{
  int sock, port, mode, nread, buf_len = 1024, sin_s;
  struct hostent *host;
  struct sockaddr_in remote;
  unsigned char string[buf_len];

  fd_set fdset;
  memset(string, 0, buf_len);
  FD_ZERO(&fdset);

  fprintf(stderr, "Sock v0.3 by spockie@brigade.dhs.org\n");

  if (argc != 3) {
    fprintf(stderr, "Usage: %s hostname|-l port\n", argv[0]);
    exit(1);
  }

  if (argv[2])
    port = atoi(argv[2]);

  if ((strcmp(argv[1], "-l")) == 0)
    mode = 1;
  else
    mode = 0;

  if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
    perror("socket");
    exit(1);
  }

  if (mode == 0) {
    if ((host = gethostbyname(argv[1])) == NULL) {
      herror("gethostbyname");
      exit(1);
    }

    remote.sin_family = AF_INET;
    remote.sin_addr.s_addr = *(long *) (host->h_addr);
    remote.sin_port = htons(port);

    if ((connect(sock, (struct sockaddr *) &remote, sizeof(remote))) < 0) {
      perror("connect");
      exit(1);
    }
    fprintf(stderr, "Connected to %s.\n", inet_ntoa(remote.sin_addr));
  }

  if (mode == 1) {
    struct sockaddr_in local;
    local.sin_family = AF_INET;
    local.sin_addr.s_addr = INADDR_ANY;
    local.sin_port = htons(port);

    if (bind(sock, (struct sockaddr *) &local, sizeof(struct sockaddr)) == -1) {
      perror("bind");
      exit(1);
    }

    if (listen(sock, 1) == -1) {
      perror("listen");
      exit(1);
    }
    sin_s = sizeof(struct sockaddr_in);
    fprintf(stderr, "Waiting for connection..\n");

    if ((sock = accept(sock, (struct sockaddr *) &remote, &sin_s)) == -1)
      perror("accept");

    fprintf(stderr, "Connection from %s\n", inet_ntoa(remote.sin_addr));
  }

  for (;;) {
    FD_SET(sock, &fdset);
    FD_SET(0, &fdset);

    if (select(sock + 1, &fdset, NULL, NULL, NULL) < 0) {
      fprintf(stderr, "Selective error!\n");
      exit(1);
    }

    if (FD_ISSET(0, &fdset)) {
      if ((nread = read(0, string, buf_len)) < 0) {
        fprintf(stderr, "Stdin read error\n");
        break;
      } else if (nread == 0) {
        fprintf(stderr, "Connection closed.\n");
        break;
      }

      send(sock, string, strlen(string), 0);
      memset(string, 0, buf_len);
    }

    if (FD_ISSET(sock, &fdset)) {
      if ((nread = recv(sock, string, buf_len, 0)) < 0) {
        fprintf(stderr, "Network read error\n");
        break;
      } else if (nread == 0) {
        fprintf(stderr, "Connection closed by foreign host.\n");
        break;
      }

      printf("%s", string);
      memset(string, 0, buf_len);
    }
  }

  if ((close(sock)) < 0) {
    perror("close");
    exit(1);
  }

  fprintf(stderr, "Program finished.\n");
}