#include"stoon.h"

#ifdef _WIN32
#include<winsock2.h>
#include<ws2tcpip.h>
#include<ntsecapi.h>
#else
#include<netdb.h>
#include<sys/socket.h>
#include<sys/random.h>
#include<arpa/inet.h>
#endif
#include<unistd.h>
#include<sys/types.h>
#include<sys/time.h>
#include<stdint.h>
#include<string.h>
#include<stdio.h>
#include<errno.h>
#include<ifaddrs.h>

#define STUN_BINDING_REQUEST 0x0001
#define STUN_BINDING_INDICATION 0x1100
#define STUN_BINDING_RESPONSE 0x0101
#define STUN_XOR_MAPPED_ADDRESS 0x0020

#define STUN_MAGIC 0x2112A442

#define MAX_ATTRIB_BUFFER_SIZE 128

#define STUN_NETFAM_IPV4 1
#define STUN_NETFAM_IPV6 2

#ifdef _WIN32
#define RAND(b, i) RtlGenRandom(b, i)
#else
#define RAND(b, i) getrandom(b, i, 0)
#endif

struct StunMsg {
	uint16_t type;
	uint16_t len;
	uint32_t magic;
	uint8_t id[12];
};

static int stoon_init_mini(struct Stoon *this, struct addrinfo *serv, uint16_t myport) {
#ifdef _WIN32
	errno = 0;
	int fd = socket(serv->ai_family, serv->ai_socktype, serv->ai_protocol);
	ioctlsocket(fd, FIONBIO, &(unsigned long) {1});
#else
	int fd = socket(serv->ai_family, serv->ai_socktype | SOCK_NONBLOCK, serv->ai_protocol);
#endif
	
	if(serv->ai_family == AF_INET6) {
		setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, (void*) &(int) {1}, sizeof(int));
	}
	
	char p[6];
	sprintf(p, "%u", myport);
	
	struct addrinfo *myaddrinfo;
	if(getaddrinfo(NULL, p, &(struct addrinfo) {.ai_family = serv->ai_family, .ai_socktype = SOCK_DGRAM, .ai_protocol = IPPROTO_UDP, .ai_flags = AI_PASSIVE}, &myaddrinfo)) {
		close(fd);
		return -1;
	}
	if(bind(fd, myaddrinfo->ai_addr, myaddrinfo->ai_addrlen)) {
		close(fd);
		return -2;
	}
	
	freeaddrinfo(myaddrinfo);
	
#ifdef _WIN32
	IP_ADAPTER_ADDRESSES *addrs = malloc(1024 * 64);
	if(GetAdaptersAddresses(serv->ai_family, GAA_FLAG_SKIP_MULTICAST, NULL, &addrs) == NO_ERROR) {
		for(IP_ADAPTER_ADDRESSES *ifa = addrs; ifa; ifa = ifa->Next) {
			if(serv->ai_family == AF_INET6) {
				uint8_t *addr = (void*) &((struct sockaddr_in6*) ifa->FirstUnicastAddress->Address.lpSockaddr)->sin6_addr;
				if(addr[0] == 0 && addr[1] == 0 && addr[2] == 0 && addr[3] == 0 && addr[4] == 0 && addr[5] == 0 && addr[6] == 0 && addr[7] == 0 && addr[8] == 0 && addr[9] == 0 && addr[10] == 0 && addr[11] == 0 && addr[12] == 0 && addr[13] == 0 && addr[14] == 0 && (addr[15] == 0 || addr[15] == 1)) {
					continue;
				}
				
				memcpy(this->peercode.localV6, addr, 16);
				memcpy(this->peercode.localP6, &myport, 2);
			} else {
				uint8_t *addr = (void*) &((struct sockaddr_in*) ifa->FirstUnicastAddress->Address.lpSockaddr)->sin_addr;
				if(addr[0] == 0 && addr[1] == 0 && addr[2] == 0 && addr[3] == 0 || addr[0] == 127 && addr[1] == 0 && addr[2] == 0 && addr[3] == 1) {
					continue;
				}
				
				memcpy(this->peercode.localV4, addr, 4);
				memcpy(this->peercode.localP4, &myport, 2);
			}
		}
	}
	free(addrs);
#else
	struct ifaddrs *ifaddr;
	if(getifaddrs(&ifaddr) >= 0) {
		for(struct ifaddrs *ifa = ifaddr; ifa; ifa = ifa->ifa_next) {
			if(ifa->ifa_addr->sa_family == serv->ai_family) {
				if(serv->ai_family == AF_INET6) {
					uint8_t *addr = (void*) &((struct sockaddr_in6*) ifa->ifa_addr)->sin6_addr;
					if(addr[0] == 0 && addr[1] == 0 && addr[2] == 0 && addr[3] == 0 && addr[4] == 0 && addr[5] == 0 && addr[6] == 0 && addr[7] == 0 && addr[8] == 0 && addr[9] == 0 && addr[10] == 0 && addr[11] == 0 && addr[12] == 0 && addr[13] == 0 && addr[14] == 0 && (addr[15] == 0 || addr[15] == 1)) {
						continue;
					}
					
					memcpy(this->peercode.localV6, addr, 16);
					memcpy(this->peercode.localP6, &myport, 2);
				} else {
					uint8_t *addr = (void*) &((struct sockaddr_in*) ifa->ifa_addr)->sin_addr;
					if(addr[0] == 0 && addr[1] == 0 && addr[2] == 0 && addr[3] == 0 || addr[0] == 127 && addr[1] == 0 && addr[2] == 0 && addr[3] == 1) {
						continue;
					}
					
					memcpy(this->peercode.localV4, addr, 4);
					memcpy(this->peercode.localP4, &myport, 2);
				}
				break;
			}
		}
		
		freeifaddrs(ifaddr);
	}
#endif
	
	return fd;
}
struct Stoon stoon_init(const char *stunhost, uint16_t stunport, uint16_t myport) {
	struct addrinfo *stunaddrs;
	struct addrinfo *v4 = NULL;
	struct addrinfo *v6 = NULL;
	
	{
		char p[6];
		sprintf(p, "%u", stunport);
		
		getaddrinfo(stunhost, p, &(struct addrinfo) {.ai_family = AF_UNSPEC, .ai_socktype = SOCK_DGRAM, .ai_protocol = IPPROTO_UDP}, &stunaddrs);
		
		for(struct addrinfo *addr = stunaddrs; addr; addr = addr->ai_next) {
			if(addr->ai_family == AF_INET6) {
				v6 = addr;
			} else {
				v4 = addr;
			}
		}
	}
	
	struct Stoon ret = {};
	
	if(v4) {
		memcpy(&ret.stu4, v4->ai_addr, v4->ai_addrlen);
		ret.fd4 = stoon_init_mini(&ret, v4, myport);
	} else ret.fd4 = -1;
	
	if(v6) {
		memcpy(&ret.stu6, v6->ai_addr, v6->ai_addrlen);
		ret.fd6 = stoon_init_mini(&ret, v6, myport);
	} else ret.fd6 = -1;
	
	ret.poonchstage = 0;
	
	freeaddrinfo(stunaddrs);
	
	return ret;
}

static int stoon_req_mini(struct Stoon *this, int *fd, struct sockaddr *serv, int servlen) {
	struct StunMsg req = {.type = htons(STUN_BINDING_REQUEST), .len = htons(0), .magic = htonl(STUN_MAGIC)};
	RAND(req.id, sizeof(req.id));
	
	if(sendto(*fd, (void*) &req, sizeof(req), 0, serv, servlen) == -1) {
		if(errno == ENETUNREACH) {
			close(*fd);
			*fd = -1;
			return false;
		}
	}
	
	return true;
}
int stoon_req(struct Stoon *this) {
	if(this->fd4 >= 0) stoon_req_mini(this, &this->fd4, (struct sockaddr*) &this->stu4, sizeof(struct sockaddr_in));
	if(this->fd6 >= 0) stoon_req_mini(this, &this->fd6, (struct sockaddr*) &this->stu6, sizeof(struct sockaddr_in6));
	
	if(this->fd4 < 0 && this->fd6 < 0) {
		return false;
	}
	
	return true;
}

static int stoon_listen_mini(struct Stoon *this, int fd, struct sockaddr *serv, int servlen) {
	struct {
		struct StunMsg base;
		uint8_t attribs[MAX_ATTRIB_BUFFER_SIZE];
	} res;
	
	if(recvfrom(fd, (void*) &res, sizeof(res), 0, NULL, 0) < 0) {
		return -1;
	}
	
	if(ntohs(res.base.type) != STUN_BINDING_RESPONSE) {
		puts("Not binding response");
		return -3;
	}
	
	uint16_t attribsLen = ntohs(res.base.len);
	if(attribsLen > MAX_ATTRIB_BUFFER_SIZE) attribsLen = MAX_ATTRIB_BUFFER_SIZE;
	
	for(uint8_t *d = res.attribs;;) {
		// Overflow check
		if((uintmax_t) (d - res.attribs) > MAX_ATTRIB_BUFFER_SIZE - 4) break;
		
		uint16_t attribType = ntohs(((uint16_t*) d)[0]);
		uint16_t attribLen = ntohs(((uint16_t*) d)[1]);
		
		// Overflow check
		if((uintmax_t) (d + 4 + attribLen - res.attribs) > MAX_ATTRIB_BUFFER_SIZE) break;
		
		if(attribType == STUN_XOR_MAPPED_ADDRESS) {
			uint16_t netfam = ntohs(((uint16_t*) d)[2]);
			uint16_t publicPort = ntohs(((uint16_t*) d)[3]) ^ (STUN_MAGIC >> 16);
			if(netfam == STUN_NETFAM_IPV4) { 
				uint32_t publicIp = htonl(ntohl(((uint32_t*) d)[2]) ^ STUN_MAGIC);
				
				memcpy(this->peercode.publicV4, &publicIp, 4);
				memcpy(this->peercode.publicP4, &publicPort, 2);
				
				return 1;
			} else if(netfam == STUN_NETFAM_IPV6) {
				uint32_t publicIp[4];
				publicIp[0] = ((uint32_t*) d)[2] ^ htonl(STUN_MAGIC);
				publicIp[1] = ((uint32_t*) d)[3] ^ ((uint32_t*) res.base.id)[0];
				publicIp[2] = ((uint32_t*) d)[4] ^ ((uint32_t*) res.base.id)[1];
				publicIp[3] = ((uint32_t*) d)[5] ^ ((uint32_t*) res.base.id)[2];
				
				memcpy(this->peercode.publicV6, &publicIp, 16);
				memcpy(this->peercode.publicP6, &publicPort, 2);
				
				return 1;
			} else {
				puts("Neither ipv4 nor ipv6???");
			}
		}
		
		d += 4 + attribLen;
	}
	
	return false;
}
int stoon_listen(struct Stoon *this) {
	if(this->fd4 >= 0) if(stoon_listen_mini(this, this->fd4, (struct sockaddr*) &this->stu4, sizeof(struct sockaddr_in)) < 0) return 0;
	if(this->fd6 >= 0) if(stoon_listen_mini(this, this->fd6, (struct sockaddr*) &this->stu6, sizeof(struct sockaddr_in6)) < 0) return 0;
	return true;
}

void stoon_keepalive(struct Stoon *this) {
	struct StunMsg req = {.type = htons(STUN_BINDING_INDICATION), .len = htons(0), .magic = htonl(STUN_MAGIC)};
	RAND(&req.id, sizeof(req.id));
	
	if(this->fd4 >= 0) {
		sendto(this->fd4, (void*) &req, sizeof(req), 0, (struct sockaddr*) &this->stu4, sizeof(struct sockaddr_in));
	}
	
	if(this->fd6 >= 0) {
		sendto(this->fd6, (void*) &req, sizeof(req), 0, (struct sockaddr*) &this->stu6, sizeof(struct sockaddr_in6));
	}
}

void stoon_kill(struct Stoon *this) {
	if(this->fd4 >= 0) close(this->fd4);
	this->fd4 = -1;
	
	if(this->fd6 >= 0) close(this->fd6);
	this->fd6 = -1;
}

/*int stoon_poonch(struct Stoon *this, const uint8_t peerdata[static STOON_CONN_INFO_SIZE]) {
	struct {
		uint32_t addr4;
		uint16_t port4;
		
		uint8_t addr6[16];
		uint16_t port6;
	} peer;
	
	memcpy(&peer.addr4, peerdata + 0, 4);
	memcpy(&peer.port4, peerdata + 4, 2);
	
	memcpy(&peer.addr6, peerdata + 6, 16);
	memcpy(&peer.port6, peerdata + 22, 2);
	
	int myfd = -1;
	struct sockaddr_storage peeraddr = {};
	int peersz;
	
	// Prioritize IPv6
	if(this->my6.sin6_family && peer.port6) {
		myfd = this->fd6;
		
		struct sockaddr_in6 *pee = (void*) &peeraddr;
		pee->sin6_family = AF_INET6;
		memcpy(&pee->sin6_addr, &peer.addr6, 16);
		pee->sin6_port = peer.port6;
		peersz = sizeof(*pee);
	} else if(this->my4.sin_family && peer.port4) {
		myfd = this->fd4;
		
		struct sockaddr_in *pee = (void*) &peeraddr;
		pee->sin_family = AF_INET;
		memcpy(&pee->sin_addr, &peer.addr4, 4);
		pee->sin_port = peer.port4;
		peersz = sizeof(*pee);
	} else {
		return STOON_POONCH_INCOMPATIBLE_NETFAMS;
	}
	
	if(this->poonchstage == POONCH_STAGE_KNOCKING) {
		char buf[64];
		inet_ntop(peeraddr.ss_family, peeraddr.ss_family == AF_INET ? peerdata + 0 : peerdata + 6, buf, sizeof(buf));
		printf("Sending knock to %s:%u\n", buf, peeraddr.ss_family == AF_INET ? ntohs(peer.port4) : ntohs(peer.port6));
		struct Poonch cmd = {.magic = POONCH_MAGIC, .stage = POONCH_STAGE_KNOCKING};
		sendto(myfd, (void*) &cmd, sizeof(cmd), 0, (struct sockaddr*) &peeraddr, peersz);
	}
	
	struct Poonch comin;
	if(recvfrom(myfd, (void*) &comin, sizeof(comin), 0, NULL, NULL) == sizeof(comin) && comin.magic == POONCH_MAGIC) {
		puts("received somethign");
		if(comin.stage == POONCH_STAGE_KNOCKING) {
			puts("Received knock");
			struct Poonch ret = {.magic = POONCH_MAGIC, .stage = POONCH_STAGE_ACK};
			this->poonchstage = POONCH_STAGE_ACK;
			sendto(myfd, (void*) &ret, sizeof(ret), 0, (struct sockaddr*) &peeraddr, peersz);
			
			return STOON_POONCH_ACKED;
		}
	}
	
	return STOON_POONCH_NO_KNOCK;
}*/

#ifdef STOON_STANDALONE
int main() {
#ifdef _WIN32
	WSADATA wsadata;
	WSAStartup(MAKEWORD(1, 1), &wsadata);
#endif
	
	struct Stoon stoon = stoon_init("stun.l.google.com", 19305, 26656);
	
	stoon_req(&stoon);
	
	usleep(1000000);
	stoon_listen(&stoon);
	
	stoon_keepalive(&stoon);
	
	char buf[64];
	
	fputs("IPv4: ", stdout);
	if(stoon.my4.sin_family) {
		inet_ntop(AF_INET, &stoon.my4.sin_addr, buf, sizeof(buf));
		printf("%s:%u\n", buf, ntohs(stoon.my4.sin_port));
	} else {
		puts("unavailable");
	}
	
	fputs("IPv6: ", stdout);
	if(stoon.my6.sin6_family) {
		inet_ntop(AF_INET6, &stoon.my6.sin6_addr, buf, sizeof(buf));
		printf("[%s]:%u\n", buf, ntohs(stoon.my6.sin6_port));
	} else {
		puts("unavailable");
	}
	
	uint8_t r[STOON_CONN_INFO_SIZE];
	stoon_serialize(&stoon, r);
	
	fputs("Serialized: ", stdout);
	for(int i = 0; i < STOON_CONN_INFO_SIZE; i++) {
		printf("%02x", r[i]);
	}
	fputc(10, stdout);
	
	puts("Enter peer's connection string:");
	
	char peerdatareadable[128];
	fgets(peerdatareadable, sizeof(peerdatareadable), stdin);
	
	uint8_t peerdata[STOON_CONN_INFO_SIZE] = {};
	for(int i = 0;; i += 2) {
		if(peerdatareadable[i] == 0) break;
		
		int b = 0;
		
		int c = peerdatareadable[i];
		if(c >= '0' && c <= '9') {
			b |= (c - '0') << 4;
		} else if(c >= 'a' && c <= 'f') {
			b |= (c - 'a' + 10) << 4;
		}
		
		c = peerdatareadable[i + 1];
		if(c >= '0' && c <= '9') {
			b |= (c - '0');
		} else if(c >= 'a' && c <= 'f') {
			b |= (c - 'a' + 10);
		}
		
		peerdata[i / 2] = b;
	}
	
	int got = 0;
	while(1) {
		if(stoon_poonch(&stoon, peerdata) == STOON_POONCH_NO_KNOCK) {
			got++;
			if(got >= 10 && stoon.poonchstage == POONCH_STAGE_ACK) {
				puts("Gape successful!");
				break;
			}
		} else {
			got = 0;
		}
		usleep(500000);
	}
	
	stoon_kill(&stoon);
}
#endif