/*
 * Trevor's utility to find an empeg on a network using a broadcast.
 *
 * Copyright (C) 2003 Trevor Man
 * 
 * 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 <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <unistd.h>
#include <string.h>
#include <sys/time.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <netdb.h>
#include <sys/socket.h>

#define DISCOVERY_DATA		"?"
#define DISCOVERY_PORT		8300
#define MAX_PACKET_SIZE		65500
#define MAX_NAME_LENGTH		255
#define MAX_STRING_LENGTH	8196

struct empeg_data
{
  char name[MAX_NAME_LENGTH];
  int serial_number;
  unsigned short addr[4];
};

int edit_hosts_file(char *filename, struct empeg_data *empeg, char *name);
int get_listen_socket(int port_number);
int send_msg(int sock, char *addr, int port_number, char *msg, size_t msg_len);
int get_msg(int sock, char *buf, size_t buf_size, int timeout, struct sockaddr_in *addr);
int decode_empeg_reply(char *buf, struct empeg_data *empeg, struct sockaddr_in *addr);

int main(int argc, char *argv[])
{
  char *chosen_empeg_name = NULL;
  char *host_file_name = NULL;
  char *broadcast_addr = "255.255.255.255"; // Default broadcast address
  int required_empeg_serial = 0;
  int identify_all = 0;
  struct empeg_data empeg;
  struct sockaddr_in tmpaddr;
  char buf[MAX_PACKET_SIZE];
  int tmp, sock, timeout = 5;

  fprintf(stderr, "Trevor's Empeg/RioCar Finder v1.3 - Copyright 2002 - Trevor Man\n");

  for (tmp = 1; tmp < argc; tmp++) {
    if ((argv[tmp][0] == '-') && (strlen(argv[tmp]) >= 2)) {
      switch(argv[tmp][1]) {
      case 'b':
      case 'B':
        broadcast_addr = argv[tmp] + 2;
        if (strlen(broadcast_addr) == 0) {
          fprintf(stderr, "Invalid broadcast address.\n");
          return 1;
        }
        fprintf(stderr, "Broadcasting to: %s\n", broadcast_addr);
        break;

      case 's':
      case 'S':
        required_empeg_serial = atoi(argv[tmp] + 2);
        if (required_empeg_serial == 0) {
          fprintf(stderr, "Invalid serial number.\n");
          return 1;
        }
        fprintf(stderr, "Searching for serial number %d\n", required_empeg_serial);
        break;

      case 'h':
      case 'H':
        host_file_name = argv[tmp] + 2;
        if (strlen(host_file_name) == 0) {
          fprintf(stderr, "Invalid hosts filename.\n");
          return 1;
        }
        fprintf(stderr, "Host file: %s\n", host_file_name);
        break;

      case 'n':
      case 'N':
        chosen_empeg_name = argv[tmp] + 2;
        if (strlen(chosen_empeg_name) == 0) {
          fprintf(stderr, "Invalid empeg name.\n");
          return 1;
        }
        fprintf(stderr, "Name to give to empeg: %s\n", chosen_empeg_name);
        break;

      case 't':
      case 'T':
        timeout = atoi(argv[tmp] + 2);
        if (timeout == 0) {
          fprintf(stderr, "Invalid timeout.\n");
          return 1;
        }
        fprintf(stderr, "Setting timeout to be %d second%s.\n", timeout, timeout == 1 ? "" : "s");
        break;

      case 'i':
      case 'I':
        identify_all = 1;
        fprintf(stderr, "Will now identify all empegs on network.\n");
        break;

      default:
        fprintf(stderr, "Unknown option: %s\n\n", argv[tmp]);
      case '?':
        fprintf(stderr, "Usage: %s <options>\n", argv[0]);
        fprintf(stderr, "\n"
                        "-b<addr>   - Set broadcast address if needed e.g. 192.168.1.255\n"
                        "-s<serial> - Set serial number to look for e.g. 12345678\n"
                        "-n<name>   - Name to give empeg in hosts file e.g. empeg.mynetwork.com\n"
                        "-h<file>   - Hosts file to edit\n"
                        "-t<number> - Number of seconds before timing out\n"
                        "-i         - Identify all empegs on network\n");
        return 1;
      }
    }
  }

  if (host_file_name == NULL && chosen_empeg_name) {
    fprintf(stderr, "You must specify a hosts file.\n");
    return 1;
  }


  if ((sock = get_listen_socket(DISCOVERY_PORT)) == -1) {
    fprintf(stderr, "Unable to open listen socket.\n");
    return 1;
  }

  if (!send_msg(sock, broadcast_addr, DISCOVERY_PORT, DISCOVERY_DATA, sizeof(DISCOVERY_DATA))) {
    fprintf(stderr, "Unable to send broadcast.\n");

    close(sock);

    return 1;
  }

  for (;;) {
    if (!get_msg(sock, buf, MAX_PACKET_SIZE, timeout, &tmpaddr)) {
      printf("No more replies.  Giving up...\n");
      break;
    }

    if (!decode_empeg_reply(buf, &empeg, &tmpaddr)) // Is this a valid reply?
      continue;

    fprintf(stderr, "Found an empeg with serial %0d called \"%s\" at %d.%d.%d.%d\n",
            empeg.serial_number, empeg.name,
            empeg.addr[0], empeg.addr[1], empeg.addr[2], empeg.addr[3]);

    if ((required_empeg_serial == 0) ||
        (required_empeg_serial == empeg.serial_number) ||
        (identify_all == 1)) {
      if (host_file_name) {
        if (!edit_hosts_file(host_file_name, &empeg, chosen_empeg_name)) {
          fprintf(stderr, "Unable to modify %s\n", host_file_name);
          close(sock);
          return 1;
        }
      }

      if (!identify_all)
        break;
    }
  }

  close(sock);

  return 0;
}

int get_listen_socket(int port_number)
{
  struct sockaddr_in listen_addr;
  int sock;
  int sockopt;

  assert(port_number != 0);

  // Open a socket
  if ((sock = socket(AF_INET, SOCK_DGRAM, PF_UNSPEC)) == -1) {
    perror("socket");
    return -1;
  }

  memset(&listen_addr, 0, sizeof(listen_addr));

  listen_addr.sin_family = AF_INET;
  listen_addr.sin_port = htons(port_number);
  listen_addr.sin_addr.s_addr = INADDR_ANY;

  // Bind to socket
  if (bind(sock, (struct sockaddr *)&listen_addr, sizeof(listen_addr)) == -1) {
    perror("bind");
    return -1;
  }

  sockopt = 1;
  if (setsockopt(sock, SOL_SOCKET, SO_BROADCAST, &sockopt, sizeof(sockopt))) {
    perror("setsockopt");
    return -1;
  }

  return sock;
}

int send_msg(int sock, char *addr, int port_number, char *msg, size_t msg_len)
{
  struct hostent *host_entry;
  struct sockaddr_in send_addr;
  int len;

  assert(addr != NULL);
  assert(port_number != 0);
  assert(msg != NULL);
  assert(msg_len != 0);

  memset(&send_addr, 0, sizeof(send_addr));
  len = sizeof(struct sockaddr);
  send_addr.sin_family = AF_INET;
  send_addr.sin_port = htons(DISCOVERY_PORT);

  if ((host_entry = gethostbyname(addr)) == NULL) {
    perror("gethostbyname");
    return 0;
  }
  send_addr.sin_addr = *((struct in_addr *)host_entry->h_addr);

  if (sendto(sock, msg, msg_len, 0, (struct sockaddr *)&send_addr, len) == -1) {
    perror("sendto");
    return 0;
  }

  return 1;
}

int get_msg(int sock, char *buf, size_t buf_size, int timeout, struct sockaddr_in *addr)
{
  fd_set read_fds;
  struct timeval tv;
  int len;

  assert(buf != NULL);
  assert(buf_size != 0);
  assert(timeout != 0);
  assert(addr != NULL);

  len = sizeof(struct sockaddr);

  FD_ZERO(&read_fds);
  FD_SET(sock, &read_fds);

  memset(buf, 0, buf_size);

  tv.tv_sec = timeout;
  tv.tv_usec = 0;
  select(sock + 1, &read_fds, NULL, NULL, &tv);

  if (!FD_ISSET(sock, &read_fds))
    return 0;

  if (recvfrom(sock, buf, buf_size, 0, (struct sockaddr *)addr, &len) == -1) {
    return 0;
  } else {
    return 1;
  }
}

int decode_empeg_reply(char *buf, struct empeg_data *empeg, struct sockaddr_in *addr)
{
  char *empeg_serial;
  char *empeg_name;

  assert(buf != NULL);
  assert(empeg != NULL);
  assert(addr != NULL);

  empeg_serial = strstr(buf, "id=");
  empeg_name = strstr(buf, "name=");

  if ((empeg_serial == NULL) || (empeg_name == NULL))
    return 0;

  empeg->serial_number = atoi(empeg_serial + 3);
  strcpy(empeg->name, empeg_name + 5);
  empeg->addr[0] = (addr->sin_addr.s_addr & 0x000000FF) >> 0;
  empeg->addr[1] = (addr->sin_addr.s_addr & 0x0000FF00) >> 8;
  empeg->addr[2] = (addr->sin_addr.s_addr & 0x00FF0000) >> 16;
  empeg->addr[3] = (addr->sin_addr.s_addr & 0xFF000000) >> 24;

  return 1;
}

int edit_hosts_file(char *filename, struct empeg_data *empeg, char *name)
{
  char empeg_name[MAX_STRING_LENGTH];
  char buf[MAX_STRING_LENGTH];
  char tmp[MAX_STRING_LENGTH];
  char tmp1[MAX_STRING_LENGTH];
  FILE *fp;
  int added_entry = 0;

  assert(filename != NULL);
  assert(empeg != NULL);

  if ((fp = fopen(filename, "r+t")) == NULL) {
    perror("fopen");
    return 0;
  }

  memset(tmp, 'X', MAX_NAME_LENGTH - 1);
  tmp[MAX_NAME_LENGTH] = '\0';

  sprintf(empeg_name, "EMPEG-%d-EMPEG", empeg->serial_number);

  sprintf(tmp1, "%d.%d.%d.%d %s empeg-%d # Empeg: Serial=%d Name=%s %s",
          empeg->addr[0], empeg->addr[1], empeg->addr[2], empeg->addr[3],
          name ? name : "", empeg->serial_number, empeg->serial_number, empeg->name, empeg_name);
  strncpy(tmp, tmp1, strlen(tmp1));

  while (!feof(fp)) {
    fgets(buf, MAX_STRING_LENGTH - 1, fp);
    if (feof(fp))
      break;
    if (strstr(buf, empeg_name) != NULL) {
#if defined(LINUX) // LF
      fseek(fp, -MAX_NAME_LENGTH, SEEK_CUR);
#else // CRLF
      fseek(fp, -(MAX_NAME_LENGTH + 1), SEEK_CUR);
#endif
      if (name)
        fprintf(stderr, "  Modifying entry for %s in %s\n", name, filename);
      else
        fprintf(stderr, "  Modifying default entry for %d in %s\n", empeg->serial_number, filename);
      fprintf(fp, "%s\n", tmp);
      added_entry = 1;
    }
  }

  if (!added_entry) {
    if (name)
      fprintf(stderr, "  Appending entry for %s to %s\n", name, filename);
    else
      fprintf(stderr, "  Appending default entry for %d to %s\n", empeg->serial_number, filename);
    fprintf(fp, "\n%s\n", tmp);
  }

  fclose(fp);

  return 1;
}