/************************************/
/* DHCP Client                      */
/* 2006/9/1 first release           */
/*                                  */
/************************************/

#include <macro.h>
#include <mes2.h>
#include "dhcp.h"

unsigned int DhcpServerIP;	///< IP address of the DHCP server that offered lease
unsigned int DhcpClientIP;
unsigned int DhcpNetmask;
unsigned int DhcpGatewayIP;
unsigned int DhcpDNSServerIP;
unsigned int DhcpTransactID;    ///< Unique transaction ID that identifies DHCP request/replies
unsigned int DhcpLeaseTime;     ///< Number of seconds left in DHCP lease
unsigned char macaddr[6];

#undef NET_DEBUG

int dhcpIn(char *dev, int sock, int len, struct netDhcpHeader *packet) {
	unsigned char	msgtype;
	unsigned char	*optptr;
	unsigned int	val;

#if NET_DEBUG >= 3
	dhcpPrintHeader(packet);
#endif
	// check that this is a reply, and for me
	if(packet->bootp.op != BOOTP_OP_BOOTREPLY ||
	   packet->bootp.htype != BOOTP_HTYPE_ETHERNET ||
	   packet->bootp.hlen != BOOTP_HLEN_ETHERNET ||
	   be32toh(packet->cookie) != 0x63825363 ||
	   be32toh(packet->bootp.xid) != DhcpTransactID) return -1;

	// process incoming packet
	// check reply type
	dhcpGetOption(packet->options, DHCP_OPT_DHCPMSGTYPE, 1, &msgtype);
#if NET_DEBUG >= 2
	printf("DHCP: Received msgtype = %d\r\n", msgtype);
#endif

	if(msgtype == DHCP_MSG_DHCPOFFER) {
		// get DHCP server ID
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_SERVERID, 4, &val);
		if(val == 0) {
			val = packet->bootp.siaddr;
		}
		DhcpServerIP = be32toh(val);
		DhcpGatewayIP = DhcpServerIP;
		DhcpDNSServerIP = DhcpServerIP;
		DhcpClientIP = be32toh(packet->bootp.yiaddr);
#ifdef NET_DEBUG
		printf("DHCP: Got offer from server ");
		val = DhcpServerIP;
		printf("%d.%d.%d.%d\r", (val >> 24) & 0xff, (val >> 16) & 0xff, (val >> 8) & 0xff, val & 0xff); 
		printf("DHCP: Sending request in response to offer\r");
#endif
		// get netmask
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_NETMASK, 4, &val);
		DhcpNetmask = htobe32(val);
		if(DhcpNetmask == 0) DhcpNetmask = 0xffffff00;
		// get gateway
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_ROUTERS, 4, &val);
		if(val != 0) DhcpGatewayIP = htobe32(val);
		// get gateway
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_DNSSERVERS, 4, &val);
		if(val != 0) DhcpDNSServerIP = htobe32(val);
	} else if(msgtype == DHCP_MSG_DHCPACK) {
		// get netmask
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_NETMASK, 4, &val);
		DhcpNetmask = htobe32(val);
		if(DhcpNetmask == 0) DhcpNetmask = 0xffffff00;
		// get gateway
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_ROUTERS, 4, &val);
		if(val != 0) DhcpGatewayIP = htobe32(val);
		// get gateway
		val = 0;
		dhcpGetOption(packet->options, DHCP_OPT_DNSSERVERS, 4, &val);
		if(val != 0) DhcpDNSServerIP = htobe32(val);
		// get lease time
		dhcpGetOption(packet->options, DHCP_OPT_LEASETIME, 4, &val);
		DhcpLeaseTime = htobe32(val);
		// assign new network info
		ifconfig(dev, be32toh(packet->bootp.yiaddr), DhcpNetmask);
#ifdef NET_DEBUG
		// print netmask
		val = DhcpNetmask;
		printf("Mask : %d.%d.%d.%d\r", (val >> 24) & 0xff, (val >> 16) & 0xff, (val >> 8) & 0xff, val & 0xff); 
		// print gateway
		val = DhcpGatewayIP;
		printf("GW : %d.%d.%d.%d\r", (val >> 24) & 0xff, (val >> 16) & 0xff, (val >> 8) & 0xff, val & 0xff); 
		// print dns server
		val = DhcpDNSServerIP;
		printf("DNS : %d.%d.%d.%d\r", (val >> 24) & 0xff, (val >> 16) & 0xff, (val >> 8) & 0xff, val & 0xff); 
		printf("DHCP: Got request ACK, bind complete\r");
#endif
	} else {
		return -1;	
	}
	return msgtype;
}

int dhcpRequest(char *dev, int sock) {
	struct sockaddr		target;
	struct netDhcpHeader*	packet;
	unsigned char		*options;
	unsigned int		ip, mask, val;
	int			net, ret, count;
	unsigned char		buf[8];

	// get interface ip address
	net = get_netnum(dev);
	if(net == -1) return -1;
	ret = get_ipinfo(net, &ip, &mask, macaddr);
	if(ret == -1) return -1;

	// get buffer area
	packet = (struct netDhcpHeader*)malloc(DHCP_HEADER_LEN + 512);
	if(packet == 0) return -1;
	options = packet->options;
	bzero((char*)packet, DHCP_HEADER_LEN + 512);

	// build BOOTP/DHCP header
	packet->bootp.op = BOOTP_OP_BOOTREQUEST;        // request type
	packet->bootp.htype = BOOTP_HTYPE_ETHERNET;
	packet->bootp.hlen = BOOTP_HLEN_ETHERNET;
	packet->bootp.hops = 0;
	packet->bootp.xid = htobe32(DhcpTransactID);
	packet->bootp.secs = 0;
	packet->bootp.ciaddr = htobe32(ip);
	packet->bootp.flags = (ip == 0) ? 0 : htobe16(0x8000);
	packet->bootp.yiaddr = 0;
	packet->bootp.siaddr = 0;
	packet->bootp.giaddr = 0;
	memcpy(&packet->bootp.chaddr[0], macaddr, 6); // fill client hardware address
	// build DHCP request
	// begin with magic cookie
	packet->cookie = htobe32(0x63825363);
	// set operation
	count = 0;
	if(ip != 0) {
		buf[0] = DHCP_MSG_DHCPRELEASE;
	} else if(DhcpServerIP == 0) {
		buf[0] = DHCP_MSG_DHCPDISCOVER;
	} else {
		val = htobe32(DhcpServerIP);
		count += 2 + 4;
		options = dhcpSetOption(options, DHCP_OPT_SERVERID, 4, &val);
		if(DhcpClientIP != 0) {
			val = htobe32(DhcpClientIP);
			count += 2 + 4;
			options = dhcpSetOption(options, DHCP_OPT_REQUESTEDIP, 4, &val);
		}
		buf[0] = DHCP_MSG_DHCPREQUEST;
	}
	count += 2 + 1;
	options = dhcpSetOption(options, DHCP_OPT_DHCPMSGTYPE, 1, buf);
	buf[0] = BOOTP_HTYPE_ETHERNET;
	memcpy(&buf[1], macaddr, 6);
	count += 2 + 7;
	options = dhcpSetOption(options, DHCP_OPT_CLIENT_IDENT, 7, buf);
	if(ip == 0) {
		buf[0] = DHCP_OPT_NETMASK;
		buf[1] = DHCP_OPT_ROUTERS;
		buf[2] = DHCP_OPT_DNSSERVERS;
		buf[3] = DHCP_OPT_DOMAINNAME;
		count += 2 + 4;
		options = dhcpSetOption(options, DHCP_OPT_PARAMREQLIST, 4, buf);
	}
	count++;
	options = dhcpSetOption(options, DHCP_OPT_END, 0, 0);
#ifdef NET_DEBUG
	printf("DHCP: Sending Query\r");
#endif
	// send request
	target.sin_port = DHCP_UDP_SERVER_PORT;
	target.sin_addr = 0xffffffff;
	if(sendto(sock, (unsigned char*)packet, DHCP_HEADER_LEN + count, &target) == -1) {
		free((char*)packet);
		return -1;
	}
	return 0;
}

void dhcpTimer(void) {
	// this function to be called once per second

	// decrement lease time
	if(DhcpLeaseTime) DhcpLeaseTime--;
}

unsigned char dhcpGetOption(unsigned char* options, unsigned char optcode, unsigned char optlen, void* optvalptr) {
	unsigned char i;

	// parse for desired option
	for (;;) {
		// skip pad characters
		if(*options == DHCP_OPT_PAD) options++;
		// break if end reached
		else if(*options == DHCP_OPT_END) break;
		// check for desired option
		else if(*options == optcode) {
			// found desired option
			// limit size to actual option length
			optlen = (optlen < *(options+1)) ? optlen : *(options+1);
			//if(*(options+1) < optlen)
			//  optlen = *(options+1);
			// copy contents of option
			for(i = 0; i < optlen; i++)
				*(((unsigned char*)optvalptr)+i) = *(options+i+2);
			// return length of option
			return *(options+1);
		} else {
			// skip to next option
			options++;
			options+=*options;
			options++;
		}
	}
	// failed to find desired option
	return 0;
}

unsigned char* dhcpSetOption(unsigned char* options, unsigned char optcode, unsigned char optlen, void* optvalptr) {
	// use current options address as write point
	// set optcode
	*options++ = optcode;
	// set optlen
	*options++ = optlen;
	// copy in argument/data
	while(optlen--) {
		*options++ = *(unsigned char*)optvalptr++;
	}
	// write end marker
	*options = DHCP_OPT_END;

	// return address of end marker, to be used as a future write point
	return options;
}

#ifdef NET_DEBUG
void dhcpPrintHeader(struct netDhcpHeader* packet) {
	int	ip;

	printf("DHCP Packet:\r");
	// print op
	printf("Op      : ");
	switch(packet->bootp.op) {
		case BOOTP_OP_BOOTREQUEST: printf("BOOTREQUEST\r"); break;
		case BOOTP_OP_BOOTREPLY: printf("BOOTREPLY\r"); break;
		default: printf("UNKNOWN\r"); break;
	}
	// print transaction ID
	printf("XID     : 0x%08x\r", packet->bootp.xid);

	// print client IP address
	ip = be32toh(packet->bootp.ciaddr);
	printf("%d.%d.%d.%d\r", (ip >> 24) & 0xff, (ip >> 16) & 0xff, (ip >> 8) & 0xff, ip & 0xff); 
	// print 'your' IP address
	ip = be32toh(packet->bootp.yiaddr);
	printf("%d.%d.%d.%d\r", (ip >> 24) & 0xff, (ip >> 16) & 0xff, (ip >> 8) & 0xff, ip & 0xff); 
	// print server IP address
	ip = be32toh(packet->bootp.siaddr);
	printf("%d.%d.%d.%d\r", (ip >> 24) & 0xff, (ip >> 16) & 0xff, (ip >> 8) & 0xff, ip & 0xff); 
	// print gateway IP address
	ip = be32toh(packet->bootp.giaddr);
	printf("%d.%d.%d.%d\r", (ip >> 24) & 0xff, (ip >> 16) & 0xff, (ip >> 8) & 0xff, ip & 0xff); 
	// print client hardware address
	printf("%s\r", packet->bootp.chaddr);
}
#endif

int main(int argc, char **argv) {
	struct sockaddr	myaddr, target;
	int		net, sock, size, ret, retry, msgtype, retry2;
	unsigned int	ip, mask;
	char		dev[8], *buf;
	char		flag;

#ifdef NET_DEBUG
	printf("DHCP Client 2007/1/1 0:00 coded.\r");
#endif
	if(argc == 1 || argc >= 4) return -1;
	flag = 0;
	if(argc == 3 && argv[2][0] == 'd') flag = 1;
	strcpy(dev, argv[1]);
	buf = malloc(1024);
	if(buf == 0) return -1;
	for(retry = 4;retry >= 0;retry--) {
#ifdef NET_DEBUG
		printf("-----Retry %d-----\r", 3 - retry);
#endif
		net = get_netnum(dev);
		if(net == -1) {
#ifdef NET_DEBUG
			printf("NIC is disabled\r");
#endif
		} else {
			ret = get_ipinfo(net, &ip, &mask, macaddr);
			if(ip != 0) {
				// set transaction ID based on mac address
				DhcpTransactID = *((unsigned int*)&macaddr) - retry;
				sock = udp_socket();
				if(sock == -1) continue;
				myaddr.sin_addr = 0;
				myaddr.sin_port = DHCP_UDP_CLIENT_PORT;
				udp_bind(sock, &myaddr);
				dhcpRequest(dev, sock);
				udp_free(sock);
			}
		}
		net = ifconfig(dev, 0, 0);
		if(flag == 1) return 0;
		if(net == -1) {
#ifdef NET_DEBUG
			printf("Fail(1)\r");
#endif
			continue;
		}
		// get interface mac address
		ret = get_ipinfo(net, &ip, &mask, macaddr);
		if(ret == -1) {
#ifdef NET_DEBUG
			printf("Fail(2)\r");
#endif
			continue;
		}
		// set transaction ID based on mac address
		DhcpTransactID = *((unsigned int*)&macaddr) - retry;
		// reset lease time
		DhcpLeaseTime = 0;

		// socket
		sock = udp_socket();
		if(sock == -1) {
#ifdef NET_DEBUG
			printf("Fail(3)\r");
#endif
			continue;
		}
		myaddr.sin_addr = 0;
		myaddr.sin_port = DHCP_UDP_CLIENT_PORT;
		udp_timeout(sock, 4000);
		if(udp_bind(sock, &myaddr) == -1) {
			udp_free(sock);
#ifdef NET_DEBUG
			printf("Fail(4)\r");
#endif
			continue;
		}
		DhcpServerIP = DhcpClientIP = 0;

		/***********************/
		/**** SEND DISCOVER ****/
		/***********************/

		if(dhcpRequest(dev, sock) == -1) {
			udp_free(sock);
#ifdef NET_DEBUG
			printf("Fail(5)\r");
#endif
			continue;
		}
		for(retry2 = 0;retry2 < 4;) {
			bzero(buf, 1024);
			size = recvfrom(sock, buf, 1024, &target);
			if(size <= 0) {
				retry2++;
				dhcpRequest(dev, sock);
			} else {
				msgtype = dhcpIn(dev, sock, size, (struct netDhcpHeader *)buf);
				if(msgtype != DHCP_MSG_DHCPACK) break;
			}
		}
		if(size <= 0) {
			udp_free(sock);
#ifdef NET_DEBUG
			printf("Fail(6)\r");
#endif
			continue;
		}
		if(msgtype != DHCP_MSG_DHCPOFFER) {
			udp_free(sock);
#ifdef NET_DEBUG
			printf("Fail(7)\r");
#endif
			continue;
		}

		/**********************/
		/**** SEND REQUEST ****/
		/**********************/

		if(dhcpRequest(dev, sock) == -1) {
			udp_free(sock);
#ifdef NET_DEBUG
			printf("Fail(8)\r");
#endif
			continue;
		}

		for(retry2 = 0;retry2 < 4;) {
			bzero(buf, 1024);
			size = recvfrom(sock, buf, 1024, &target);
			if(size <= 0) {
				retry2++;
				dhcpRequest(dev, sock);
			} else {
				msgtype = dhcpIn(dev, sock, size, (struct netDhcpHeader *)buf);
				if(msgtype != DHCP_MSG_DHCPOFFER) break;
			}
		}
		if(size <= 0) {
#ifdef NET_DEBUG
			printf("Fail(9)\r");
			printf("DHCP-ACK not received, but set ip.\r");
#endif
			// assign new network info
			ifconfig(dev, DhcpClientIP, DhcpNetmask);
		} else if(msgtype != DHCP_MSG_DHCPACK) {
#ifdef NET_DEBUG
			printf("Fail(10)\r");
			printf("DHCP-ACK not received, but set ip.\r");
#endif
			// assign new network info
			ifconfig(dev, DhcpClientIP, DhcpNetmask);
		}
		udp_free(sock);
		set_gateway(DhcpGatewayIP);
		dns_server(DhcpDNSServerIP);
		free(buf);
		return 0;
	}
	free(buf);
	return retry;
}