/*
 * PJ Waskiewicz
 * 4/23/2000 port_scanner.c
 *
 * A port scanner, scans for open
 * TCP and UDP ports.
 *
 */

#include <netinet/in.h>
#include <netdb.h>
#include <stdlib.h>
#include <stdio.h>
#include <fcntl.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>
#include <errno.h>

#define EXIT_GOOD 0
#define EXIT_BAD 1
#define BUF_MAX 1024

void usage();

main(int argc, char *argv[]) {
  int sock_fd;
  struct timeval tv;
  int retval;
  struct sockaddr_in *p_sock;
  struct hostent *p_host;
  char *hostname;
  char *type = NULL;
  char buf[BUF_MAX];
  u_short b_port, b_port_temp;
  u_short f_port;
  char opt;
  fd_set rfds;
  int size;
  opterr = 0;
  while(1) {
    if((opt = getopt(argc, argv, "h:b:f:t:")) == '?') usage();
    else if(opt == -1) {
      if((optind == 1) || (optind < 9)) usage();
      else break;
    }
    if(opt == 'h') {
      /* Grab the hostname */
      hostname = malloc(strlen(optarg));
      memcpy(hostname, optarg, strlen(optarg));
    }
    else if(opt == 'b') {
      /* Grab the starting port */
      b_port = (u_short)atoi(optarg);
      b_port_temp = b_port;
    }
    else if(opt == 'f') {
      /* Grab the ending port */
      f_port = (u_short)atoi(optarg);
    }
    else if(opt == 't') {
      /* Make sure we have a valid type */
      if((strcmp(optarg, "tcp") == 0) || (strcmp(optarg, "udp") == 0)
         || (strcmp(optarg, "both") == 0)) {
        /* We're good to go */
        type = malloc(strlen(optarg));
        memcpy(type, optarg, strlen(optarg));
      }
      else usage();
    }
  }
  if(b_port > f_port) {
    free(hostname);
    free(type);
    usage();
  }
  /* Now check which service to scan */
  if(strcmp(type, "both") == 0) {
    /* Do both of them */
    /* First, tcp, since it is more work */
    do {
      p_sock->sin_family = AF_INET;
      p_sock->sin_port = htons(b_port);
      p_host = gethostbyname(hostname);
      p_sock->sin_addr = *(struct in_addr *)p_host->h_addr;
      tv.tv_sec = 1;
      tv.tv_usec = 0;
      sock_fd = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
      FD_ZERO(&rfds);
      FD_SET(sock_fd, &rfds);
      memset(buf, '\0', BUF_MAX);
      fcntl(sock_fd, F_SETFL, O_NONBLOCK);
      connect(sock_fd, p_sock, sizeof(struct sockaddr_in));
      retval = select(sock_fd + 1, &rfds, NULL, NULL, &tv);
      size = recv(sock_fd, buf, BUF_MAX, MSG_NOSIGNAL);
      if((errno != ETIMEDOUT) && (errno != ECONNREFUSED)
         && (errno != EHOSTDOWN) && (errno != EHOSTUNREACH)) {
        /* Port is there */
        if(strlen(buf) > 0) {
          if((buf[strlen(buf) - 2] == '\n') ||
             (buf[strlen(buf) - 1] == '\n'))
            printf("%i\tTCP\t%s", b_port, buf);
          else printf("%i\tTCP\t%s\n", b_port, buf);
        }
        else printf("%i\tTCP\t\n", b_port);
      }
      close(sock_fd);
      /* Now go for udp */
      p_sock->sin_family = AF_INET;
      p_sock->sin_port = htons(b_port);
      p_host = gethostbyname(hostname);
      p_sock->sin_addr = *(struct in_addr *)p_host->h_addr;
      tv.tv_sec = 1;
      tv.tv_usec = 0;
      sock_fd = socket(PF_INET, SOCK_DGRAM, IPPROTO_IP);
      FD_ZERO(&rfds);
      FD_SET(sock_fd, &rfds);
      memset(buf, '\0', BUF_MAX);
      fcntl(sock_fd, F_SETFL, O_NONBLOCK);
      //      connect(sock_fd, p_sock, sizeof(struct sockaddr_in));
      //      send(sock_fd, "\n", strlen("\n"), MSG_NOSIGNAL);
      sendto(sock_fd, "\n", strlen("\n"), MSG_NOSIGNAL,
             p_sock, sizeof(struct sockaddr_in));
      //      retval = select(sock_fd + 1, &rfds, NULL, NULL, &tv);
      //      size = recvfrom(sock_fd, buf, BUF_MAX, MSG_NOSIGNAL,
      //                      p_sock, (int *)sizeof(struct sockaddr_in));
      //      size = recv(sock_fd, buf, BUF_MAX, MSG_NOSIGNAL);
      //perror("**UDP**");
      if((errno != ETIMEDOUT) && (errno != ECONNREFUSED)
         && (errno != EHOSTDOWN) && (errno != EHOSTUNREACH)) {
        /* Port is there */
        if(strlen(buf) > 0) {
          if((buf[strlen(buf) - 2] == '\n') ||
             (buf[strlen(buf) - 1] == '\n'))
            printf("%i\tUDP\t%s", b_port, buf);
          else printf("%i\tUDP\t%s\n", b_port, buf);
        }
        else printf("%i\tUDP\t\n", b_port);
      }
      close(sock_fd);
      b_port++;
    } while(b_port <= f_port);
  }
  else if(strcmp(type, "tcp") == 0) {
    do {
      p_sock->sin_family = AF_INET;
      p_sock->sin_port = htons(b_port);
      p_host = gethostbyname(hostname);
      p_sock->sin_addr = *(struct in_addr *)p_host->h_addr;
      tv.tv_sec = 1;
      tv.tv_usec = 0;
      sock_fd = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
      FD_ZERO(&rfds);
      FD_SET(sock_fd, &rfds);
      memset(buf, '\0', BUF_MAX);
      fcntl(sock_fd, F_SETFL, O_NONBLOCK);
      connect(sock_fd, p_sock, sizeof(struct sockaddr_in));
      retval = select(sock_fd + 1, &rfds, NULL, NULL, &tv);
      size = recv(sock_fd, buf, BUF_MAX, MSG_NOSIGNAL);
      if((errno != ETIMEDOUT) && (errno != ECONNREFUSED)
         && (errno != EHOSTDOWN) && (errno != EHOSTUNREACH)) {
        /* Port is there */
        if(strlen(buf) > 0) {
          if((buf[strlen(buf) - 2] == '\n') ||
             (buf[strlen(buf) - 1] == '\n'))
            printf("%i\tTCP\t%s", b_port, buf);
          else printf("%i\tTCP\t%s\n", b_port, buf);
        }
        else printf("%i\tTCP\t\n", b_port);
      }
      close(sock_fd);
      b_port++;
    } while(b_port <= f_port);
  }
  else if(strcmp(type, "udp") == 0) {
    fprintf(stderr, "UDP not implemented yet.\n");
    exit(EXIT_GOOD);
  }
}

void usage() {
  printf("usage: port_scan -h <hostname> -b <begin-port>");
  printf(" -f <finish-port> [-t <type>]\n");
  printf("\t-t <type>: tcp for TCP, udp for UDP, or both for BOTH\n");
  exit(EXIT_BAD);
}