/* 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 "opt.h"
#include "scan_udp.h"

#include <cstdio>
#include <cstring>
#include <string>
#include <unistd.h>

// scan_udp_single::handler

scan_udp_single::handler::~handler()
{
}

// scan_udp_single::send_trigger

scan_udp_single::send_trigger::send_trigger(scan_udp_single& s)
    : scanner(s)
{
}

void
scan_udp_single::send_trigger::connect(event_queue& q, scan_trigger& t,
                                       ipv4_t host)
{
  sender *s = new sender(scanner, host);
  scanner.senders[host] = s;
}

void
scan_udp_single::send_trigger::all_connected()
{
  scanner.the_listener->unwatch();
  close(scanner.listen_fd);
}

// scan_udp_single::sender

scan_udp_single::sender::sender(scan_udp_single& s, ipv4_t addr)
  : event_queue::handler(s.queue), scanner(s), host(addr),
    retries(scanner.retries)
{
  address.sin_family = AF_INET;
  address.sin_addr.s_addr = htonl(addr);
  address.sin_port = htons(opt_port);
  the_handler = scanner.handler_template.clone();
}

scan_udp_single::sender::~sender()
{
  delete the_handler;
  scanner.senders.erase(host);
  scanner.trigger.finish();
}

bool
scan_udp_single::sender::do_send()
{
  the_handler->create_query(host, data);
  int result = sendto(scanner.listen_fd, data.data(), data.size(),
                      0,        // flags
                      reinterpret_cast<sockaddr *>(&address),
                      sizeof(address));
  if (result == -1) {
    if (opt_net_errors) {
      results_add (ticks_get(), host, errno, 0, 0);
    }

    // Only some errors are fatal; in the case of others, we retry.
    switch (errno) {
    case ENETUNREACH:
    case EHOSTUNREACH:
    case EPERM:
    case EACCES:
      return false;
    }
  }
  return true;
}

void
scan_udp_single::sender::stop_sending()
{
  retries = -1;
  set_immediate_timeout();
}

bool
scan_udp_single::sender::on_timeout(ticks_t ticks)
{
  switch (retries) {
  case -1:
    return false;
  case 0:
    the_handler->timeout(host, ticks);
    return false;
  }

  do_send();
  --retries;
  set_relative_timeout(scanner.timeout);
  return true;
}

// scan_udp_single::listneer

scan_udp_single::listener::listener(scan_udp_single& s)
  : fd_handler(s.queue, s.listen_fd, watch_read), scanner(s)
{
  set_infinite_timeout();
}

bool
scan_udp_single::listener::on_activity(activity)
{
  char data[65536];
  sockaddr_in address;
  socklen_t len = sizeof(address);

  int result = recvfrom(scanner.listen_fd, data, sizeof(data), 0,
                        reinterpret_cast<sockaddr*>(&address),
                        &len);
  if (result < 0) {
    fprintf(stderr, "%s: could receive from UDP socket, error was: %s\n",
            opt_program, strerror(errno));
    exit (EXIT_FAILURE);
  }

  const ipv4_t host = ntohl(address.sin_addr.s_addr);
  unsigned short port = ntohs(address.sin_port);
  senders_t::iterator p = scanner.senders.find(host);
  if (p == scanner.senders.end()) {
    if (opt_verbose && opt_net_errors) {
      ipv4_string_t a;
      ipv4_address_to_string (host, a);

      fprintf(stderr, "%s: warning: stray UDP packet from %s:%u\n",
              opt_program, a, static_cast<unsigned>(port));
    }
  } else {
    p->second->the_handler->reply_received(host, port,
                                           std::string(data, result));
    p->second->stop_sending();
  }
  return true;
}

bool
scan_udp_single::listener::on_timeout(ticks_t)
{
  abort();                      // can't happen
}

// scan_udp_single

scan_udp_single::scan_udp_single(event_queue& q, subnets& n, handler& h,
                                 unsigned maximum, unsigned burst_delay,
                                 unsigned burst_size,
                                 unsigned retry_count, unsigned retry_timeout)
  : queue(q), handler_template(h),
    trigger_handler(*this),
    trigger(q, n, trigger_handler, maximum, burst_delay, burst_size),
    retries(retry_count), timeout(retry_timeout)
{
  listen_fd = socket(PF_INET, SOCK_DGRAM, 0);
  if (listen_fd < 0) {
    fprintf(stderr, "%s: could not create UDP socket, error was: %s\n",
            opt_program, strerror(errno));
    exit (EXIT_FAILURE);
  }

  sockaddr_in address;
  address.sin_family = AF_INET;
  address.sin_port = htons(0);
  address.sin_addr.s_addr = htonl(0);
  int result = bind(listen_fd, reinterpret_cast<sockaddr*>(&address),
                    sizeof(address));
  if (result < 0) {
    fprintf(stderr, "%s: could not listen on UDP socket, error was: %s\n",
            opt_program, strerror(errno));
    exit (EXIT_FAILURE);
  }

  new send_trigger(*this);
  the_listener = new listener(*this);
}
