/* 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
 */

/* Implementation notes
 *
 * This file contains the main scanning engine.  There is some room
 * for improvement.  We certainly should use a slow-start algorithm.
 * The current one suits the name of the tool rather nicely, but it
 * has a tendency of missing hosts, too. */

#include "config.h"
#include "opt.h"
#include "proto.h"
#include "results.h"
#include "scan.h"
#include "subnets.h"
#include "ticks.h"

#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <unistd.h>

static struct pollfd *polls = 0;
static scan_host_t *hosts = 0;

static unsigned hosts_used;
/* Number of entries in hosts actually touched so far. */

static unsigned hosts_active;
/* Hosts actually active (grows at the beginning, shrinks at the end).
   Always less than or equal to hosts_used. */

static void allocate (void);
/* Allocates polls, hosts and performs basic initialization
   (hosts[j].poll, hosts_used). */

ticks_t timeout;
/* Timeout for next poll */

static void update_timeout (ticks_t);
/* Update timeout value if argument is sooner. */

static void add_host (subnets& nets);
/* Adds a single host to the hosts array. */

static void replace_host (subnets& nets, scan_host_t *);
/* Replaces a given entry in the hosts array. */

static void add_more_hosts (subnets& nets);
/* Add more hosts if necessary. */

static void process_hosts (subnets& nets);
/* Scans the lists of hosts and invokes callbacks as necessary.
   Timed-out entries are removed. */

static unsigned relative_timeout (ticks_t base);
/* Returns the relativ time from timeout to base (truncated at
   zero). */

static void invoke_poll (void);
/* Invoke poll as required. */

static void show_progress (subnets& nets, int);
/* Show progress, if necessary. */

void
scan (subnets& nets)
{
  allocate ();
  if (!proto_start (nets)) {
    return;
  }

  for (;;) {
    ticks_get ();               /* update cache */
    timeout = TICKS_LAST;

    add_more_hosts (nets);

    process_hosts (nets);
    if (hosts_active == 0) {
      break;
    }

    show_progress (nets, 0);

    invoke_poll ();
  }
  show_progress (nets, 1);
}

static void
allocate (void)
{
  unsigned j;

  if (polls == 0) {
    polls = new pollfd[opt_fd_count];
  }
  if (hosts == 0) {
    hosts = new scan_host_t[opt_fd_count];
  }

  hosts_used = 0;

  for (j = 0; j < opt_fd_count; j++) {
    polls[j].fd = -1;
    polls[j].events = 0;
    polls[j].revents = 0;
    hosts[j].poll = polls + j;
    hosts[j].state = 0;
  }
}

static void
update_timeout (ticks_t new_timeout)
{
  if (new_timeout < timeout) {
    timeout = new_timeout;
  }
}


static void
add_host (subnets& nets)
{
  replace_host (nets, hosts + hosts_used);

  if (hosts[hosts_used].poll->fd != -1) {
    hosts_used++;
  }
}

static void
replace_host (subnets& nets, scan_host_t *s)
{
  /* Host was active, close socket. */
  if (s->poll->fd != -1) {
    if (s->close_callback) {
      s->close_callback (s);
    }
    close (s->poll->fd);
    hosts_active--;
  }

  s->poll->fd = -1;
  s->poll->events = 0;
  s->poll->revents = 0;

  while (s->poll->fd == -1) {
    if (nets.finished ()) {
      return;
    }

    s->host = nets.next ();
    s->poll->fd = -1;
    proto_open (s);
  }
  if (s->poll->events == 0) {
    abort ();
  }

  hosts_active++;

  if (s->timeout == 0) {
    abort ();
  }
  if (s->io_callback == 0) {
    abort ();
  }
  if (s->poll->events == 0) {
    abort ();
  }

  update_timeout (s->timeout);
}

static void
add_more_hosts (subnets& nets)
{
  static ticks_t add_timeout = 0;

  if (hosts_used < opt_fd_count) {
    if (!hosts_active || (ticks_get_cached () > add_timeout)) {
      unsigned to_add, j;

      to_add = opt_fd_count - hosts_used;
      if (to_add > opt_add_burst) {
        to_add = opt_add_burst;
      }

      for (j = 0; j < to_add; j++) {
        add_host (nets);
      }

      if (hosts_used != opt_fd_count) {
        add_timeout = ticks_get_cached () + opt_add_timeout;
        update_timeout (add_timeout);
      }
    } else {
      update_timeout (add_timeout);
    }
  }
}

static void
process_hosts (subnets& nets)
{
  unsigned j;
  ticks_t ticks = ticks_get_cached ();

  for (j = 0; j < hosts_used; j++) {
    if (hosts[j].poll->fd == -1) {
      /* unused entry */
      continue;
    }

    if (hosts[j].poll->revents) {
      hosts[j].io_callback (hosts + j);

    } else {
      if (ticks > hosts[j].timeout) {
        /* Entry expired. */
        replace_host (nets, hosts + j);
      }
    }

    if ((hosts[j].timeout == 0) || (hosts[j].io_callback == 0)) {
      /* Close socket as requested by callback. */
      replace_host (nets, hosts + j);
    }

    update_timeout (hosts[j].timeout);
  }
}

static unsigned
relative_timeout (ticks_t base)
{
  if (base > timeout) {
    return 0;
  } else {
    return timeout - base;
  }
}


static void
invoke_poll (void)
{
  unsigned poll_timeout = relative_timeout (ticks_get_cached ());
  unsigned min_poll_timeout;

  if (timeout == TICKS_LAST) {
    /* must not happen, all operations must be controlled by
       timeouts */
    abort ();
  }

  /* Limit the poll timeout so that we can take a few timeouts at
     once.  If we are still adding new hosts, use opt_add_timeout as
     limit, otherwise use min_poll_timeout. */
  if (hosts_used == opt_fd_count) {
    min_poll_timeout = opt_add_timeout;
  } else {
    min_poll_timeout = 100;
  }
  if (poll_timeout < min_poll_timeout) {
    poll_timeout = min_poll_timeout;
  }

  while (poll (polls, hosts_used, poll_timeout) == -1) {
    int err = errno;

    if (err == EINTR) {
      poll_timeout = relative_timeout (ticks_get ());
      continue;
    }

    fprintf (stderr, "%s: poll failed: %s",
             opt_program, strerror (err));
    exit (EXIT_FAILURE);
  }
}

static void
show_progress (subnets& nets, int final)
{
  static ticks_t progress_timeout = 0;
  int show = 0;

  if (! opt_indicator) {
    return;
  }

  if (progress_timeout < ticks_get_cached ()) {
    /* Update progress indicator at most every 500 milliseconds. */
    progress_timeout = ticks_get_cached () + 500;
    show = 1;
  }

  if (final) {
    show = 1;
  }

  if (show) {
    fprintf (stderr, "%u of %u connected, %u active, %u reported     %s",
             nets.hosts_processed(), nets.hosts_total(),
             hosts_active, results_count (), final ? "\n" : "\r");
    fflush (stderr);
  }
}
