#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <netinet/ip.h>
#include <netinet/ip6.h>
#include <netinet/tcp.h>
#include <netinet/ether.h>
#include <pcap/sll.h>
#include <pcap/pcap.h>
#include "gtest/gtest.h"

#include "socks_decoder_internal.h"
extern "C" {
#include "stellar/session.h"
#include "dummy.h"
}

#define LINUX_COOKED_CAPTURE 20

extern int socks_process(struct socks_decoder_info *socks_decoder_info, struct session *sess, struct socks_tunnel_stream *stream, const char *payload, size_t payload_len);
struct pcap_loop_arg
{
	pcap_t *pcap_handle;
	struct socks_tunnel_stream *stream;
};

struct socks_decoder_info *socks_decoder_info = NULL;

static void pcap_handle_cb(pcap_loop_arg *userarg, const struct pcap_pkthdr *pkthdr, const u_char *packet)
{
	struct pcap_loop_arg *arg = (struct pcap_loop_arg *)userarg;
	struct session sess;
	int payload_len = 0;
	char *payload = NULL;
	unsigned short sport, dport;
	unsigned short eth_proto_type;
	unsigned char *ip_header;

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

	int data_link_type = pcap_datalink(arg->pcap_handle);
	switch (data_link_type) {
		case DLT_EN10MB:
			eth_proto_type = ntohs(*(unsigned short *)(packet + 12));
			ip_header = (unsigned char *)(packet + sizeof(struct ethhdr));
			break;
		case 276://DLT_LINUX_SLL2
			eth_proto_type = ntohs(*(unsigned short *)packet);
			ip_header = (unsigned char *)(packet + LINUX_COOKED_CAPTURE);
			break;
		default:
			return;
	}

	if (eth_proto_type == ETH_P_IP) {
		int l4_proto = *(unsigned char *)(ip_header + 9);
		if (l4_proto == IPPROTO_TCP) {
			int ip_total_len = ntohs(*(unsigned short *)(ip_header + 2));
			int ip_header_len = (*(unsigned char *)ip_header & 0x0f) * 4;
			int tcp_header_len = 4 * ((*(unsigned char *)(ip_header + sizeof(iphdr) + 12) & 0xf0) >> 4);
			payload_len = ip_total_len - ip_header_len - tcp_header_len;
			payload = (char *)ip_header + ip_header_len + tcp_header_len;
		} else {
			return;
		}

		sport = ntohs(*(unsigned short *)(ip_header + sizeof(iphdr) + 0));
		dport = ntohs(*(unsigned short *)(ip_header + sizeof(iphdr) + 2));
		
	} else if (eth_proto_type == ETH_P_IPV6) {
		int l4_proto = *(unsigned char *)(ip_header + 6);
		if (l4_proto == IPPROTO_TCP) {
			int tcp_header_len = 4 * ((*(unsigned char *)(ip_header + sizeof(struct ip6_hdr) + 12) & 0xf0) >> 4);
			payload_len = pkthdr->caplen - (ip_header - packet) - sizeof(struct ip6_hdr) - tcp_header_len;
			payload = (char *)ip_header + sizeof(struct ip6_hdr) + tcp_header_len;
		} else {
			return;
		}

		sport = ntohs(*(unsigned short *)(ip_header + sizeof(struct ip6_hdr) + 0));
		dport = ntohs(*(unsigned short *)(ip_header + sizeof(struct ip6_hdr) + 2));
	} else {
		return;
	}

	session_set_current_state(&sess, SESSION_STATE_ACTIVE);

	if (sport > dport)
	{
		dummy_set_direction(FLOW_TYPE_C2S);
	} else {
		dummy_set_direction(FLOW_TYPE_S2C);
	}
	socks_process(socks_decoder_info, &sess, arg->stream, payload, payload_len);

}

TEST(socks_decoder, socks4)
{
	char error[100];
	struct pcap_loop_arg arg;
	struct socks_tunnel_stream *stream = (struct socks_tunnel_stream *)calloc(1, sizeof(struct socks_tunnel_stream));
	pcap_t * handle = pcap_open_offline("pcap/socks4.pcap", error);
	ASSERT_NE(handle, nullptr);

	memset(&arg, 0, sizeof(arg));
	arg.stream = stream;
	arg.pcap_handle = handle;

	pcap_loop(handle, -1, (pcap_handler)pcap_handle_cb, (u_char *)&arg);

	EXPECT_EQ(stream->info.version, SOCKS_VERSION_4);
	EXPECT_EQ(stream->info.dst_addr.type, SOCKS_ADDR_IPV4);
	EXPECT_EQ(stream->info.dst_addr.ipv4, inet_addr("93.184.216.119"));
	EXPECT_EQ(stream->info.dst_addr.port, htons(80));
	EXPECT_EQ(stream->info.user_name.iov_len, 0);
	EXPECT_EQ(stream->info.password.iov_len, 0);
	EXPECT_EQ(stream->client_state, SS_END);
	EXPECT_EQ(stream->server_state, SS_END);

	pcap_close(handle);
	free(stream);
}

TEST(socks_decoder, socks4a_domain)
{
	char error[100];
	struct pcap_loop_arg arg;
	struct socks_tunnel_stream *stream = (struct socks_tunnel_stream *)calloc(1, sizeof(struct socks_tunnel_stream));
	pcap_t * handle = pcap_open_offline("pcap/socks4a_domain.pcap", error);
	ASSERT_NE(handle, nullptr);

	memset(&arg, 0, sizeof(arg));
	arg.stream = stream;
	arg.pcap_handle = handle;

	pcap_loop(handle, -1, (pcap_handler)pcap_handle_cb, (u_char *)&arg);

	EXPECT_EQ(stream->info.version, SOCKS_VERSION_4);
	EXPECT_EQ(stream->info.dst_addr.type, SOCKS_ADDR_FQDN);
	EXPECT_EQ(stream->info.dst_addr.fqdn.iov_len, strlen("www.example.com"));
	EXPECT_STREQ("www.example.com", (char *)stream->info.dst_addr.fqdn.iov_base);
	EXPECT_EQ(stream->info.dst_addr.port, htons(80));
	EXPECT_EQ(stream->info.user_name.iov_len, 0);
	EXPECT_EQ(stream->info.password.iov_len, 0);
	EXPECT_EQ(stream->client_state, SS_END);
	EXPECT_EQ(stream->server_state, SS_END);

	pcap_close(handle);
	free(stream->info.dst_addr.fqdn.iov_base);
	free(stream);
}

TEST(socks_decoder, socks5_no_auth)
{
	char error[100];
	struct pcap_loop_arg arg;
	struct socks_tunnel_stream *stream = (struct socks_tunnel_stream *)calloc(1, sizeof(struct socks_tunnel_stream));
	pcap_t * handle = pcap_open_offline("pcap/socks5_no_auth.pcap", error);
	ASSERT_NE(handle, nullptr);

	memset(&arg, 0, sizeof(arg));
	arg.stream = stream;
	arg.pcap_handle = handle;

	pcap_loop(handle, -1, (pcap_handler)pcap_handle_cb, (u_char *)&arg);

	EXPECT_EQ(stream->info.version, SOCKS_VERSION_5);
	EXPECT_EQ(stream->info.dst_addr.type, SOCKS_ADDR_IPV4);
	EXPECT_EQ(stream->info.dst_addr.ipv4, inet_addr("93.184.216.119"));
	EXPECT_EQ(stream->info.dst_addr.port, htons(80));
	EXPECT_EQ(stream->info.user_name.iov_len, 0);
	EXPECT_EQ(stream->info.password.iov_len, 0);
	EXPECT_EQ(stream->client_state, SS_END);
	EXPECT_EQ(stream->server_state, SS_END);

	pcap_close(handle);
	free(stream);
}

TEST(socks_decoder, socks5_auth_success)
{
	char error[100];
	struct pcap_loop_arg arg;
	struct socks_tunnel_stream *stream = (struct socks_tunnel_stream *)calloc(1, sizeof(struct socks_tunnel_stream));
	pcap_t * handle = pcap_open_offline("pcap/socks5_auth_success.pcap", error);
	ASSERT_NE(handle, nullptr);

	memset(&arg, 0, sizeof(arg));
	arg.stream = stream;
	arg.pcap_handle = handle;

	pcap_loop(handle, -1, (pcap_handler)pcap_handle_cb, (u_char *)&arg);

	EXPECT_EQ(stream->info.version, SOCKS_VERSION_5);
	EXPECT_EQ(stream->info.dst_addr.type, SOCKS_ADDR_IPV4);
	EXPECT_EQ(stream->info.dst_addr.ipv4, inet_addr("93.184.216.119"));
	EXPECT_EQ(stream->info.dst_addr.port, htons(80));
	EXPECT_EQ(stream->info.user_name.iov_len, strlen("testuser"));
	EXPECT_STREQ("testuser", (char *)stream->info.user_name.iov_base);
	EXPECT_EQ(stream->info.password.iov_len, strlen("testuser"));
	EXPECT_STREQ("testuser", (char *)stream->info.password.iov_base);
	EXPECT_EQ(stream->client_state, SS_END);
	EXPECT_EQ(stream->server_state, SS_END);

	pcap_close(handle);
	free(stream->info.user_name.iov_base);
	free(stream->info.password.iov_base);
	free(stream);
}

TEST(socks_decoder, socks5_auth_fail)
{
	char error[100];
	struct pcap_loop_arg arg;
	struct socks_tunnel_stream *stream = (struct socks_tunnel_stream *)calloc(1, sizeof(struct socks_tunnel_stream));
	pcap_t * handle = pcap_open_offline("pcap/socks5_auth_failed.pcap", error);
	ASSERT_NE(handle, nullptr);

	memset(&arg, 0, sizeof(arg));
	arg.stream = stream;
	arg.pcap_handle = handle;

	pcap_loop(handle, -1, (pcap_handler)pcap_handle_cb, (u_char *)&arg);

	EXPECT_EQ(stream->client_state, SS_SUB);
	EXPECT_EQ(stream->server_state, SS_FAILED);

	pcap_close(handle);
	free(stream->info.user_name.iov_base);
	free(stream->info.password.iov_base);
	free(stream);
}

int main(int argc, char **argv)
{
	testing::InitGoogleTest(&argc, argv);

	socks_decoder_info = (struct socks_decoder_info *)calloc(1, sizeof(struct socks_decoder_info));
	
	int result = RUN_ALL_TESTS();

	free(socks_decoder_info);

	return result;
}