/*
 * ZMap Copyright 2013 Regents of the University of Michigan 
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 */

#define _GNU_SOURCE

#include <stdlib.h>
#include <stdio.h>

#include <assert.h>
#include <sched.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <net/if.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <pcap/pcap.h>
#include <pthread.h>

#include "../lib/logger.h"
#include "../lib/random.h"

#include "zopt.h"
#include "send.h"
#include "recv.h"
#include "state.h"
#include "monitor.h"
#include "output_modules/output_modules.h"
#include "probe_modules/probe_modules.h"
#include "get_gateway.h"

pthread_mutex_t cpu_affinity_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t recv_ready_mutex = PTHREAD_MUTEX_INITIALIZER;

static void set_cpu(void)
{
	pthread_mutex_lock(&cpu_affinity_mutex);
	static int core=0;
	int num_cores = sysconf(_SC_NPROCESSORS_ONLN);
	cpu_set_t cpuset;	
	CPU_ZERO(&cpuset);
	CPU_SET(core, &cpuset);
	if (pthread_setaffinity_np(pthread_self(),
				sizeof(cpu_set_t), &cpuset) != 0) {
		log_error("zmap", "can't set thread CPU affinity");
	}
	log_trace("zmap", "set thread %u affinity to core %d",
			pthread_self(), core);
	core = (core + 1) % num_cores;
	pthread_mutex_unlock(&cpu_affinity_mutex);
}	

static void* start_send(__attribute__((unused)) void *arg)
{
	set_cpu();
	send_run();
	return NULL;
}

static void* start_recv(__attribute__((unused)) void *arg)
{
	set_cpu();
	recv_run(&recv_ready_mutex);
	return NULL;
}

static void *start_mon(__attribute__((unused)) void *arg)
{
	set_cpu();
	monitor_run();
	return NULL;
}

#define SI(w,x,y) printf("%s\t%s\t%i\n", w, x, y);
#define SD(w,x,y) printf("%s\t%s\t%f\n", w, x, y);
#define SU(w,x,y) printf("%s\t%s\t%u\n", w, x, y);
#define SLU(w,x,y) printf("%s\t%s\t%lu\n", w, x, y);
#define SS(w,x,y) printf("%s\t%s\t%s\n", w, x, y);
#define STRTIME_LEN 1024

static void summary(void)
{
	char send_start_time[STRTIME_LEN+1];
	assert(dstrftime(send_start_time, STRTIME_LEN, "%c", zsend.start));
	char send_end_time[STRTIME_LEN+1];
	assert(dstrftime(send_end_time, STRTIME_LEN, "%c", zsend.finish));
	char recv_start_time[STRTIME_LEN+1];
	assert(dstrftime(recv_start_time, STRTIME_LEN, "%c", zrecv.start));
	char recv_end_time[STRTIME_LEN+1];
	assert(dstrftime(recv_end_time, STRTIME_LEN, "%c", zrecv.finish));
	double hitrate = ((double) 100 * zrecv.success_unique)/((double)zsend.sent);

	SU("cnf", "target-port", zconf.target_port);
	SU("cnf", "source-port-range-begin", zconf.source_port_first);
	SU("cnf", "source-port-range-end", zconf.source_port_last);
	SS("cnf", "source-addr-range-begin", zconf.source_ip_first);
	SS("cnf", "source-addr-range-end", zconf.source_ip_last);
	SU("cnf", "maximum-targets", zconf.max_targets);
	SU("cnf", "maximum-runtime", zconf.max_runtime);
	SU("cnf", "maximum-results", zconf.max_results);
	SU("cnf", "permutation-seed", zconf.seed);
	SI("cnf", "cooldown-period", zconf.cooldown_secs);
	SS("cnf", "send-interface", zconf.iface);
	SI("cnf", "rate", zconf.rate);
	SLU("cnf", "bandwidth", zconf.bandwidth);
	SU("env", "nprocessors", (unsigned) sysconf(_SC_NPROCESSORS_ONLN));
	SS("exc", "send-start-time", send_start_time);
	SS("exc", "send-end-time", send_end_time);
	SS("exc", "recv-start-time", recv_start_time);
	SS("exc", "recv-end-time", recv_end_time);
	SU("exc", "sent", zsend.sent);
	SU("exc", "blacklisted", zsend.blacklisted);
	SU("exc", "first-scanned", zsend.first_scanned);
	SD("exc", "hit-rate", hitrate);
	SU("exc", "success-total", zrecv.success_total);
	SU("exc", "success-unique", zrecv.success_unique);
	SU("exc", "success-cooldown-total", zrecv.cooldown_total);
	SU("exc", "success-cooldown-unique", zrecv.cooldown_unique);
	SU("exc", "failure-total", zrecv.failure_total);
	SU("exc", "sendto-failures", zsend.sendto_failures);
	SU("adv", "permutation-gen", zconf.generator);
	SS("exc", "scan-type", zconf.probe_module->name);
}

static void start_zmap(void)
{
	log_init(stderr, zconf.log_level);
	log_info("zmap", "started");

	// finish setting up configuration
	if (zconf.iface == NULL) {
		char errbuf[PCAP_ERRBUF_SIZE];
		char *iface = pcap_lookupdev(errbuf);
		if (iface == NULL) {
			log_fatal("zmap", "could not detect default network interface "
					"(e.g. eth0). Try running as root or setting"
					" interface using -i flag.");
		}
		log_debug("zmap", "no interface provided. will use %s", iface);
		zconf.iface = iface;
	}
	if (zconf.source_ip_first == NULL) {
		struct in_addr default_ip;
		zconf.source_ip_first = malloc(INET_ADDRSTRLEN);
		zconf.source_ip_last = zconf.source_ip_first;
		if (get_iface_ip(zconf.iface, &default_ip) < 0) {
			log_fatal("zmap", "could not detect default IP address for for %s."
					" Try specifying a source address (-S).", zconf.iface);
		}
		inet_ntop(AF_INET, &default_ip, zconf.source_ip_first, INET_ADDRSTRLEN);
		log_debug("zmap", "no source IP address given. will use %s",
				zconf.source_ip_first);
	}
	if (!zconf.gw_mac_set) {
		struct in_addr gw_ip;
		char iface[IF_NAMESIZE];
		if (get_default_gw(&gw_ip, iface) < 0) {
			log_fatal("zmap", "could not detect default gateway address for %i."
					" Try setting default gateway mac address (-G).");
		}
		log_debug("zmap", "found gateway IP %s on %s", inet_ntoa(gw_ip), iface); 
		if (get_hw_addr(&gw_ip, iface, zconf.gw_mac) < 0) {
			log_fatal("zmap", "could not detect GW MAC address for %s on %s."
					" Try setting default gateway mac address (-G).",
					inet_ntoa(gw_ip), zconf.iface);
		}
		zconf.gw_mac_set = 1;
		log_debug("zmap", "using default gateway MAC %02x:%02x:%02x:%02x:%02x:%02x",
				  zconf.gw_mac[0], zconf.gw_mac[1], zconf.gw_mac[2],
				  zconf.gw_mac[3], zconf.gw_mac[4], zconf.gw_mac[5]);
	}

	// initialization
	if (zconf.output_module && zconf.output_module->init) {
		zconf.output_module->init(&zconf);
	}
	if (send_init()) {
		exit(EXIT_FAILURE);
	}
	if (zconf.output_module && zconf.output_module->start) {
		zconf.output_module->start(&zconf, &zsend, &zrecv);
	}

	// start threads
	pthread_t *tsend, trecv, tmon;
	int r = pthread_create(&trecv, NULL, start_recv, NULL);
	if (r != 0) {
		log_fatal("zmap", "unable to create recv thread");
		exit(EXIT_FAILURE);
	}
	for (;;) {
		pthread_mutex_lock(&recv_ready_mutex);
		if (zconf.recv_ready) {
			break;
		}
		pthread_mutex_unlock(&recv_ready_mutex);
	}
	tsend = malloc(zconf.senders * sizeof(pthread_t));
	assert(tsend);
	log_debug("zmap", "using %d sender threads", zconf.senders);
	for (int i=0; i < zconf.senders; i++) {
		r = pthread_create(&tsend[i], NULL, start_send, NULL);
		if (r != 0) {
			log_fatal("zmap", "unable to create send thread");
			exit(EXIT_FAILURE);
		}
	}
	if (!zconf.quiet) {
		r = pthread_create(&tmon, NULL, start_mon, NULL);
		if (r != 0) {
			log_fatal("zmap", "unable to create monitor thread");
			exit(EXIT_FAILURE);
		}
	}

	// wait for completion
	for (int i=0; i < zconf.senders; i++) {
		pthread_join(tsend[i], NULL);
		if (r != 0) {
			log_fatal("zmap", "unable to join send thread");
			exit(EXIT_FAILURE);
		}
	}
	log_debug("zmap", "senders finished");
	pthread_join(trecv, NULL);
	if (r != 0) {
		log_fatal("zmap", "unable to join recv thread");
		exit(EXIT_FAILURE);
	}
	if (!zconf.quiet) {
		pthread_join(tmon, NULL);
		if (r != 0) {
			log_fatal("zmap", "unable to join monitor thread");
			exit(EXIT_FAILURE);
		}
	}

	// finished
	if (zconf.summary) {
		summary();
	}
	if (zconf.output_module && zconf.output_module->close) {
		zconf.output_module->close(&zconf, &zsend, &zrecv);
	}
	if (zconf.probe_module && zconf.probe_module->close) {
		zconf.probe_module->close(&zconf, &zsend, &zrecv);
	}
	log_info("zmap", "completed");
}

static void enforce_range(const char *name, int v, int min, int max)
{
	if (v < min || v > max) {
	  	fprintf(stderr, "%s: argument `%s' must be between %d and %d\n",
			CMDLINE_PARSER_PACKAGE, name, min, max);
		exit(EXIT_FAILURE);
	}
}

static int file_exists(char *name)
{
	FILE *file = fopen(name, "r");
	if (!file)
		return 0;
	fclose(file);
	return 1;
}

#define MAC_LEN IFHWADDRLEN
int parse_mac(macaddr_t *out, char *in)
{
	if (strlen(in) < MAC_LEN*3-1)
		return 0;
	char octet[3];
	octet[2] = '\0';
	for (int i=0; i < MAC_LEN; i++) {
		if (i < MAC_LEN-1 && in[i*3+2] != ':') {
			return 0;
		}
		strncpy(octet, &in[i*3], 2);
		char *err = NULL;
		long b = strtol(octet, &err, 16);
		if (err && *err != '\0') {
			return 0;
		}
		out[i] = b & 0xFF;
	}
	return 1;
}

#define SET_IF_GIVEN(DST,ARG) \
	{ if (args.ARG##_given) { (DST) = args.ARG##_arg; }; }
#define SET_BOOL(DST,ARG) \
	{ if (args.ARG##_given) { (DST) = 1; }; }

int main(int argc, char *argv[])
{
	struct gengetopt_args_info args;
	struct cmdline_parser_params *params;
	params = cmdline_parser_params_create();
	params->initialize = 1;
	params->override = 0;
	params->check_required = 0;
	if (cmdline_parser_ext(argc, argv, &args, params) != 0) {
		exit(EXIT_SUCCESS);
	}
	if (args.help_given) {
		cmdline_parser_print_help();
		exit(EXIT_SUCCESS);
	}
	if (args.version_given) {
		cmdline_parser_print_version();
		exit(EXIT_SUCCESS);
	}
	if (args.list_output_modules_given) {
		print_output_modules();
		exit(EXIT_SUCCESS);
	}
	if (args.list_probe_modules_given) {
		print_probe_modules();
		exit(EXIT_SUCCESS);
	}
	if (args.config_given || file_exists(args.config_arg)) {
		params->initialize = 0;
		params->override = 0;
		if (cmdline_parser_config_file(args.config_arg, &args, params) 
				!= 0) {
			exit(EXIT_FAILURE);
		}
	}
	if (cmdline_parser_required(&args, CMDLINE_PARSER_PACKAGE) != 0) {
		exit(EXIT_FAILURE);
	}

	SET_IF_GIVEN(zconf.output_filename, output_file);
	SET_IF_GIVEN(zconf.blacklist_filename, blacklist_file);
	SET_IF_GIVEN(zconf.whitelist_filename, whitelist_file);
	SET_IF_GIVEN(zconf.probe_args, probe_args);
	SET_IF_GIVEN(zconf.output_args, output_args);
	SET_IF_GIVEN(zconf.iface, interface);
	SET_IF_GIVEN(zconf.max_runtime, max_runtime);
	SET_IF_GIVEN(zconf.max_results, max_results);
	SET_IF_GIVEN(zconf.rate, rate);
	SET_IF_GIVEN(zconf.packet_streams, probes);
	SET_BOOL(zconf.dryrun, dryrun);
	SET_BOOL(zconf.quiet, quiet);
	SET_BOOL(zconf.summary, summary);
	zconf.cooldown_secs = args.cooldown_time_arg;
	zconf.senders = args.sender_threads_arg;
	zconf.log_level = args.verbosity_arg;

	zconf.output_module = get_output_module_by_name(args.output_module_arg);
	if (!zconf.output_module) {
	  fprintf(stderr, "%s: specified output module (%s) does not exist\n",
		  args.output_module_arg, CMDLINE_PARSER_PACKAGE);
	  exit(EXIT_FAILURE);
	}
	zconf.probe_module = get_probe_module_by_name(args.probe_module_arg);
	if (!zconf.probe_module) {
		fprintf(stderr, "%s: specified probe module (%s) does not exist\n",
				CMDLINE_PARSER_PACKAGE, args.probe_module_arg);
	  exit(EXIT_FAILURE);
	}
	if (zconf.probe_module->port_args) {
		if (args.source_port_given) {
			char *dash = strchr(args.source_port_arg, '-');
			if (dash) { // range
				*dash = '\0';
				zconf.source_port_first = atoi(args.source_port_arg);
				enforce_range("starting source-port", zconf.source_port_first, 0, 0xFFFF);
				zconf.source_port_last = atoi(dash+1);
				enforce_range("ending source-port", zconf.source_port_last, 0, 0xFFFF);
				if (zconf.source_port_first > zconf.source_port_last) {
					fprintf(stderr, "%s: invalid source port range: "
						"last port is less than first port\n",
						CMDLINE_PARSER_PACKAGE);
					exit(EXIT_FAILURE);
				}
			} else { // single port
				int port = atoi(args.source_port_arg);
				enforce_range("source-port", port, 0, 0xFFFF);
				zconf.source_port_first = port;
				zconf.source_port_last = port;
			}
		}
	  	if (!args.target_port_given) {
			fprintf(stderr, "%s: target port is required for this type of probe\n",
				CMDLINE_PARSER_PACKAGE);
			exit(EXIT_FAILURE);
		}
		enforce_range("target-port", args.target_port_arg, 0, 0xFFFF);
		zconf.target_port = args.target_port_arg;
	}
	if (args.source_ip_given) {
		char *dash = strchr(args.source_ip_arg, '-');
		if (dash) { // range
			*dash = '\0';
			zconf.source_ip_first = args.source_ip_arg;
			zconf.source_ip_last = dash+1;
		} else { // single address
			zconf.source_ip_first = args.source_ip_arg;
			zconf.source_ip_last = args.source_ip_arg;
		}
	}
	if (args.gateway_mac_given) {
		if (!parse_mac(zconf.gw_mac, args.gateway_mac_arg)) {
			fprintf(stderr, "%s: invalid MAC address `%s'\n", 
				CMDLINE_PARSER_PACKAGE, args.gateway_mac_arg);
			exit(EXIT_FAILURE);
		}
		zconf.gw_mac_set = 1;
	}
	if (args.seed_given) {
		zconf.seed = args.seed_arg;
		zconf.use_seed = 1;
	}
	if (args.bandwidth_given) {
		// Supported: G,g=*1000000000; M,m=*1000000 K,k=*1000 bits per second
		zconf.bandwidth = atoi(args.bandwidth_arg);
		char *suffix = args.bandwidth_arg;
		while (*suffix >= '0' && *suffix <= '9') {
			suffix++;
		}
		if (*suffix) {
			switch (*suffix) {
			case 'G': case 'g':
				zconf.bandwidth *= 1000000000;
				break;
			case 'M': case 'm':
				zconf.bandwidth *= 1000000;
				break;
			case 'K': case 'k':
				zconf.bandwidth *= 1000;
				break;
			default:
			  	fprintf(stderr, "%s: unknown bandwidth suffix '%s' "
					"(supported suffixes are G, M and K)\n",
					CMDLINE_PARSER_PACKAGE, suffix);
				exit(EXIT_FAILURE);
			}
		}
	}
	if (args.max_targets_given) {
		errno = 0;
		char *end;
	  	double v = strtod(args.max_targets_arg, &end);
		if (end == args.max_targets_arg || errno != 0) {
			fprintf(stderr, "%s: can't convert max-targets to a number\n",
					CMDLINE_PARSER_PACKAGE);
			exit(EXIT_FAILURE);
		}
		if (end[0] == '%' && end[1] == '\0') {
			// treat as percentage
			v = v * (1L << 32) / 100.;
		} else if (end[0] != '\0') {
			fprintf(stderr, "%s: extra characters after max-targets\n",
				  CMDLINE_PARSER_PACKAGE);
			exit(EXIT_FAILURE);
		}
		if (v <= 0) {
			zconf.max_targets = 0;
		}
		else if (v >= (1L << 32)) {
			zconf.max_targets = 0xFFFFFFFF;
		} else {
			zconf.max_targets = v;
		}
	}

	start_zmap();

	cmdline_parser_free(&args);
	free(params);
	return EXIT_SUCCESS;
}