/* 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 "ipv4.h"
#include "opt.h"
#include "proto.h"
#include "results.h"
#include "scan_tcp.h"
#include "scan_trigger.h"
#include "utils.h"

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

#include <pcre.h>

static bool proto_tcp_start(subnets&);
static void proto_tcp_open(scan_host_t *);

static const char *send_buffer = 0;
static unsigned send_length = 0;

static pcre *receive_regexp = 0;
static int receive_regexp_terminate_on_match;
static int receive_regexp_capture_count;

void
proto_tcp_register(void)
{
  proto_register("tcp", proto_tcp_start, proto_tcp_open);
}

class proto_tcp : public tcp_client_handler {
  int state;
  std::string receive_buffer;
  unsigned receive_offset;
  unsigned send_offset;

  void handle_close(int error, ticks_t = ticks_get_cached());

public:
  proto_tcp(event_queue&, ipv4_t);

  virtual bool on_activity(activity);
  virtual bool on_timeout(ticks_t);
};

proto_tcp::proto_tcp(event_queue& q, ipv4_t host)
  : tcp_client_handler(q, host, opt_port),
    state(0), receive_buffer(opt_banner_size, char()),
    receive_offset(0), send_offset(0)
{
  set_relative_timeout(opt_connect_timeout);
}

bool
proto_tcp::on_activity(activity act)
{
  int error, result;

  switch (state) {
  case 0:
    // We just got connected.

    // Check for error.

    error = get_error();
    if (error) {
      if (opt_net_errors) {
        results_add(ticks_get_cached(), host(), error, 0, 0);
      }
      return false;
    }

    // Regular processing.  We can disconnect immediately if we do not
    // collect banners.

    if (opt_banner_size == 0) {
      results_add(ticks_get_cached(), host(), 0, 0, 0);
      return false;
    }

    set_relative_timeout(opt_read_timeout);

    if (send_length > 0) {
      watch(watch_read_write);
    } else {
      watch(watch_read);
    }

    ++state;
    return true;

  case 1:
    // Established connection.  First check for error.

    if (act == activity_error) {
      handle_close(get_error());
      return false;
    }

    // Handle data which is available for reading.

    if (act == activity_read || act == activity_read_write) {
      result = read(fd(), &receive_buffer[0] + receive_offset,
                    receive_buffer.size() - receive_offset);

      if (result == -1) {
        handle_close(errno);
        return false;
      }

      receive_offset += result;

      if (receive_regexp_terminate_on_match) {
        int rc;

        rc = pcre_exec(receive_regexp, 0,
                       receive_buffer.data(), receive_offset, 0, 0, 0, 0);
        switch (rc) {
        case 0:
        case 1:
        case 2:
          // We have a match and can therefore drop the connection.
          handle_close(0);
          return false;

        case PCRE_ERROR_NOMATCH:
          /* No match, we have to continue. */
          break;

        default:
          fprintf(stderr, "%s: PCRE error %d during receive",
                  opt_program, rc);
          exit(EXIT_FAILURE);
        }
      }

      if (result == 0 || receive_offset == receive_buffer.size()) {
        handle_close(0);
        return false;
      }
    }


    if (act == activity_write || act == activity_read_write) {
      result = write(fd(), send_buffer + send_offset,
                     send_length - send_offset);

      if (result == -1) {
        handle_close(errno);
        return false;
      }

      send_offset += result;

      if (send_offset == send_length) {
        // No more data to send, just keep reading.
        watch(watch_read);
      }
    }

    set_relative_timeout(opt_read_timeout);

    // Stay in this state.

    return true;
  }

  abort();
}

bool
proto_tcp::on_timeout(ticks_t ticks)
{
  switch (state) {
  case 0:
    // Connection timeouts are normal.
    return false;

  case 1:
    handle_close(0, ticks);
    return false;
  }

  abort();
}

void
proto_tcp::handle_close(int error, ticks_t ticks)
{
  if (receive_regexp_capture_count > 0) {
    int rc;
    int ovector[6];

    rc = pcre_exec(receive_regexp, 0, receive_buffer.data(), receive_offset,
                   0, 0, ovector, 6);

    switch (rc) {
    case 2:
      results_add(ticks, host(), error,
                  receive_buffer.data() + ovector[2], ovector[3] - ovector[2]);
      break;

    case PCRE_ERROR_NOMATCH:
      results_add(ticks, host(),
                  error ? error : RESULTS_ERROR_NOMATCH,
                  receive_buffer.data(), receive_offset);
      break;

    case 0:
    case 1:
      fprintf(stderr, "%s: internal PCRE error after receive, "
              "rc %d is invalid", opt_program, rc);
      exit(EXIT_FAILURE);
      break;

    default:
      fprintf(stderr, "%s: PCRE error %d after receive",
              opt_program, rc);
      exit(EXIT_FAILURE);
    }

  } else {
    results_add(ticks, host(), error,
                receive_buffer.data(), receive_offset);
  }
}

static void
proto_tcp_open (scan_host_t *s)
{
}

static bool
proto_tcp_start(subnets& n)
{
  if (opt_receive && (opt_receive[0] != 0) && (opt_banner_size == 0)) {
    fprintf (stderr, "%s: --receive requires --banner option\n",
             opt_program);
    exit (EXIT_FAILURE);
  }

  if (opt_send && (opt_send[0] != 0)) {
    char *sb;
    string_dequote (opt_send, &sb, &send_length,
                    "--send option");
    send_buffer = sb;
  } else {
    send_buffer = 0;
    send_length = 0;
  }

  if (opt_receive && (opt_receive[0] != 0)) {
    /* Parse the regular expression passed to --receive. */

    const char* error;
    int offset;

    receive_regexp = pcre_compile (opt_receive, PCRE_ANCHORED | PCRE_DOLLAR_ENDONLY | PCRE_DOTALL, &error, &offset, 0);
    if (receive_regexp == 0) {
      fprintf (stderr, "%s: cannot compile '--receive' regexp\n",
               opt_program);
      fprintf (stderr, "%s: %s (at position %d)\n",
               opt_program, error, offset + 1);
      exit (EXIT_FAILURE);
    } else {
      receive_regexp_terminate_on_match = opt_receive[strlen (opt_receive) - 1] == '$';
      pcre_fullinfo (receive_regexp, 0, PCRE_INFO_CAPTURECOUNT, &receive_regexp_capture_count);
      if (receive_regexp_capture_count > 1) {
        fprintf (stderr, "%s: too many captures in '--receive' regexp\n",
                 opt_program);
        exit (EXIT_FAILURE);
      }
    }
  } else {
    receive_regexp_terminate_on_match = 0;
    receive_regexp_capture_count = 0;
  }

  std::auto_ptr<event_queue> q(event_queue::create(opt_fd_count));
  scan_trigger::default_handler<proto_tcp> th;
  scan_trigger t(*q, n, th, opt_fd_count, opt_add_timeout, opt_add_burst);

  q->run();
  return false;
}

// arch-tag: 5a507a50-e12b-4209-b6bc-fafa09da3a89