/*
 * 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
 */

// probe module for performing TCP SYN scans extended for probe response attacks

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <unistd.h>
#include <string.h>
#include <assert.h>

#include "../../lib/includes.h"
#include "../fieldset.h"
#include "probe_modules.h"
#include "packet.h"

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

#include <inttypes.h>

#define ICMP_SMALLEST_SIZE 5
#define ICMP_TIMXCEED_UNREACH_HEADER_SIZE 

probe_module_t module_tcp_synscan_pra_proberesponse;
static uint32_t num_ports;

int synscan_pra_global_initialize(struct state_conf *state)
{
	num_ports = state->source_port_last - state->source_port_first + 1;
	return EXIT_SUCCESS;
}

int synscan_pra_init_perthread(void* buf, macaddr_t *src,
		macaddr_t *gw, port_h_t dst_port,
		__attribute__((unused)) void **arg_ptr)
{
	memset(buf, 0, MAX_PACKET_SIZE);
	struct ether_header *eth_header = (struct ether_header *) buf;
	make_eth_header(eth_header, src, gw);
	struct ip *ip_header = (struct ip*)(&eth_header[1]);
	uint16_t len = htons(sizeof(struct ip) + sizeof(struct tcphdr));
	make_ip_header(ip_header, IPPROTO_TCP, len);
	struct tcphdr *tcp_header = (struct tcphdr*)(&ip_header[1]);
	make_tcp_header(tcp_header, dst_port);
	return EXIT_SUCCESS;
}

int synscan_pra_make_packet(void *buf, ipaddr_n_t src_ip, ipaddr_n_t dst_ip,
		uint32_t *validation, int probe_num, __attribute__((unused)) void *arg)
{
	// ipaddr_n_t = uint32_t
	struct ether_header *eth_header = (struct ether_header *)buf;
	struct ip *ip_header = (struct ip*)(&eth_header[1]);
	struct tcphdr *tcp_header = (struct tcphdr*)(&ip_header[1]);
	uint32_t tcp_seq = validation[0];

	// needed to reach the correct destination
	ip_header->ip_dst.s_addr = dst_ip;
	uint32_t marker_value = 0;

	// how many bits to be used for IP and checksum
	int amount_bits_marker_value = zconf.markerbits_ip;
	int amount_bits_checksum = zconf.markerbits_checksum;

	if(!zconf.use_markervalue) {
		// further calculations are based on host byte order
		// lib/types.h:typedef uint32_t ipaddr_n_t; // IPv4 address network order
		// 1.2.3.4 -> 1.2 (based on amount of bits used)
		marker_value = ntohl(dst_ip) >> (32 - amount_bits_marker_value);
	}
	// take given IP address as marker value instead of destination address (Stage > 1)
	else {
		marker_value = (uint32_t)zconf.markervalue;
	}

	//log_debug("encoder", "marker value: 0x%.16" PRIX32, marker_value);

	// stores [destination IP][checksum(IP) | 0]
	/*
	291 -> b"00000123" (hex) -> b"000123" (marker value: eg 3 bytes)
	    -> chk = checksum(b"000123" -> b"00012300") = b"abcdefgh"
	    -> b"000123" (marker value) + b"abcdefgh"[:checksum_len] (marker checksum)
	*/
	uint64_t markervalue_and_checksum = 0;
	markervalue_and_checksum |= marker_value;
	// encode marker value:
	// 0x00000000 00AAAAAA -> 0xAAAAAA00 00000000
	markervalue_and_checksum <<= (64 - amount_bits_marker_value);

	if (amount_bits_checksum != 0) {
		// encode marker checksum:
		//log_info("encoder", "encoding checksum bits: %d", amount_bits_checksum);
		// 0xAAAAAA00 00000000 -> 0x00000000 AAAAAA00 -> 0xAAAAAA00
		uint32_t marker_value_for_checksum = (markervalue_and_checksum >> 32);
		//log_debug("encoder", "marker value for checksum: 0x%.8" PRIX32, marker_value_for_checksum);
		marker_value_for_checksum = htonl(marker_value_for_checksum);
		// checksum is calculated oder 4 bytes (padded from right with 0)
		// nw order -> fletcher32 -> host order -> nw order
		// return value: we need more space to shift more than 32 Bit (return value of fletcher32 if uint32_t)
		uint64_t sum_marker_value = (uint64_t)(fletcher32((uint16_t*)&marker_value_for_checksum, 2) & 0x00000000FFFFFFFF);
		//log_debug("encoder", "checksum: 0x%.16" PRIX64, sum_marker_value);

		// encode checksum:
		// 0x00000000 CCCCCCCC
		// 0xAAAAAA00 00000000 -> 0xAAAAAACC CCCCCC00
		// Unneeded trailing checksum bits don't need to be removed: just stop reading after x bits
		markervalue_and_checksum |= (sum_marker_value << (32 - amount_bits_marker_value));
	}

	int amount_bits_to_encode = amount_bits_marker_value + amount_bits_checksum;
	//printf("total bits to encode : %d\n", amount_bits_to_encode);
	//printf("%u\n", markervalue_and_checksum >> 32);
	//printf("Value to encode: 0x%.16" PRIX64 "\n", markervalue_and_checksum);

	// will be overwritten if marker type is given
	ip_header->ip_src.s_addr = src_ip;
	// will be overwritten if marker type is given
	tcp_header->th_sport = htons(get_src_port(num_ports, probe_num, validation));

	int bitpos = 1;
	// encoded bits so far
	int marker_bitcount = 0;

	//log_debug("encoder", "encoding %d bits: 0x%.16" PRIX64, amount_bits_to_encode, markervalue_and_checksum);

	// check for every available encoder marker to place destination address (+checksum)
	// order of encoding: dport, source ip, sport (source ip should not be equal to destination to avoid filtering)
	if (amount_bits_to_encode)
	while (bitpos <= 31 && marker_bitcount < amount_bits_to_encode) {
		if ((zconf.marker_encoding & bitpos) == 1) {
			//log_debug("encoder", "encoding to dport");

			if (!zconf.marker_encoding_dst_small) {
				// encode into destination port
				tcp_header->th_dport = htons((uint16_t)( (markervalue_and_checksum & 0xFFFF000000000000) >> 48) );
			} else {
				// if marker_encoding_mitigation: assume destination port is the only marker and limited
				// to 10 bits: 0->1023 (0000 0000 00 -> 1111 1111 11)
				tcp_header->th_dport = htons((uint16_t)( (markervalue_and_checksum & 0xFFC0000000000000) >> (64-10)) );
				// just use destination port, skip all other
				break;
			}
			marker_bitcount += 16;
		}
		else if ((zconf.marker_encoding & bitpos) == 2) {
			//log_debug("encoder", "encoding to src ip");
			// encode into source address
			ip_header->ip_src.s_addr = htonl((uint32_t)( (markervalue_and_checksum & (0xFFFFFFFF00000000 >> marker_bitcount)) >> (32 - marker_bitcount)) );
			//printf("source IP: %u\n", ip_header->ip_src.s_addr);
			marker_bitcount += 32;
		}
		else if ((zconf.marker_encoding & bitpos) == 4) {
			//log_debug("encoder", "encoding to sport");
			// encode into source port
			tcp_header->th_sport = htons((uint16_t)( (markervalue_and_checksum & (0xFFFF000000000000 >> marker_bitcount)) >> (48 - marker_bitcount)) );
			marker_bitcount += 16;
		}
		/*
		else if ((zconf.marker_encoding & bitpos) == 8) {
		}
		else if ((zconf.marker_encoding & bitpos) == 16) {
		}
		*/
		bitpos <<= 1;
	}

	tcp_header->th_seq = tcp_seq;
	tcp_header->th_sum = 0;
	tcp_header->th_sum = tcp_checksum(sizeof(struct tcphdr),
			ip_header->ip_src.s_addr, ip_header->ip_dst.s_addr, tcp_header);

	ip_header->ip_sum = 0;
	ip_header->ip_sum = zmap_ip_checksum((unsigned short *) ip_header);

	return EXIT_SUCCESS;
}

void synscan_pra_print_packet(FILE *fp, void* packet)
{
	struct ether_header *ethh = (struct ether_header *) packet;
	struct ip *iph = (struct ip *) &ethh[1];
	struct tcphdr *tcph = (struct tcphdr *) &iph[1];
	fprintf(fp, "tcp { source: %u | dest: %u | seq: %u | checksum: %#04X }\n",
			ntohs(tcph->th_sport),
			ntohs(tcph->th_dport),
			ntohl(tcph->th_seq),
			ntohs(tcph->th_sum));
	fprintf_ip_header(fp, iph);
	fprintf_eth_header(fp, ethh);
	fprintf(fp, "------------------------------------------------------\n");
}

// scanner feedback is meant for first stage only
int synscan_pra_validate_packet(const struct ip *ip_hdr, uint32_t len,
		__attribute__((unused))uint32_t *src_ip,
		uint32_t *validation)
{
	//log_debug("decoder", "got feedback packet");

	if (zconf.use_markervalue) {
		//log_debug("decoder", "not collecting feedback: using marker value");
		return 0;
	}

	if ((zconf.marker_encoding & 2) == 2) {
		//log_debug("decoder", "not collecting feedback: spoofed address");
		// we won't get any answers for spoofed addresses ->  only check in case of modes 1, 4, 5 (1+4)
		return 0;
	}

	if ((4*ip_hdr->ip_hl + sizeof(struct tcphdr)) > len) {
		//log_debug("decoder", "not collecting feedback: packet to large");
		// buffer not large enough to contain expected tcp header
		return 0;
	}

	// arriving packet
	uint16_t sport_remote = 0;
	uint16_t dport_remote = 0;
	uint32_t source_ip_remote = 0;
	uint32_t ack_remote = 0;
	// sent packet
	uint16_t sport_local = 0;
	//uint16_t sport_local_viamarker = 0;
	uint16_t dport_local = htons((uint16_t)zconf.target_port);

	if (ip_hdr->ip_p == IPPROTO_TCP) {
		//log_debug("decoder", "got TCP feedback");
		struct tcphdr *tcp = (struct tcphdr*)((char *) ip_hdr + 4*ip_hdr->ip_hl);
		sport_remote = tcp->th_sport;
		dport_remote = tcp->th_dport;
		ack_remote = tcp->th_ack;
		source_ip_remote = ip_hdr->ip_src.s_addr;
	}
	// in case of ICMP responses: extract IP and TCP structure
	else if (ip_hdr->ip_p == IPPROTO_ICMP) {
		//log_debug("decoder", "got ICMP feedback");

		if (((uint32_t) 4 * ip_hdr->ip_hl + ICMP_SMALLEST_SIZE) > len) {
			// buffer not large enough to contain expected icmp header
			//log_debug("decoder", "buffer not large enough to contain expected icmp header");
			return 0;
		}

		struct icmp *icmp_hdr = (struct icmp *) ((char *) ip_hdr + 4*ip_hdr->ip_hl);

		// for some packet types we must look inside the payload
		// TODO: add redirect?
		if (icmp_hdr->icmp_type == ICMP_TIMXCEED || icmp_hdr->icmp_type == ICMP_UNREACH) {
			//log_debug("decoder", "icmp type: %d, code: %d, checksum: %d",
			//				icmp_hdr->icmp_type,
			//				icmp_hdr->icmp_code,
			//				ntohs(icmp_hdr->icmp_cksum));
			// get inner IP by skipping 4 bytes ICMP header, increment by unit size did not work
			struct ip *ip_inner_icmp = (struct ip *)(((char*)icmp_hdr) + 8);

			//log_debug("decoder", "ip hl in icmp: %d", ip_inner_icmp->ip_hl);

			struct tcphdr *tcp_inner_icmp = (struct tcphdr*)((char *) ip_inner_icmp + 4*ip_inner_icmp->ip_hl);

			// replace port and adress information
			// as we got the original message back we need to take the reverse order (src => dst, dst => srs)
			sport_remote = tcp_inner_icmp->th_dport;
			dport_remote = tcp_inner_icmp->th_sport;
			source_ip_remote = ip_inner_icmp->ip_dst.s_addr;
			// Avoid filtering duplicates on ICMP repsonses
			// ZMap sets source (remote) IP itself to filter out duplicates (see src/recv.c -> handle_packet())
			// ICMP responses can arrive from same source but "belong" to different targets -> replace source ip with
			// destination IP extracted from ICMP payload (the real intended target)
			//
			*src_ip = source_ip_remote;
			//log_debug("decoder", "extracted ICMP data: remote port source <-> dst: %d <-> %d",
			//							ntohs(sport_remote),
			//							ntohs(dport_remote));
		}
		else {
			//log_debug("decoder", "icmp message type was not checkable (other than ICMP_TIMXCEED or ICMP_UNREACH): %d",
			//											icmp_hdr->icmp_type);
			return 0;
		}
	}
	else {
		// validation is based on IP, TCP and ICMP -> ignore all other
		//log_debug("decoder", "can't validate this protocol: %d", ip_hdr->ip_p);
		return 0;
	}


	uint16_t bitpos = 1;
	uint16_t marker_bitcount = 0;
	// extract source/destination port from source address (attack target)
	// destination IP = [[dport]...[sport]] -> sport (and dport) could be ommitted
	// we had at maximum markerbits_ip Bits to encode sport abd dport
	if (zconf.markerbits_ip) {
	uint32_t source_ip_remote_ho = ntohl(source_ip_remote);
	// remove not encoded bits if any
	int bits_to_remove = 32 - zconf.markerbits_ip;
	source_ip_remote_ho = (source_ip_remote_ho >> bits_to_remove) << bits_to_remove;

	while (bitpos <= 31 && marker_bitcount < zconf.markerbits_ip) {
		if ((zconf.marker_encoding & bitpos) == 1) {
			//log_debug("decoder", "decoding dport");

			if (!zconf.marker_encoding_dst_small) {
				// assume full destination port was encoded
				dport_local = (source_ip_remote_ho & 0xFFFF0000) >> 16;
				//log_debug("decoder", "decoded dport: %d", dport_local);
				marker_bitcount += 16;
			} else {
				//log_debug("decoder", "small encoding");
				// if marker encoding mitigation: assume destination port is the only marker (see encoding)
				dport_local = (source_ip_remote_ho & 0xFFC00000) >> (32-10);
				// just destination port, skip all other
				break;
			}
			dport_local = htons(dport_local);
		}
		else if ((zconf.marker_encoding & bitpos) == 4) {
			//log_debug("decoder", "decoding sport");
			sport_local = (source_ip_remote_ho & (0xFFFF0000 >> marker_bitcount)) >> (16 - marker_bitcount);
			//log_debug("decoder", "decoded sport: %d", sport_local);
			// [dst][src] -> [src][dst]: remove checksum bits from remote dst
			uint32_t dport_remote_ho = ntohs(dport_remote);
			dport_remote_ho = (dport_remote_ho >> zconf.markerbits_checksum) << zconf.markerbits_checksum;
			dport_remote = htons(dport_remote_ho);
			//log_debug("decoder", "changed remote dport to: %d", dport_remote_ho);

			marker_bitcount += 16;
			//sport_local_viamarker = 1;
			sport_local = htons(sport_local);
		}
		/*
		else if ((zconf.marker_encoding & bitpos) == 8) {
		}
		else if ((zconf.marker_encoding & bitpos) == 16) {
		}
		*/
		bitpos <<= 1;
	}
	}
	//log_debug("decoder", "local src/dst -> remote src/dst: %d %d -> %d %d", ntohs(sport_local), ntohs(dport_local), ntohs(sport_remote), ntohs(dport_remote));


	// validate remote source port. Check via "&" as source port could be only partially encoded
	if ((dport_local & sport_remote) != dport_local) {
		//log_debug("decoder", "ports did not match: local dst != remote src ->  %d != %d", ntohs(dport_local), ntohs(sport_remote));
		return 0;
	}
	// validate remote destination port: extracted via arriving IP address or local config
	//if ((sport_local_viamarker && (sport_local != dport_remote)) || !check_dst_port(ntohs(dport_remote), num_ports, validation)) {
	if (sport_local != dport_remote) {
		/*
		log_debug("decoder", "ports did not match: local src != remote dst -> %d != %d (via marker? %d)",
										ntohs(sport_local),
										ntohs(dport_remote),
										sport_local_viamarker);
		*/
		return 0;
	}
	// validate tcp acknowledgement number
	if (ack_remote != 0 && ( htonl(ack_remote) != htonl(validation[0])+1 ) ) {
		//log_debug("decoder", "sequence number did not match");
		return 0;
	}

	//log_debug("decoder", "This was an answer packet!!!");

	return 1;
}

void synscan_pra_process_packet(const u_char *packet,
		__attribute__((unused)) uint32_t len, fieldset_t *fs)
{
	//log_debug("encoder", "writing feedback packet to file");
	// Idea: how to correlate reponse <-> marker group? This has to be done by comparing
	// response IP addresses to groups.
	struct ip *ip_hdr = (struct ip *)&packet[sizeof(struct ether_header)];
	// WARNING: fields have to be written in the same order as defined in fields[]
	// That is the reason why this code looks a bit redundant

	if (ip_hdr->ip_p == IPPROTO_TCP) {
		//log_debug("encoder", "storing TCP response");
		struct tcphdr *tcp = (struct tcphdr*)((char *)ip_hdr + 4*ip_hdr->ip_hl);

		fs_add_uint64(fs, "sport", (uint64_t) ntohs(tcp->th_sport));
		fs_add_uint64(fs, "dport", (uint64_t) ntohs(tcp->th_dport));
		fs_add_uint64(fs, "seqnum", (uint64_t) ntohl(tcp->th_seq));
		fs_add_uint64(fs, "acknum", (uint64_t) ntohl(tcp->th_ack));
		fs_add_uint64(fs, "window", (uint64_t) ntohs(tcp->th_win));
		fs_add_string(fs, "daddr_inner_icmp", (char*) "(None)", 0);

		if (tcp->th_flags & TH_RST) { // RST packet
			fs_add_string(fs, "classification", (char*) "rst", 0);
			fs_add_uint64(fs, "success", 0);
		} else { // SYNACK packet
			fs_add_string(fs, "classification", (char*) "synack", 0);
			fs_add_uint64(fs, "success", 1);
		}
	}
	else {
	// prefiltering was already done in synscan_pra_validate_packet()
		//log_debug("encoder", "storing ICMP response");

		struct icmp *icmp_hdr = (struct icmp *) ((char *) ip_hdr + 4*ip_hdr->ip_hl);
		struct ip *ip_hdr_from_icmp = (struct ip *)(((char * )icmp_hdr) + 8);
		struct tcphdr *tcp = (struct tcphdr*)((char *)ip_hdr_from_icmp + 4*ip_hdr_from_icmp->ip_hl);

		// switch source/destination port: source is always the attacker
		fs_add_uint64(fs, "sport", (uint64_t) ntohs(tcp->th_dport));
		fs_add_uint64(fs, "dport", (uint64_t) ntohs(tcp->th_sport));
		fs_add_uint64(fs, "seqnum", (uint64_t) 0);
		fs_add_uint64(fs, "acknum", (uint64_t) 0);
		fs_add_uint64(fs, "window", (uint64_t) 0);
		fs_add_string(fs, "daddr_inner_icmp", make_ip_str(ip_hdr_from_icmp->ip_dst.s_addr), 0);

		if (icmp_hdr->icmp_type == ICMP_UNREACH)
			fs_add_string(fs, "classification", (char*) "icmp_unreach", 0);
		else if(icmp_hdr->icmp_type == ICMP_TIMXCEED)
			fs_add_string(fs, "classification", (char*) "icmp_timeexceed", 0);
		else
			fs_add_string(fs, "classification", (char*) "icmp_other", 0);

		fs_add_uint64(fs, "success", 1);
	}
}

static fielddef_t fields[] = {
	{.name = "sport",  .type = "int", .desc = "TCP source port"},
	{.name = "dport",  .type = "int", .desc = "TCP destination port"},
	{.name = "seqnum", .type = "int", .desc = "TCP sequence number"},
	{.name = "acknum", .type = "int", .desc = "TCP acknowledgement number"},
	{.name = "window", .type = "int", .desc = "TCP window"},
	{.name = "daddr_inner_icmp", .type="string", .desc = "Destination address of IP header contained in ICMP response"},
	{.name = "classification", .type="string", .desc = "packet classification"},
	{.name = "success", .type="int", .desc = "is response considered success"}
};

probe_module_t module_tcp_synscan_proberesponse = {
	.name = "tcp_synscan_pra",
	.packet_length = 54,
	//.packet_length = 82,
	.pcap_filter = "icmp || (tcp && tcp[13] & 4 != 0 || tcp[13] == 18)",
	.pcap_snaplen = 96,
	.port_args = 1,
	.global_initialize = &synscan_pra_global_initialize,
	.thread_initialize = &synscan_pra_init_perthread,
	.make_packet = &synscan_pra_make_packet,
	.print_packet = &synscan_pra_print_packet,
	.process_packet = &synscan_pra_process_packet,
	.validate_packet = &synscan_pra_validate_packet,
	.close = NULL,
	.helptext = "Extended TCP probe module for probe response attacks. ",
	.fields = fields,
	.numfields = 8};