#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <assert.h>
#include <arpa/inet.h>

#ifdef __cplusplus
extern "C"
{


#include "cJSON.h"
#include "dns_decoder.h"

#include "stellar/stellar.h"
#include "stellar/session.h"
#include "stellar/session_exdata.h"
#include "stellar/session_mq.h"

}
#endif

struct dns_decoder_test_plugin_env
{
	int plugin_id;
	int topic_id;
	int result_index;
};

extern "C" int commit_test_result_json(cJSON *node, const char *name);

void dns_real_result_write_file(char *result_str)
{
	FILE *fp=fopen("dns_real_result.json", "a+");
	if(fp!=NULL)
	{
		fwrite(result_str, 1, strlen(result_str), fp);
		fclose(fp);
	}
}

void dns_resource_record_str2hex_append(cJSON *one_rr_object, uint8_t *str, size_t str_sz, const char *key)
{
	if(one_rr_object==NULL || str==NULL || str_sz==0 || key==NULL)
	{
		return;
	}

	size_t offset=0;
	char hex_buff[4096]={0};
	size_t hex_buff_sz=sizeof(hex_buff);
	
	for(size_t i=0; i<str_sz && i<hex_buff_sz/2; i++)
	{
		offset+=snprintf((char *)(hex_buff+offset), hex_buff_sz-offset, "%02x", str[i]);
		if(offset+2 >= hex_buff_sz)
		{
			break;
		}
	}

	if(offset>0)
	{
		cJSON_AddStringToObject(one_rr_object, key, hex_buff);
	}
}

void dns_resource_record_dstring_append(cJSON *one_rr_object, struct dstring *dstr, const char *key)
{
	if(dstr->value_sz > 0)
	{
		cJSON_AddStringToObject(one_rr_object, key, (char *)(dstr->value));
	}
}

int dns_resource_record_json_exporter(cJSON *object, struct dns_resource_record *rr_array, uint16_t n_rr, const char *rr_type, int *dns_sec)
{
	if(object==NULL || rr_array==NULL || n_rr==0 || dns_sec==NULL)
	{
		return 0;
	}

	char ip_str[128];
	cJSON *dns_rr_array=cJSON_CreateArray();
	
	for(uint16_t i=0; i<n_rr; i++)
	{
		cJSON *one_rr_object=cJSON_CreateObject();
		struct dns_resource_record *dns_rr=&(rr_array[i]);

		if(dns_rr->type == DNS_RR_TYPE_OPT)
		{			
			cJSON_AddStringToObject(one_rr_object, "name", (const char *)(dns_rr->qname.value));
			cJSON_AddNumberToObject(one_rr_object, "type", dns_rr->type);
			cJSON_AddNumberToObject(one_rr_object, "udp_payload", dns_rr->rr_class);
			cJSON_AddNumberToObject(one_rr_object, "rcode", (int)(dns_rr->ttl>>24));
			cJSON_AddNumberToObject(one_rr_object, "version", (int)((dns_rr->ttl>>16)&0xFF));
			cJSON_AddNumberToObject(one_rr_object, "Z", (int)(dns_rr->ttl&&0xFFFF));
			cJSON_AddNumberToObject(one_rr_object, "rdlength", dns_rr->rdlength);
		}
		else
		{
			cJSON_AddStringToObject(one_rr_object, "name", (const char *)(dns_rr->qname.value));
			cJSON_AddNumberToObject(one_rr_object, "type", dns_rr->type);
			cJSON_AddNumberToObject(one_rr_object, "class", dns_rr->rr_class);
			cJSON_AddNumberToObject(one_rr_object, "ttl", dns_rr->ttl);
			cJSON_AddNumberToObject(one_rr_object, "rdlength", dns_rr->rdlength);
		}

		if(dns_rr->rdlength==0)
		{
			cJSON_AddItemToArray(dns_rr_array, one_rr_object);
			continue;
		}
		
		switch(dns_rr->type)
		{
			case DNS_RR_TYPE_A:
				inet_ntop(AF_INET, (void *)(dns_rr->rdata.a.value), ip_str, sizeof(ip_str));
				cJSON_AddStringToObject(one_rr_object, "a", ip_str);
				break;
			case DNS_RR_TYPE_AAAA:
				inet_ntop(AF_INET6, dns_rr->rdata.aaaa.value, ip_str, sizeof(ip_str));
				cJSON_AddStringToObject(one_rr_object, "aaaa", ip_str);
				break;
			case DNS_RR_TYPE_NS:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.ns), "ns");
				break;
			case DNS_RR_TYPE_MD:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.md), "md");
				break;
			case DNS_RR_TYPE_MF:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.mf), "mf");
				break;
			case DNS_RR_TYPE_CNAME:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.cname), "cname");
				break;
			case DNS_RR_TYPE_MB:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.mb), "mb");
				break;
			case DNS_RR_TYPE_MG:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.mg), "mg");
				break;
			case DNS_RR_TYPE_MR:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.mr), "mr");
				break;
			case DNS_RR_TYPE_TXT:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.txt), "txt");
				cJSON_AddNumberToObject(one_rr_object, "size", dns_rr->rdata.txt.value_sz);
				break;
			case DNS_RR_TYPE_NULL:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.null), "null");
				cJSON_AddNumberToObject(one_rr_object, "size", dns_rr->rdata.null.value_sz);
				break;
			case DNS_RR_TYPE_PTR:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.ptr), "ptr");
				break;
			case DNS_RR_TYPE_HINFO:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.hinfo.cpu), "cpu");
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.hinfo.os), "os");
				break;
			case DNS_RR_TYPE_MINFO:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.minfo.rmailbx), "rmailbx");
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.minfo.emailbx), "emailbx");
				break;
			case DNS_RR_TYPE_MX:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.mx.exchange), "exchange");
				cJSON_AddNumberToObject(one_rr_object, "preference", dns_rr->rdata.mx.preference);
				break;
			case DNS_RR_TYPE_RP:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.rp.mailbox), "mailbox");
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.rp.txt_rr), "txt_rr");
				break;
			case DNS_RR_TYPE_SOA:
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.soa.mname), "mname");
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.soa.rname), "rname");
				cJSON_AddNumberToObject(one_rr_object, "serial", dns_rr->rdata.soa.serial);
				cJSON_AddNumberToObject(one_rr_object, "refresh", dns_rr->rdata.soa.refresh);
				cJSON_AddNumberToObject(one_rr_object, "retry", dns_rr->rdata.soa.retry);
				cJSON_AddNumberToObject(one_rr_object, "cname", dns_rr->rdata.soa.expire);
				cJSON_AddNumberToObject(one_rr_object, "minimum", dns_rr->rdata.soa.minimum);
				break;

			case DNS_RR_TYPE_WKS:
				cJSON_AddStringToObject(one_rr_object, "addr", ip_str);
				cJSON_AddNumberToObject(one_rr_object, "protocol", dns_rr->rdata.wks.protocol);
				cJSON_AddNumberToObject(one_rr_object, "size", dns_rr->rdata.wks.size);
				
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.wks.bitmap, dns_rr->rdata.wks.size, "bitmap");
				break;
			case DNS_RR_TYPE_OPT:
				break;
			case DNS_RR_TYPE_DS:
				*dns_sec = 2;
				cJSON_AddNumberToObject(one_rr_object, "key_tag", dns_rr->rdata.ds.key_tag);
				cJSON_AddNumberToObject(one_rr_object, "algo", dns_rr->rdata.ds.algo);
				cJSON_AddNumberToObject(one_rr_object, "digest_type", dns_rr->rdata.ds.digest_type);
				cJSON_AddNumberToObject(one_rr_object, "digest_len", dns_rr->rdata.ds.digest_len);
				
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.ds.digest, dns_rr->rdata.ds.digest_len, "digest");					
				break;
			case DNS_RR_TYPE_RRSIG:
				*dns_sec = 2;
				cJSON_AddNumberToObject(one_rr_object, "type_covered", dns_rr->rdata.rrsig.type_covered);
				cJSON_AddNumberToObject(one_rr_object, "algo", dns_rr->rdata.rrsig.algo);
				cJSON_AddNumberToObject(one_rr_object, "labels", dns_rr->rdata.rrsig.labels);
				cJSON_AddNumberToObject(one_rr_object, "original_ttl", dns_rr->rdata.rrsig.original_ttl);
				cJSON_AddNumberToObject(one_rr_object, "sig_expiration", dns_rr->rdata.rrsig.sig_expiration);
				cJSON_AddNumberToObject(one_rr_object, "sig_inception", dns_rr->rdata.rrsig.sig_inception);
				cJSON_AddNumberToObject(one_rr_object, "key_tag", dns_rr->rdata.rrsig.key_tag);

				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.rrsig.signer_name), "signer_name");
				
				cJSON_AddNumberToObject(one_rr_object, "signature_len", dns_rr->rdata.rrsig.signature_len);
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.rrsig.signature, dns_rr->rdata.rrsig.signature_len, "signature");
				break;
			case DNS_RR_TYPE_NSEC:
				*dns_sec = 2;
				dns_resource_record_dstring_append(one_rr_object, &(dns_rr->rdata.nsec.next_domain), "next_domain");
				cJSON_AddNumberToObject(one_rr_object, "maps_len", dns_rr->rdata.nsec.type_bit_maps.value_sz);
				
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.nsec.type_bit_maps.value, dns_rr->rdata.nsec.type_bit_maps.value_sz, "type_bit_maps");
				break;
			case DNS_RR_TYPE_DNSKEY:
				*dns_sec = 2;
				cJSON_AddNumberToObject(one_rr_object, "flags", dns_rr->rdata.dnskey.flags);
				cJSON_AddNumberToObject(one_rr_object, "protocol", dns_rr->rdata.dnskey.protocol);
				cJSON_AddNumberToObject(one_rr_object, "algo", dns_rr->rdata.dnskey.algo);
				cJSON_AddNumberToObject(one_rr_object, "public_key_len", dns_rr->rdata.dnskey.public_key_len);
				
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.dnskey.public_key, dns_rr->rdata.dnskey.public_key_len, "public_key");
				break;
			case DNS_RR_TYPE_NSEC3:
				*dns_sec = 2;
				cJSON_AddNumberToObject(one_rr_object, "hash_algo", dns_rr->rdata.nsec3.hash_algo);
				cJSON_AddNumberToObject(one_rr_object, "flags", dns_rr->rdata.nsec3.flags);
				cJSON_AddNumberToObject(one_rr_object, "iteration", dns_rr->rdata.nsec3.iteration);

				cJSON_AddNumberToObject(one_rr_object, "salt_len", dns_rr->rdata.nsec3.salt_len);
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.nsec3.salt_value, dns_rr->rdata.nsec3.salt_len,  "salt_value");
				
				cJSON_AddNumberToObject(one_rr_object, "hash_len", dns_rr->rdata.nsec3.hash_len);
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.nsec3.next_hash_owner, dns_rr->rdata.nsec3.hash_len, "next_hash_owner");
				
				cJSON_AddNumberToObject(one_rr_object, "maps_len", dns_rr->rdata.nsec3.type_bit_maps.value_sz);
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.nsec3.type_bit_maps.value,dns_rr->rdata.nsec3.type_bit_maps.value_sz, "type_bit_maps");
				break;
			case DNS_RR_TYPE_NSEC3PARAM:
				cJSON_AddNumberToObject(one_rr_object, "hash_algo", dns_rr->rdata.nsec3param.hash_algo);
				cJSON_AddNumberToObject(one_rr_object, "flags", dns_rr->rdata.nsec3param.flags);
				cJSON_AddNumberToObject(one_rr_object, "iteration", dns_rr->rdata.nsec3param.iteration);
				cJSON_AddNumberToObject(one_rr_object, "salt_len", dns_rr->rdata.nsec3param.salt_len);
				
				dns_resource_record_str2hex_append(one_rr_object, dns_rr->rdata.nsec3param.salt_value, dns_rr->rdata.nsec3param.salt_len, "salt_value");
				break;
			default:
				break;
		}
		
		cJSON_AddItemToArray(dns_rr_array, one_rr_object);
	}

	cJSON_AddItemToObject(object, rr_type, dns_rr_array);
	
	return 1;
}

void dns_decoder_test_message_cb(struct session *ss, int topic_id, const void *msg, void *per_session_ctx, void *plugin_env_str)
{
	struct dns_message *dns_msg=(struct dns_message *)msg;
	enum dns_message_type type=dns_message_type_get(dns_msg);
	if(type==DNS_MESSAGE_TRANSACTION_BEGIN || type==DNS_MESSAGE_TRANSACTION_END)
	{
		return;
	}

	//uint16_t id=dns_message_header_id_get(dns_msg);
	struct dns_decoder_test_plugin_env *plugin_env=(struct dns_decoder_test_plugin_env *)plugin_env_str;

    cJSON *real_result=cJSON_CreateObject();
    cJSON_AddStringToObject(real_result, "Tuple4", session_get0_readable_addr(ss));

	uint16_t n_question=0;
	struct dns_query_question *question=NULL;
	dns_message_question_get0(dns_msg, &question, &n_question);
	if(n_question>0 && question!=NULL)
	{
		cJSON_AddStringToObject(real_result, "dns_qname", (char *)question->qname);
		cJSON_AddNumberToObject(real_result, "dns_qtype", question->qtype);
		cJSON_AddNumberToObject(real_result, "dns_qclass", question->qclass);
	}

	struct dns_flag *flag=dns_message_header_flag_get0(dns_msg);
	if(flag!=NULL)
	{
		cJSON_AddNumberToObject(real_result, "dns_qr", (double)(flag->qr));
		cJSON_AddNumberToObject(real_result, "dns_opcode", (double)(flag->opcode));
		cJSON_AddNumberToObject(real_result, "dns_rd", (double)(flag->rd));
	}

	cJSON *rr_object=cJSON_CreateObject();
	uint16_t n_answer_rr=0;
	struct dns_resource_record *answer_rr=NULL;
	dns_message_answer_resource_record_get0(dns_msg, &answer_rr, &n_answer_rr);

	uint16_t n_authority_rr=0;
	struct dns_resource_record *authority_rr=NULL;
	dns_message_authority_resource_record_get0(dns_msg, &authority_rr, &n_authority_rr);

	uint16_t n_additional_rr=0;
	struct dns_resource_record *additional_rr=NULL;
	dns_message_additional_resource_record_get0(dns_msg, &additional_rr, &n_additional_rr);

	int dns_sec=1;
	dns_resource_record_json_exporter(rr_object, answer_rr, n_answer_rr, "answer", &dns_sec);
	dns_resource_record_json_exporter(rr_object, authority_rr, n_authority_rr, "authority", &dns_sec);
	dns_resource_record_json_exporter(rr_object, additional_rr, n_additional_rr, "additional", &dns_sec);

	cJSON_AddItemToObject(real_result, "rr", rr_object);

    char result_name[16]="";
    sprintf(result_name, "DNS_RESULT_%d", plugin_env->result_index++);

	char *real_result_str=cJSON_Print(real_result);
	dns_real_result_write_file(real_result_str);
	free(real_result_str);

    commit_test_result_json(real_result, result_name);
}

void *dns_decoder_test_per_session_context_new(struct session *sess, void *plugin_env)
{
    return NULL;
}

void dns_decoder_test_per_session_context_free(struct session *sess, void *session_ctx, void *plugin_env)
{

}

extern "C" void *dns_decoder_test_init(struct stellar *st)
{
	struct dns_decoder_test_plugin_env *plugin_env=(struct dns_decoder_test_plugin_env *)calloc(1, sizeof(struct dns_decoder_test_plugin_env));
	plugin_env->result_index=1;

    plugin_env->plugin_id=stellar_session_plugin_register(st, dns_decoder_test_per_session_context_new, dns_decoder_test_per_session_context_free, plugin_env);
    if(plugin_env->plugin_id<0)
    {
        printf("dns_decoder_test_init: stellar_plugin_register failed !!!\n");
		exit(-1);
    }

	plugin_env->topic_id=stellar_session_mq_get_topic_id(st, DNS_MESSAGE_TOPIC);
	if(plugin_env->topic_id<0)
	{
		printf("stellar_session_mq_get_topic_id failed, topic: %s \n", DNS_MESSAGE_TOPIC);
		exit(-1);
	}

	stellar_session_mq_subscribe(st, plugin_env->topic_id, dns_decoder_test_message_cb, plugin_env->plugin_id);

    return (void *)plugin_env;
}

extern "C" void dns_decoder_test_exit(void *plugin_env_str)
{
	if(plugin_env_str!=NULL)
	{
		free(plugin_env_str);
	}
}