/* doscan - Denial Of Service Capable Auditing of Networks
 * Copyright (C) 2003 Florian Weimer
 *
 * 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 2 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, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#include "config.h"
#include "engine_tcp.h"
#include "opt.h"
#include "results.h"
#include "scan.h"

#include <cerrno>
#include <fcntl.h>
#include <netinet/in.h>
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>

static void post_connect (scan_host_t *);
/* Initializes poll->fd and error fields. */

static void post_receive (struct scan_host_t *s, char *buffer, unsigned size,
                          engine_tcp_receive_callback_t completion,
                          scan_callback_t scan_callback);
/* Post a receive request, with flexible callback. */

static void cb_connect (scan_host_t *);
static void cb_send (scan_host_t *);
static void cb_receive (scan_host_t *);
static void cb_receive_packet (scan_host_t *);
static void cb_close (scan_host_t *);
static void free_buffer (scan_host_t *);

void
engine_tcp_open (struct scan_host_t *s,
                 engine_tcp_callback_t completion_callback,
                 engine_tcp_callback_t close_callback)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  if (s->poll->fd != -1) {
    close(s->poll->fd);
  }

  /* Initialize engine structure. */

  e->buffer = 0;
  e->size = 0;
  e->offset = 0;
  e->regexp = 0;
  e->copy = static_cast<engine_tcp_copy_t>(0);
  e->completion_callback = completion_callback;
  e->close_callback = close_callback;
  e->connect_time = ticks_get_cached ();

  post_connect (s);
  if (s->poll->fd == -1) {
    return;
  }

  s->timeout = ticks_get_cached () + opt_connect_timeout;
  s->close_callback = cb_close;
  s->io_callback = cb_connect;
}

void
engine_tcp_send (struct scan_host_t *s, const char *buffer, unsigned size,
                 engine_tcp_copy_t copy, engine_tcp_callback_t completion)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  /* Make a copy of the buffer if necessary. */

  e->copy = copy;
  e->offset = 0;
  switch (copy) {
  case ENGINE_TCP_MAKE_COPY:
    e->buffer = new char[size];
    memcpy (e->buffer, buffer, size);
    break;

  case ENGINE_TCP_NO_COPY:
  case ENGINE_TCP_NO_COPY_AND_FREE:
    e->buffer = (char *)buffer;
    break;

  default:
    abort ();
  }
  e->size = size;

  s->poll->events = POLLOUT;
  e->completion_callback = completion;
  s->io_callback = cb_send;
  s->timeout = ticks_get_cached () + opt_write_timeout;
}

void
engine_tcp_receive (struct scan_host_t *s, char *buffer, unsigned size,
                    engine_tcp_receive_callback_t completion)
{
  post_receive (s, buffer, size, completion, cb_receive);
}

void
engine_tcp_receive_until_match (struct scan_host_t *s, pcre *regexp,
                                char *buffer, unsigned size,
                                engine_tcp_receive_callback_t completion)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  engine_tcp_receive (s, buffer, size, completion);
  e->regexp = regexp;
}

void engine_tcp_receive_packet (struct scan_host_t *s, unsigned size,
                                engine_tcp_receive_callback_t completion)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  e->buffer = 0;
  e->size = size;
  e->offset = 0;
  e->regexp = 0;
  e->copy = static_cast<engine_tcp_copy_t>(0);

  s->poll->events = POLLIN;
  e->completion_callback = 0;
  e->receive_callback = completion;
  s->io_callback = cb_receive_packet;
  s->timeout = ticks_get_cached () + opt_read_timeout;
}

ticks_t
engine_tcp_connect_time (struct scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  return e->connect_time;
}

void
engine_tcp_close (struct scan_host_t *s)
{
  free_buffer (s);

  s->io_callback = 0;
  s->timeout = 0;
}


/* Internal functions (mostly callbacks for proto_tcp.c) */

static void
post_connect (scan_host_t *s)
{
  int sockfd, flags;
  struct sockaddr_in sa;

  s->poll->fd = -1;

  sockfd = socket (PF_INET, SOCK_STREAM, 0);
  if (sockfd == -1) {
    int err = errno;
    ipv4_string_t a;

    /* If we encounter an error at this point, it is not actually
       network-related, so we bail out immediately. */

    ipv4_address_to_string (s->host, a);
    fprintf (stderr, "%s: could not create socket for %s, error was: %s\n",
             opt_program, a, strerror (err));
    fprintf (stderr, "%s: try the '--connections' option with a smaller value\n",
             opt_program);
    exit (EXIT_FAILURE);
  }

  /* Make socket non-blocking. */
  flags = fcntl (sockfd, F_GETFL, 0);
  if (fcntl (sockfd, F_SETFL, flags | O_NONBLOCK) == -1) {
    int err = errno;
    ipv4_string_t a;

    /* Again, this error is not network-related. */

    ipv4_address_to_string (s->host, a);
    fprintf (stderr, "%s: could not set non-blocking mode for %s, error was: %s\n",
             opt_program, a, strerror (err));
    exit (EXIT_FAILURE);
  }

  memset (&sa, 0, sizeof (sa));
  sa.sin_family = AF_INET;
  sa.sin_port = htons (opt_port);
  sa.sin_addr.s_addr = htonl (s->host);

  if (connect (sockfd, (struct sockaddr *)&sa, sizeof (sa)) == -1) {
    int err = errno;

    if (err != EINPROGRESS) {
      close (sockfd);
      if (opt_net_errors) {
        results_add (ticks_get_cached (), s->host, err, 0, 0);
      }
      return;
    }
  }

  s->poll->fd = sockfd;
  s->poll->events = POLLIN | POLLOUT;
  s->poll->revents = 0;
}

static void
post_receive (struct scan_host_t *s, char *buffer, unsigned size,
              engine_tcp_receive_callback_t completion,
              scan_callback_t scan_callback)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  if (buffer == 0) {
    e->copy = ENGINE_TCP_NO_COPY_AND_FREE;
    e->buffer = new char[size];
  } else {
    e->copy = ENGINE_TCP_NO_COPY;
    e->buffer = buffer;
  }
  e->size = size;
  e->offset = 0;
  e->regexp = 0;

  s->poll->events = POLLIN;
  e->completion_callback = 0;
  e->receive_callback = completion;
  s->io_callback = scan_callback;
  s->timeout = ticks_get_cached () + opt_read_timeout;
}

static void
cb_connect (scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  /* Fetch error status in a portable way.
     We cannot depend on POLLERR because Solaris does not set it. */

  int error = 0;
  socklen_t len = sizeof (error);

  errno = 0;
  if (getsockopt(s->poll->fd, SOL_SOCKET, SO_ERROR, &error, &len) >= 0) {
    /* non-Solaris case: error is overwritten with the error code,
       nothing to do. */

  } else {
    /* Solaris case: errno has the error code. */
    error = errno;
  }

  if (error) {
    if (opt_net_errors) {
      results_add (ticks_get_cached (), s->host, error, 0, 0);
    }
    s->io_callback = 0;
    s->timeout = 0;
    return;
  }

  e->connect_time = ticks_get_cached ();

  s->io_callback = 0;
  s->timeout = 0;
  e->completion_callback (s);
}

static void
cb_send (scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);
  int result;

  result = write (s->poll->fd, e->buffer + e->offset,
                  e->size - e->offset);

  if (result == -1) {
    /* Log error and close. */

    free_buffer (s);

    results_add (e->connect_time, s->host, errno, 0, 0);
    s->io_callback = 0;
    s->timeout = 0;
    return;
  }

  e->offset += result;

  if (e->offset == e->size) {
    /* Send operation complete. */

    free_buffer (s);

    s->io_callback = 0;
    s->timeout = 0;
    e->completion_callback (s);

  } else {
    s->timeout = ticks_get_cached () + opt_write_timeout;
  }
}

static void
cb_receive (scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);
  int result;
  int match = 0;

  result = read (s->poll->fd, e->buffer + e->offset,
                 e->size - e->offset);

  if (result == -1) {
    /* Log error and close. */

    results_add (e->connect_time, s->host, errno, 0, 0);
    s->io_callback = 0;
    s->timeout = 0;
    free_buffer (s);
    return;
  }

  e->offset += result;

  if (e->regexp) {
    int rc = pcre_exec (e->regexp, 0, e->buffer, e->offset, 0, 0, 0, 0);

    if (rc < 0) {
      if (rc == PCRE_ERROR_NOMATCH) {
        match = 0;
      } else {
        fprintf (stderr, "%s: internal PCRE error after receive, "
                 "rc %d is invalid",
                 opt_program, rc);
        exit (EXIT_FAILURE);
      }
    } else {
      match = 1;
    }
  }

  /* Check for full buffer. */

  if ((result == 0) || (e->offset == e->size)) {

    s->io_callback = 0;
    s->timeout = 0;

    if (e->regexp && !match) {
      /* Receive failed, returned data does not match. */

      results_add (e->connect_time, s->host, RESULTS_ERROR_NOMATCH,
                   e->buffer, e->offset);
    } else {
      e->receive_callback (s, e->buffer, e->offset);
    }

    free_buffer (s);
    return;
  }

  /* Check for match. */

  if (match) {
    s->io_callback = 0;
    s->timeout = 0;
    e->receive_callback (s, e->buffer, e->offset);
    free_buffer (s);
    return;
  }

  /* Otherwise continue listening. */

  s->timeout = ticks_get_cached () + opt_read_timeout;
}

static void
cb_receive_packet (scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);
  char buffer[e->size];
  int result;

  result = read (s->poll->fd, buffer, e->size);

  if (result == -1) {
    /* Log error and close. */

    results_add (e->connect_time, s->host, errno, 0, 0);
    s->io_callback = 0;
    s->timeout = 0;
    free_buffer (s);
    return;
  }

  s->io_callback = 0;
  s->timeout = 0;

  if (result == 0) {
    /* Peer closed the connection. */

    results_add (e->connect_time, s->host, RESULTS_ERROR_NOMATCH, 0, 0);

  } else {
    e->receive_callback (s, buffer, result);
  }
}

static void
cb_close (scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  if (e->close_callback) {
    e->close_callback (s);
  }
  free_buffer (s);
}

static void
free_buffer (scan_host_t *s)
{
  engine_tcp_t *e = static_cast<engine_tcp_t*>(s->state);

  switch (e->copy) {
  case 0:
  case ENGINE_TCP_NO_COPY:
    /* No free() necessary. */
    break;

  case ENGINE_TCP_MAKE_COPY:
  case ENGINE_TCP_NO_COPY_AND_FREE:
    delete [] e->buffer;
    break;

  default:
    abort ();
  }

  e->copy = static_cast<engine_tcp_copy_t>(0);
  e->buffer = 0;
  e->size = 0;
  e->offset = 0;
}

/* arch-tag: e20f7aab-e6a3-48e7-8b01-18153aeda3f6
 */