#include <stdio.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <netpacket/packet.h>
#include <net/ethernet.h>       /* the L2 protocols */
#include <net/if.h>
#include <sys/time.h>
#include <getopt.h>
#include <stdlib.h>
#include <string.h>
#include <netinet/in.h>
#include <errno.h>
#include "upslug.h"

static void putdigit(int i) {
	const char hexdigit[16] = "0123456789abcdef";
	putchar(hexdigit[i&0xf]);
	fflush(stdout);
}

static unsigned short nextseq(unsigned short *sequence) {
	return ++*sequence & 0xffff;
}

void command_print(int rc, const DLC *pdlc) {
	int j;

	fprintf(stderr, "\nListen: received %d bytes:"
			" cmd(0x%x) seq(0x%x) index(0x(%x+%x)) len(%d):\n",
			rc, pdlc->wcmd, pdlc->wsequence, pdlc->wsegment*16, pdlc->woffset, pdlc->wLen);

	for (j=0; j<rc-DLC_LEN; j+=8) {
		fprintf(stderr, " %02x%02x%02x%02x%02x%02x%02x%02x",
				pdlc->bData[j+0], pdlc->bData[j+1], pdlc->bData[j+2], pdlc->bData[j+3],
				pdlc->bData[j+4], pdlc->bData[j+5], pdlc->bData[j+6], pdlc->bData[j+7]);
		if ((j & 0x1f) == 0x18 || j+8 >= rc-DLC_LEN)
			fprintf(stderr, "\n");
	}
}

int command_listen(int p, const struct sockaddr_ll slug_addr, DLC *pdlc, unsigned short cmd,
		int egain_ok) {
	/* Listen for a response. */
	int rc;

	do {
		memset(pdlc, 0, sizeof *pdlc);
		rc = recv(p, pdlc, sizeof *pdlc, 0);
		if (rc < 0) {
			if (errno == EAGAIN) {
				if (egain_ok) {
					putchar('*');
					fflush(stdout);
					break;
				}

				fprintf(stderr, "\nSlug does not answer (EGAIN)\n");
				exit(1);
			} else {
				perror("\nupslug: recv(comand listen)");
				exit(1);
			}
		} else {

			/* Check for the success return. */
			if (rc > DLC_LEN && rc > DLC_LEN+pdlc->wLen && pdlc->wcmd == cmd)
				break;

			/* This is an unexpected command response or a zero byte response. */
			command_print(rc, pdlc);
		}
	} while (1);

	return rc;
}

void down_request(unsigned short *sequence, int p, const struct sockaddr_ll slug_addr) {
	int rc;
	DLC dlc;

	memset(&dlc, 0, sizeof dlc);
	dlc.wcmd = DOWN_REQUEST;
	dlc.wsequence = nextseq(sequence);

	/* Send away ! */
	printf("Requesting upgrade...");
	rc = sendto(p, &dlc, DLC_LEN, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);

	if (rc < 0) {
		perror("upslug: sendto(upgrade)");
		exit(1);
	}

	rc = command_listen(p, slug_addr, &dlc, DOWN_REQUEST, 0);

	printf(" confirmed\n");
	printf("Erasing Flash...");
	fflush(stdout);

	/* Listen for the end of the flash erase. */
	rc = command_listen(p, slug_addr, &dlc, DOWN_REQUEST, 1);
	if (rc > 0)
		printf(" done\n");
}

void eall(unsigned short *sequence, int p, const struct sockaddr_ll slug_addr) {
	int rc;
	DLC dlc;
	
	memset(&dlc, 0, sizeof dlc);
	dlc.wcmd = DOWN_EALL;
	dlc.wsequence = nextseq(sequence);

	/* Send away ! */
	printf("Requesting upgrade (complete reflash)...");
	rc = sendto(p, &dlc, DLC_LEN, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);

	if (rc == -1) {
		perror("upslug: sendto(full upgrade)");
		exit(1);
	}

	rc = command_listen(p, slug_addr, &dlc, DOWN_EALL, 0);

	printf(" confirmed\n");
	printf("Erasing Flash...");
	fflush(stdout);

	/* Listen for the end of the flash erase. */
	rc = command_listen(p, slug_addr, &dlc, DOWN_EALL, 1);
	if (rc > 0)
		printf(" done\n");
}

void verify_file_contents(char * upload_file) {
	FILE * fp;
	unsigned char binary_block[512];
	int counted_blocks = 0;

	fp = fopen(upload_file, "rb");
	if (fp == NULL) {
		fprintf(stderr, "%s: %s: could not open file to verify contents\n",
			upload_file, strerror(errno));
		exit(1);
	}

	/*Verify File Length*/
	while(fread(binary_block, sizeof binary_block, 1, fp) == 1)
	 	counted_blocks++;

	if (counted_blocks != NUM_IMAGE_BLOCKS) {
		printf("Counted Blocks:%i\n", counted_blocks);
		printf("Needed Blocks:%i\n", NUM_IMAGE_BLOCKS);
		printf("Error: File not large enough, use -f to override\n");
		exit(1);
	}

	/*Look for Sercomm signature*/
	if (binary_block[512-8] != 0x65 || binary_block[512-7] != 0x52 ||
			binary_block[512-6] != 0x63 || binary_block[512-5] != 0x4f ||
			binary_block[512-4] != 0x6d || binary_block[512-3] != 0x4d) {
		printf("Error: eRcOmM signature does not exist in image (will not work with stock redboot) -f to override\n");
		exit(1);
	}

	fclose(fp);
}

void write_image(unsigned short *sequence, int p, const struct sockaddr_ll slug_addr,
		char *upload_file) {
	unsigned short i;
	FILE *fp;

	fp = fopen(upload_file, "rb");
	if(fp == NULL) {
		fprintf(stderr, "%s: %s: could not open file to write\n",
			upload_file, strerror(errno));
		exit(1);
	}

	/*Send out the data to be written to the flash*/
	for(i=0; i<NUM_IMAGE_BLOCKS; i++) {
		int rc;
		DLC dlcv;
		unsigned char binary_block[512];

		if (fread(binary_block, sizeof binary_block, 1, fp) != 1) {
			printf("\nUnexpected File Read Error\n");
			break;
		}

		memset(&dlcv, 0, sizeof dlcv);
		dlcv.wcmd = DOWN_DATA;
		dlcv.wsequence = nextseq(sequence);
		dlcv.woffset = 0;
		dlcv.wsegment = i*32;
		dlcv.wLen = sizeof binary_block;
		
		memcpy(dlcv.bData, binary_block, sizeof binary_block);

		do {
			DLC dlc;

			rc = sendto(p, &dlcv, DLC_LEN+dlcv.wLen, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);
	
			if (rc < 0) {
				perror("upslug: sendto(write)");
				exit(1);
			}

			/* At the start of a write we get EAGAIN until the slug responds to the first
			 * write, after that we may get more EAGAIN because writing takes some time.
			 * Handle this by resending the request but only after the first one succeeds.
			 */
			do {
				rc = command_listen(p, slug_addr, &dlc, DOWN_DATA, 1);
			} while (rc < 0 && i == 0);
			if (rc < 0)
				continue;

			if (i == 0) {
				printf(" finished\nWriting Flash...");
				fflush( stdout );
			}

			if (dlc.wLen != 2) {
				fprintf(stderr, "\nwrite: unexpected response (%d bytes, 0x%x, 0x%x)\n",
						dlc.wLen, dlc.bData[0], dlc.bData[1]);
				command_print(rc, &dlc);
				exit(1);
			} else if (dlc.bData[0] != 0) {
				fprintf(stderr, "\nwrite: failed(%d,%d) at byte 0x%x:\n"
						"\tsequence: 0x%x\n\toffset: 0x%x\n\tsegment: 0x%x\n",
						dlc.bData[0], dlc.bData[1], i*512,
						dlc.wsequence, dlc.woffset, dlc.wsegment);
				command_print(rc, &dlc);
				exit(1);
			} else if (dlc.woffset != dlcv.woffset || dlc.wsegment != dlcv.wsegment) {
				fprintf(stderr, "\nwrite: out of sequence expected 0x(%x+%x), got 0x(%x+%x)\n",
						dlcv.wsegment*16, dlcv.woffset, dlc.wsegment*16, dlc.woffset);
				command_print(rc, &dlc);
				exit(1);
			} else if (dlc.wsequence != (0xffff & *sequence)) {
				fprintf(stderr, "\nwrite: mismatched sequence (gave %x, got %x)\n",
						*sequence, dlc.wsequence);
				command_print(rc, &dlc);
				exit(1);
			}
		} while (rc < 0);

		if ((i & 0xff) == 0xff)
			putdigit(i >> 8);
	}

	printf(" completed\n");
	fclose(fp);
}

void verify_image(unsigned short *sequence, int p, const struct sockaddr_ll slug_addr) {
	unsigned short i;
	
	printf("Verify  Flash...");

	/*Verify that the image was written properly*/	
	for (i=0; i<NUM_IMAGE_BLOCKS; i++) {
		int rc;
		DLC dlcv;

		memset(&dlcv, 0, sizeof dlcv);
		dlcv.wcmd = DOWN_VERIFY;
		dlcv.wsequence = nextseq(sequence);
		dlcv.woffset = 0x0000;
		dlcv.wsegment = i*32;
		dlcv.wLen = 0;
		
		do {
			DLC dlc;

			rc = sendto(p, &dlcv, DLC_LEN+dlcv.wLen, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);
				
			if (rc < 0) {
				perror("upslug: verify: sendto");
				exit(1);
			}

			/* Listen for a response, if we get EGAIN retransmit the request. */
			rc = command_listen(p, slug_addr, &dlc, DOWN_VERIFY, 1);
			if (rc < 0)
				/*retransmit*/;
			else if (dlc.wLen != 2) {
				fprintf(stderr, "\nverify: unexpected response (%d bytes, 0x%x, 0x%x)\n",
						dlc.wLen, dlc.bData[0], dlc.bData[1]);
				command_print(rc, &dlc);
				exit(1);
			} else if (dlc.bData[0] != 0 || dlc.bData[1] != 0) {
				fprintf(stderr, "\nverify: failed(%d,%d) at byte 0x%x:\n"
						"\tsequence: 0x%x\n\toffset: 0x%x\n\tsegment: 0x%x\n",
						dlc.bData[0], dlc.bData[1], i*512,
						dlc.wsequence, dlc.woffset, dlc.wsegment);
				command_print(rc, &dlc);
				exit(1);
			} else if (dlc.woffset != dlcv.woffset || dlc.wsegment != dlcv.wsegment) {
				fprintf(stderr, "\nverify: out of sequence expected 0x(%x+%x), got 0x(%x+%x)\n",
						dlcv.wsegment*16, dlcv.woffset, dlc.wsegment*16, dlc.woffset);
				command_print(rc, &dlc);
				exit(1);
			} else if (dlc.wsequence != (0xffff & *sequence)) {
				fprintf(stderr, "\nverify: mismatched sequence (gave %x, got %x)\n",
						*sequence, dlc.wsequence);
				command_print(rc, &dlc);
				exit(1);
			}
		} while (rc < 0);

		if ((i & 0xff) == 0xff)
			putdigit(i >> 8);

	}
	printf(" completed\n");
}

void reset_slug(int p, const struct sockaddr_ll slug_addr) {
	int rc;
	DLC dlc;
	
	memset(&dlc, 0, sizeof dlc);
	dlc.wcmd = DOWN_RESET;

	/* Send away ! */
	printf("Requesting Reset...");
	rc = sendto(p, &dlc, DLC_LEN, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);

	if (rc < 0) {
		perror("upslug: sendto(reset)");
		exit(1);
	}

	rc = command_listen(p, slug_addr, &dlc, DOWN_RESET, 0);
	if (rc <= 0) {
		fprintf(stderr, "DOWN_RESET: no response\n");
		exit(1);
	}

	printf(" confirmed\n");
}

void showhelp(void) {
	printf(
	"Options:\n"
	"-i [interface]  interface is set to the interface specified. Defaults to eth0.\n"
	"-s [slug_id]    will communicate with the slug specified instead of letting the user choose\n"
	"-u [image file] removes sercomm signature, reboots slug, and then loads the firmware specified.\n"
	"-U [image file] same as -u but erases the entire contents of the flash (including redboot).\n"
	"-v              verify the written flash contents.\n"
	/*"-w [file]     writes from the file specified.\n"
	"-b [size]       the size of the block of data to write, defaults to 512.\n"*/
	"-r              resets the slug.\n"
	"-f              does not give warning messages.\n");
	/*"-e            writes zeros to the address specified for block size.\n\n\n");*/
}

int main(int argc, char **argv) {

	unsigned short sequence = 0x720;
	struct sockaddr_ll to_addr[5], slug_addr;
	struct ifreq if_eth;
	struct timeval timeout;
	DLC dlc;
	int i, j, rc, p, ifindex, slug_num = -1;
	char *eth_name = "eth0";
	char *slug_id = 0;
	char *upload_file = 0;
	char yes_no;
	int c;
	int upload = -1, erase_all = -1, write = -1, force = -1, reset = -1, erase = -1, verify = -1;
	/*Parse the options*/

	while((c = getopt(argc, argv, "i:s:u:U:w:a:b:ervfh?H")) != EOF)
	{
		switch(c)
		{
			case 'i':
				eth_name = optarg;
				break;
			case 's':
				slug_id = optarg;
				break;
			case 'u':
				upload = 1;
				upload_file = optarg;
				break;
			case 'f':
				force = 1;
				break;
			case 'w':
				upload_file = optarg;
				write = 1;
				break;
			case 'U':
				upload = 1;
				erase_all = 1;
				upload_file = optarg;
				break;
			case 'e':
				erase = 1;
				break;
			case 'r':
				reset = 1;
				break;
			case 'v':
				verify = 1;
				break;
			case 'h':
			case 'H':
			case '?':
			default:
				showhelp();
				exit(0);
				break;
		}
	}

	if (upload+reset+write+erase+verify > -3) {
		fprintf(stderr, "upslug: Cannot perform all operations requested in a single command -u -r -w -e -v -U -E must be used individually\n");
		exit(1);
	}

	printf("Using %s to find slugs...\n", eth_name);
	unsigned char bcast_addr[6] = {0xFF,0xFF,0xFF,0xFF,0xFF,0xFF};

	/* Open a raw socket */

	p = socket(PF_PACKET, SOCK_DGRAM, htons(UPGRADE_HW_ETHER));

	if (p == -1) {
		perror("upslug: socket");
		exit(1);
	}
	timeout.tv_sec = 2;
	timeout.tv_usec = 0;
	if(setsockopt(p, SOL_SOCKET, SO_RCVTIMEO,  &timeout, sizeof(timeout)) !=0)
	{
		perror("upslug: setsockopt");
		fprintf(stderr, "upslug: error Setting Socket Options!\n");
		exit(1);
	}

	/* Looking network interface index */
	strcpy(if_eth.ifr_name, eth_name);

	rc = ioctl(p, SIOCGIFINDEX, &if_eth);

	ifindex = if_eth.ifr_ifindex;

	if (rc < 0) {
		perror("upslug: ioctl");
		exit(rc);
	}

	/* Setup destination address */
	memset(&slug_addr, 0, sizeof(slug_addr));
	slug_addr.sll_family = AF_PACKET;
	slug_addr.sll_protocol = htons(UPGRADE_HW_ETHER);
	slug_addr.sll_ifindex = ifindex;
	memcpy(slug_addr.sll_addr, bcast_addr, sizeof(bcast_addr));
	slug_addr.sll_halen = sizeof(bcast_addr);

	/* Setup payload */
	memset(&dlc, 0, sizeof dlc);
	dlc.wcmd = GET_VERSION_INFO;

	/* This broadcasts a request for version info, the result will
	 * be a stream of packets from every listening slug, received
	 * below in the recvfrom.
	 */
	rc = sendto(p, &dlc, DLC_LEN, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);

	if (rc < 0) 
	{
		perror("upslug: sendto");
		exit(1);
	}

	/* Detect all slugs on the network (well, up to 5) */
	printf("Finding Slugs...");
	i=0;
	do {
		struct sockaddr_ll from_addr;
		socklen_t from_len;

		memset(&from_addr, 0, sizeof from_addr);
		from_len = sizeof from_addr;

		memset(&dlc, 0, sizeof dlc);
		dlc.wcmd = 0xffff; /* invalidate this */
		rc = recvfrom(p, &dlc, sizeof dlc, 0, (struct sockaddr*)&from_addr, &from_len);
		if (rc < 1)
			break;

		/* Must validate the dlc because non-slugs can respond! */
		if (rc >= DLC_LEN && dlc.wcmd == GET_VERSION_INFO) {
			memcpy(to_addr[i].sll_addr, from_addr.sll_addr, from_addr.sll_halen);
			to_addr[i++].sll_halen = from_addr.sll_halen;
		}
	} while (i < 5);
	printf(" %d found\n", i);

	/* Print out the results */
	if (i == 0) {
		printf("upslug: no slugs found, turkey problems?\n");
		exit(1);
	} if (slug_id != 0) {
		char address_string[20];
		
		for (j=0; j<i && slug_num < 0; j++) {
			unsigned char *o = to_addr[j].sll_addr;
			sprintf(address_string, "%02x:%02x:%02x:%02x:%02x:%02x", o[0], o[1], o[2], o[3], o[4], o[5]);
			if(strcmp(slug_id, address_string) == 0) {
				slug_num = j;
				printf("Known slug found\n");
			}
		}

		if (slug_num == -1) {
			fprintf(stderr, "%s: specified Slug not found\n", slug_id);
			exit(1);
		}
	} else {
		for(j=0; j<i; j++) {
			unsigned char *o = to_addr[j].sll_addr;
			printf("%i: %02x:%02x:%02x:%02x:%02x:%02x\n", j+1, o[0], o[1], o[2], o[3], o[4], o[5]);
			
		}

		/* Let the user choose which slug they want to upgrade */
		printf("Choose your slug [1-%d, 0 to abort]: ", i);
		slug_num = -1;
		scanf("%i", &slug_num);
		--slug_num;
		if (slug_num < 0 || slug_num >= i)
		{
			if (slug_num > 0)
				printf("upslug: %d: no such slug\n", slug_num);
			exit(slug_num > 0);
		}
	}

	/* Overwrite the broadcast address with the selected slug address. */
	memcpy(slug_addr.sll_addr, to_addr[slug_num].sll_addr, to_addr[slug_num].sll_halen);
	slug_addr.sll_halen = to_addr[slug_num].sll_halen;

	/* Repeat the 'get version' command from before. */
	memset(&dlc, 0, sizeof dlc);
	dlc.wcmd = GET_VERSION_INFO;

	rc = sendto(p, &dlc, DLC_LEN, 0, (const struct sockaddr*)&slug_addr, sizeof slug_addr);
	if (rc < 0) {
		perror("upslug: sendto(transmit)");
		exit(1);
	}

	/* This should really read all the available bytes, but it doesn't -
	 * it just relies on the result being the right size.
	 */
	memset(&dlc, 0, sizeof dlc);
	rc = recv(p, &dlc, sizeof dlc, 0);

	/* decode version info */
	if (rc >= DLC_LEN && dlc.wcmd == GET_VERSION_INFO) {
		const VCI_TABLE *vci;

		if (ntohs(dlc.wLen) < DLC_LEN + VCI_LEN) {
			printf("version info length too short (%d bytes, need %d)\n",
					dlc.wLen, DLC_LEN+VCI_LEN);
			exit(1);
		}

		vci = (const VCI_TABLE *)dlc.bData;

		printf("VerControl  = %04x\n", ntohs(vci->VerControl));
		printf("DownControl = %04x\n", ntohs(vci->DownControl));
		printf("Hid         = ");
		for (i=0; i<32; i += 2) {
			printf("%04x ", ntohs(*(unsigned short *)(vci->Hid + i)));
		}
		printf("\n");
		printf("Hver        = %04x\n", ntohs(vci->Hver));
		printf("ProdID      = %04x\n", ntohs(vci->ProdID));
		printf("ProtID      = %04x\n", ntohs(vci->ProtID));
		printf("FuncID      = %04x\n", ntohs(vci->FuncID));
		printf("Fver        = %04x\n", ntohs(vci->Fver));
		printf("Cseg        = %04x\n", ntohs(vci->Cseg));
		printf("Csize       = %04x\n", ntohs(vci->Csize));
	} else {
		printf("Unknown packet type (%04x)\n", dlc.wcmd);
	}

	if (upload == 1) {
		//Get to upgrade mode
		if (erase_all < 0) {
			if (force < 0) {
				printf("Do you really want to upgrade the slug(y/n)? ");
				verify_file_contents(upload_file);
				fflush(stdin);
				yes_no = 'a';
				while (yes_no != 'y' && yes_no != 'n') {
					scanf("%c", &yes_no);
					if(yes_no == 'n') {
						printf("Aborting...\n");
						exit(0);
					}
				}
			}
			down_request(&sequence, p, slug_addr);
			write_image(&sequence, p, slug_addr, upload_file);
			verify_image(&sequence, p, slug_addr);
			reset_slug(p, slug_addr);
			exit(0);
		} else if (erase_all == 1) {
			/* For safety this is disabled - it will overwrite redboot. */
			fprintf(stderr, "upslug: -U: disabled...\n");
			exit(1);
			if(force < 0) {
				printf("Do you really want to erase everything and upgrade the slug(y/n)? ");
				verify_file_contents(upload_file);
				fflush(stdin);
				yes_no = 'a';
				while(yes_no != 'y' && yes_no != 'n') {
					scanf("%c", &yes_no);
					if(yes_no == 'n') {
						printf("Aborting...\n");
						exit(0);
					}
				}
			}
			eall(&sequence, p, slug_addr);
			write_image(&sequence, p, slug_addr, upload_file);
			verify_image(&sequence, p, slug_addr);
			reset_slug(p, slug_addr);
			exit(0);
		}
	} else if (verify == 1) {
		verify_image(&sequence, p, slug_addr);
		exit(0);
	} else if (reset == 1) {
		if (force < 0) {
			printf("Do you really want to reset the slug(y/n)? ");
			fflush(stdin);
			yes_no = 'a';
			while (yes_no != 'y' && yes_no != 'n') {
				scanf("%c", &yes_no);
				if(yes_no == 'n') {
					printf("Aborting...\n");
					exit(0);
				}
			}
		}
		reset_slug(p, slug_addr);
		exit(0);
	} else if(write == 1) {
		if (force < 0) {
			printf("Do you really want to reset the slug(y/n)? ");
			fflush(stdin);
			yes_no = 'a';
			while(yes_no != 'y' && yes_no != 'n') {
				scanf("%c", &yes_no);
				if(yes_no == 'n') {
					printf("Aborting...\n");
					exit(0);
				}
			}
		}
		fprintf(stderr, "upslug: -w: NYI\n");
	} else if(erase == 1) {
		if (force < 0) {
			printf("Do you really want to reset the slug(y/n)? ");
			fflush(stdin);
			yes_no = 'a';
			while(yes_no != 'y' && yes_no != 'n') {
				scanf("%c", &yes_no);
				if(yes_no == 'n') {
					printf("Aborting...\n");
					exit(0);
				}
			}
		}
		fprintf(stderr, "upslug: -e: NYI\n");
	}

	fprintf(stderr, "upslug: nothing done\n");
	return 1;
}