/*
 *  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 Library 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.
 *
 *       $Id: tcp-fast-scan.c,v 1.2 2011/12/08 20:34:02 james.stevenson Exp $
 *
 * Author:
 *       NAME:   James Stevenson
 *       WWW:    http://www.stev.org
 * Usage:
 *       Compile : gcc -Wall -O2 tcp-fast-scan.c -o tcp-fast-scan
 *       Run             : ./tcp-fast-scan -h
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <time.h>
#include <netdb.h>
#include <unistd.h>
#include <fcntl.h>
#include <sysexits.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>


// Have 1024 open connections at once

#define MAX_CONN 1024  // max conn

#define STATE_NONE 0
#define STATE_WAIT 1
#define STATE_CONNECT 2
#define STATE_TIMEOUT 3
#define STATE_REFUSED 4

struct conn_t {
  int state;
  int time;
  struct sockaddr_in addr;
} conn[MAX_CONN];

int active = 0;    // Number f active connections
int maxconn = 50;  // max number of conns before we loop on select
int maxnew = 30;  // Max number of new conns before select is called
int maxfd = 0;    // Higest fd we have used
int newconn = 0;
int verbose = 0;  // verbos elevel
int hosts = 1;    // Hosts at once
int chosts = 0;    // Number of hosts we have active
int iforked = 0;  // Set to 1 if a child process
fd_set fd_write;  // Use these to spot when we are connected


void print_help(FILE *fp, char *app);
int new_conn(unsigned long ip, unsigned short port);
int kill_conn(int fd);
int wait_conn();
unsigned long host2ip(char *hostname);
int do_hosts(char *str, char *ports);
int do_ports(unsigned long ip, char *strport);
int do_fork();



int main(int argc, char **argv) {
  int i, c;
  char *ports = "1-1024";
  
  // Zero out conns
  for(i=0;i<MAX_CONN;i++) {
    conn[i].state = 0;
    conn[i].time = 0;
  }

  while((c = getopt(argc, argv, "f:hm:n:p:v")) != -1) {
    switch(c) {
      case 'f':
        hosts = atoi(optarg);
        if (optarg <= 0) {
          fprintf(stderr, "Invalid argument to -f\n");
          print_help(stderr, argv[0]);
          exit(EX_USAGE);
        }
        break;
      case 'h':
        print_help(stdout, argv[0]);
        exit(EX_OK);
        break;
      case 'm':
        maxconn = atoi(optarg);
        if (maxconn == 0) {
          fprintf(stderr, "Invalid argument to maxconn\n");
          print_help(stderr, argv[0]);
          exit(EX_USAGE);
        }
        break;
      case 'n':
        maxnew = atoi(optarg);
        if (maxnew == 0) {
          fprintf(stderr, "Invalid argument to maxnew\n");
          print_help(stderr, argv[0]);
          exit(EX_USAGE);
        }
        break;
      case 'p':
        ports = optarg;
        break;
      case 'v':
        verbose++;
        break;
      default:
        break;  
    }
  }


  if (optind < argc) {
    while(optind < argc) {
      do_hosts(argv[optind++], ports);
    }  
  } else {
    fprintf(stderr, "No hosts to scan ?\n");
    print_help(stderr, argv[0]);
    exit(EX_USAGE);
  }  

  
  // Finall loop
  while(active)
    wait_conn();

  while(chosts) {
    // Sleep until all child processes have quit
    retry_wait:  // If a signal goes off
    if (wait(&i) < 0) {
      switch(errno) {
        case EINTR:
          goto retry_wait;
          break;
        default:
          perror("wait");
          exit(-1);
          break;
      }
    }
    // A child exited
    if (WIFEXITED(i)) {
      chosts--;
    } else {
      fprintf(stderr, "Bad exitcode %d from wait()\n", i);
      exit(-1);
    }
  }
  

  return 0;  
}

void print_help(FILE *fp, char *app) {
  fprintf(fp, "%s Usage: [-h] [-m int] [-n int] [-p ports] host\n", app);
  fprintf(fp, "\n");
  fprintf(fp, "-f int   Number of hosts to scan at same time n * f sockets will be open !\n");
  fprintf(fp, "-h       Print this help\n");
  fprintf(fp, "-m int   set the max number of concurrent connections\n");
  fprintf(fp, "-n int   set the max number of connections to start per second\n");
  fprintf(fp, "-p range set the ports like 1,100-200 etc.. default is 1-1024\n");
  fprintf(fp, "-v       verbose level each v adds 1 to level\n");
  fprintf(fp, "         Level 0 - Only print successull connections or problem messages\n");
  fprintf(fp, "         Level 1 - Print timeouts, unreachable host\n");
  fprintf(fp, "         Level 2 - Print Connection refused\n");
  fprintf(fp, "hosts    Each host can be added or a combination of ranges and masks\n");
  fprintf(fp, "         192.168.1.1/24 is the same as 192.168.1.1-192.168.1.255\n");

}


int new_conn(unsigned long ip, unsigned short port) {
  int fd;
  
  while(newconn > maxnew || active > maxconn)
    wait_conn();
  
  fd = socket(PF_INET, SOCK_STREAM, 0);
  if (fd < 0) {
    switch (errno) {
      case EMFILE:
        perror("socket");
        fprintf(stderr, "This is fatel\n");
        exit(EX_SOFTWARE);
        break;
      default:
        perror("socket");
        return -1;
        break;
    }
  }
  
  if (fd >= maxfd)
    maxfd = fd+1;
  
  bzero(&conn[fd], sizeof(struct conn_t));
  conn[fd].state = 0;
  conn[fd].time = time(NULL);

  // Set it up so it does not block
  if (fcntl(fd, F_SETFL, O_NONBLOCK) < 0) {
    perror("fcntl");
    close(fd);
    return -1;
  }


  conn[fd].addr.sin_family = AF_INET;
  conn[fd].addr.sin_port = port;
  conn[fd].addr.sin_addr.s_addr = ip;
  
  if (connect(fd, (struct sockaddr *) &conn[fd].addr, sizeof(conn[fd].addr)) < 0) {
    switch(errno) {
      case EINPROGRESS:
      case EAGAIN:
        break;
      case ENETUNREACH:
        active++;
        return kill_conn(fd);
        break;
      default:
        perror("connect");
        close(fd);
        return -1;
        break;
    }
  }

  FD_SET(fd, &fd_write);
  active++;
  newconn++;  
  conn[fd].state = STATE_WAIT;
  
  return fd;
}

int kill_conn(int fd) {
  char buf[1];
  char *msg = NULL;
  int ret, level = 0;    // Level 0 to print 1 for extended 2 for really extended
  struct servent *service;

  if (conn[fd].state == STATE_TIMEOUT) {
    level = 1;
    msg = "Timeout";
  } else {
    retry:
      ret = recv(fd, &buf[0], 1, 0);
      if (ret < 0) {
        switch(errno) {
          case EINTR:
            goto retry;
            break;
          case EAGAIN:
            level = 0;
            msg = "Port Open";
            break;
          case ECONNRESET:
            level = 0;
            msg = "Open But Reset by peer";
            break;
          case ENETUNREACH:
          case EHOSTUNREACH:
            level = 1;
            msg = "Unreachable host";
            break;
          case ETIMEDOUT:
            level = 1;
            msg = "Timeout";
          case ECONNREFUSED:
            level = 2;
            msg = "Connection Refused";
            break;
          case ENOTCONN:
            level = 2;
            msg = "Unreachable / Refused";
            break;
          default:
            perror("recv");
            break;
        }
      } else {
        msg = "Port Open";
      }
  } // end if timeout
  
  conn[fd].state = 0;
  conn[fd].time = 0;
  FD_CLR(fd, &fd_write);
  if (close(fd) < 0)
    perror("Close");
  active--;

  if (level <= verbose) {
    service = getservbyport(conn[fd].addr.sin_port, "tcp");
    if (service) {
      printf("%s:%d %s (%s)\n", inet_ntoa(conn[fd].addr.sin_addr), ntohs(conn[fd].addr.sin_port), msg, service->s_name);
    } else {
      printf("%s:%d %s\n", inet_ntoa(conn[fd].addr.sin_addr), ntohs(conn[fd].addr.sin_port), msg);
    }
  }
    
  return 0;
}

int wait_conn() {
  struct timeval tv;
  fd_set fdw;
  int ret, i;
  // Do Select Code Here

  newconn = 0;  
  tv.tv_sec = 1;
  tv.tv_usec = 0;
retry:
  if (active == 0)
    return 0;
  
  fdw = fd_write;
  ret = select(maxfd, NULL, &fdw, NULL, &tv);
  if (ret < 0) {
    switch(errno) {
      case EINTR:
        goto retry;
        break;
      default:
        perror("select");
        exit(EX_SOFTWARE);
        break;
    }
  }
  
  for(i=0;i<maxfd;i++) {
    if (FD_ISSET(i, &fdw)) {
      kill_conn(i);    // Just disconnect it now
    }
  }

  if (tv.tv_sec != 0 || tv.tv_usec != 0)
    goto retry;
  
  return 0;
}

unsigned long host2ip(char *hostname) {
  static struct in_addr i;
        struct hostent *h;
        i.s_addr = inet_addr(hostname);
        if(i.s_addr == -1) {
    h = gethostbyname(hostname);
    if(h == NULL)
      return 0;
    bcopy(h->h_addr, (char *)&i.s_addr, h->h_length);
  }
  return i.s_addr;
}


/* Do the host parsing here */
int do_hosts(char *str, char *ports) {
  char *slash, *comma, *dash;
  unsigned long ip, top, bottom;
  int i, mask;
  
  comma = strchr(str, ',');
  if (comma) {
    *comma = 0;
    comma++;
    if (*comma) {
      if (do_hosts(comma, ports) < 0)
        return -1;
    } else {
      /* Sanity Check */
      fprintf(stderr, "Expected Somthing after the comma\n");
      return -1;
    }
  }


  printf("Parsing: %s\n", str);

  slash = strchr(str, '/');
  dash = strchr(str, '-');
  if (slash && dash) {
    fprintf(stderr, "Error Parsing %s\n", str);
    fprintf(stderr, "Dont use a - and / in the same range\n");
    return -1;
  }

  if (!slash && !dash) {
    /* Its just an ip address */
    ip = host2ip(str);
    if (ip == 0) {
      fprintf(stderr, "Failed On: %s\n", str);
      exit(EX_USAGE);
    }
    do_ports(ip, ports);
    return 0;
  }

  /* We have something like 192.168.0.1-192.168.0.254 */
  if (dash) {
    *dash = 0;
    dash++;
    /* Sanity Check */
    if (!*dash) {
      fprintf(stderr, "Excepted An ip after the -\n");
      return -1;
    }
    top = host2ip(str);
    bottom = host2ip(dash);
    if (top == 0 || bottom == 0) {
      fprintf(stderr, "Error Parsing %s-%s\n", str, dash);
      fprintf(stderr, "Must have an ip address range\n");
      return -1;
    }
    /* Swap these back to host format */
    top = ntohl(top);
    bottom = ntohl(bottom);
    /* Switch if wrong way round */
    if (top < bottom) {
      ip = top;
      top = bottom;
      bottom = ip;
    }
    top++; /* Make them inclusive */
    while(bottom < top) {
      do_ports(htonl(bottom), ports);
      bottom++;
    }
    return 0;
  }

  if (!slash)
    return -1;

  /* We should have something like 192.168.0.0/24 */
  *slash = 0;
  slash++;

  /* Sanity Check */
  if (!*slash) {
    fprintf(stderr, "Excected Somthing after the /\n");
    return -1;
  }

  
  ip = host2ip(str);
  if (ip == 0) {
    fprintf(stderr, "Failed on %s\n", str);
    return -1;
  }

  /* Get the number of on bits */
  mask = atoi(slash);
  if (mask <= 0 || mask > 32) {
    fprintf(stderr, "Parse Error after /\n");
    fprintf(stderr, "Must be a number Between 1-32\n");
    return -1;
  }

  /* Get the bottom */
  bottom = ntohl(ip);
  top = bottom;

  /* Set the correct bits */
  i = mask;
  if (i == 32) i = 31;
  /* FIXME: possible endian problem if arch != x86 */
  i -= 31; /* Push to a negitive */
  i *= -1; /* Put it back to + */
  while(i) {
    bottom &= ~(1 << i);
    i--;
  }
  i = mask;
  if (i == 32) i = 31;
  /* FIXME: possible endian problem if arch != x86 */
  i -= 31; /* Push to a negitive */
  i *= -1; /* Put it back to + */
  while(i) {
    top |= (1 << i);
    i--;
  }
  top++; /* Make it inclusive */
  while(bottom < top) {
    do_ports(htonl(bottom), ports);
    bottom++;
  }

  return 0;
}

int do_ports(unsigned long ip, char *strport) {
  char *str;
  char *sep; /* seperator either : or - */
  int lport, hport;
  
  if (do_fork() == 0)
    return 0;

  // If we forked we cannot ever return to the caller from this function !
  // use RETURN instead !
  #define RETURN(x) if (iforked) exit(x); else return (x);
  
  str = strdup(strport);
  if (!str) {
    perror("strdup");
    RETURN(-1);
  }

  sep = strchr(str, ':');
  if (!sep)
    sep = strchr(str, '-');

  if (!sep) { /* its a single port */
    lport = atoi(str);
    if (lport > 65535 || lport < 0) {
      fprintf(stderr, "Port outside valid range %s\n", strport);
      exit(EX_USAGE);
    }
    new_conn(ip, htons(lport));
    free(str);
    RETURN(0);
  }

  /* it has 2 values */
  *sep = 0;
  sep++;
  if (!*sep) { /* Sanity Check */
    free(str);
    RETURN(-1);
  }

  lport = atoi(str);
  hport = atoi(sep);
  if (lport == 0 || hport == 0 || hport > 65535 || lport > 65535) {
    fprintf(stderr, "Ports outside a valid range %s\n", strport);
    exit(EX_USAGE);
  }
  while(lport <= hport) {
    new_conn(ip, htons(lport));
    lport++;
  }
  free(str);
  RETURN(0);
  return 0; // Prevents compiler warning
}

// This return 0 if it is the parent and pid if it is the child
// it will exit the program and print and error on failure
// it will also block if the chosts limit is reached
int do_fork() {
  int ret;
  // This is where we fork for mutil hosts at a time
  if (hosts <= 1)
    return 1;

  if (iforked)
    return 1;
  
  if (chosts >= hosts) {  // Have we hit the limit ?
    retry_wait:  // If a signal goes off
    ret = 0;
    if (wait(&ret) < 0) {
      switch(errno) {
        case EINTR:
          goto retry_wait;
          break;
        default:
          perror("wait");
          exit(-1);
          break;
      }
    }
    // A child exited
    if (WIFEXITED(ret)) {
      chosts--;
    } else {
      fprintf(stderr, "Bad exitcode %d from wait()\n", ret);
      exit(-1);
    }
  }
  // Just before we fork we should flush all open file streams
  fflush(NULL);  
  // This is where we fork another child
  ret = fork();
  if (ret > 0) { // We are the parent return and let the scan run in child
    chosts++;
    return 0;
  }
  if (ret < 0) {
    perror("fork");
    exit(-1);
  }
  // ret == 0
  iforked = 1;
  // Lety the child run as if nothing happened
  return ret;
}