#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/stat.h>
#include <sys/wait.h>


int verbose;
int force;

struct partition {
	int id;
	char *name;
	unsigned size;
	unsigned erasesize;
	int is_nand;
	struct partition *next;
} *parts;


static int read_partitions();
static struct partition *get_part(char *name);
static struct partition *get_partno(char *arg);
static int valid_file(char *name, struct partition *part, int *size);
static int open_mtd(struct partition *part);
static void do_system(char *cmd);
static void set_force(char *arg);


static void usage(char *prog, int rc)
{
	char *p;

	if((p = strrchr(prog, '/'))) ++p; else p = prog;

	printf("usage: %s {-n partno | -p partname} file\n", p);

	exit(rc);
}


void nor_write(struct partition *part, int in, char *fname, int fsize)
{
	int out, n, wrote;
	char *buf;

	if((out = open_mtd(part)) < 0) exit(1);

	if(!(buf = malloc(part->erasesize))) {
		printf("Out of memory\n");
		exit(1);
	}

	while(fsize >= part->erasesize) {
		if((n = read(in, buf, part->erasesize)) != part->erasesize) {
			printf("Short full read! %d/%d\n", n, part->erasesize);
			exit(1);
		} else if((wrote = write(out, buf, n)) != n) {
			printf("Short full write %d/%d\n", wrote, n);
			exit(1);
		}
		fsize -= part->erasesize;
	}

	memset(buf, 0xea, part->erasesize);

	if(fsize > 0) {
		if((n = read(in, buf, part->erasesize)) != fsize) {
			printf("Short read! %d/%d\n", n, fsize);
			exit(1);
		} else if((wrote = write(out, buf, part->erasesize)) != part->erasesize) {
			printf("Short write %d/%d\n", wrote, n);
			exit(1);
		}
	}

	close(in);
	close(out);
}


void nand_write(struct partition *part, int fd, char *fname)
{
	char cmd[100];
	char mtdname[20];

	close(fd); /* don't need this. */

	sprintf(mtdname, "/dev/mtd%d", part->id);

	sprintf(cmd, "flash_eraseall %s", mtdname);
	do_system(cmd);

	sprintf(cmd, "nandwrite -p %s %s", mtdname, fname);
	do_system(cmd);
}


int main(int argc, char *argv[])
{
	int c, fd, size;
	struct partition *part = NULL;

	read_partitions();

	while((c = getopt(argc, argv, "f:hn:p:v")) != EOF)
		switch(c) {
		case 'f': set_force(optarg); break;
		case 'h': usage(argv[0], 0); break;
		case 'p': part = get_part(optarg); break;
		case 'n': part = get_partno(optarg); break;
		case 'v': ++verbose; break;
		default: puts("Sorry!"); exit(1);
		}

	if(optind >= argc) usage(argv[0], 1);

	if(part == NULL) {
		printf("No partition specified.\n");
		usage(argv[0], 1);
	}

	fd = valid_file(argv[optind], part, &size);

	if(part->is_nand)
		nand_write(part, fd, argv[optind]);
	else
		nor_write(part, fd, argv[optind], size);

	puts("Success.");
	return 0;
}


#define IMAGE_MAGIC 0x27051956
#define UBOOT_SIZE (512 * 1024)

static void check_u_boot(int fd, int size)
{
	char buf[20];
	int failed = 0;

	if(size != UBOOT_SIZE) {
		printf("u-boot file size wrong! %d != %d\n",
			   size, UBOOT_SIZE);
		++failed;
	}

	if(read(fd, buf, sizeof(buf)) != sizeof(buf)) {
		printf("Unable to read u-boot header\n");
		exit(1); /* File errors are serious. */
	}

	if(lseek(fd, 0, SEEK_SET)) {
		printf("Unable to rewind the file\n");
		exit(1); /* File errors are serious. */
	}

	if(*(unsigned*)buf != IMAGE_MAGIC) {
		printf("Bad magic %08x != %08x\n",
			   *(unsigned*)buf, IMAGE_MAGIC);
		++failed;
	}

	if(strncmp(buf + 4, "U-Boot", 6)) {
		printf("Bad magic - no U-Boot\n");
		++failed;
	}

	if(failed) {
		if(force)
			puts("FORCED!");
		else
			exit(1);
	}
}


static void check_mount(struct partition *part)
{
	FILE *fp;
	char mount[1024];
	char mtdname[12];

	if(!(fp = fopen("/proc/mounts", "r"))) {
		printf("Unable to open /proc/mounts\n");
		exit(1);
	}

	sprintf(mtdname, "mtdblock%d", part->id);

	while(fgets(mount, sizeof(mount), fp))
		if(strstr(mount, mtdname)) {
			printf("That paritition is currently mounted.\n"
				   "Please `umount /dev/%s'\n", mtdname);
			if(!force) exit(1);
		}

	fclose(fp);
}


int valid_file(char *name, struct partition *part, int *size)
{
	int fd;
	struct stat sbuf;

	if((fd = open(name, O_RDONLY)) < 0 || fstat(fd, &sbuf)) {
		perror(name);
		exit(1);
	}

	if(sbuf.st_size > part->size) {
		printf("%s: file too big for partition %s. %ld > %d\n",
			   name, part->name, sbuf.st_size, part->size);
		exit(1);
	}

	if(size) *size = sbuf.st_size;

	if(strstr(part->name, "u-boot") ||
	   strstr(part->name, "uboot")) {
		check_u_boot(fd, sbuf.st_size);

		/* Do this last since we may need to force this. */
		if(part->is_nand) {
			printf("u-boot never in NAND!\n");
			if(!force) exit(1);
		}
	} else
		check_mount(part);

	return fd;
}


int open_mtd(struct partition *part)
{
	int fd;
	char mtdname[20];

	sprintf(mtdname, "/dev/mtdblock%d", part->id);
	if((fd = open(mtdname, O_WRONLY)) < 0) perror(mtdname);

	lseek(fd, SEEK_SET, 0);

	return fd;
}


struct partition *new_partition(int is_nand)
{
	static struct partition *tail;
	struct partition *new;

	if((new = calloc(sizeof(struct partition), 1)) == NULL) {
		printf("Out of memory\n");
		exit(1);
	}

	new->is_nand = is_nand;

	if(parts)
		tail->next = new;
	else
		parts = new;
	tail = new;

	return new;
}


int read_partitions()
{
	FILE *fp;
	char line[80], name[80];
	int rc = 0;
	int is_nand = 0;

	if((fp = fopen("/proc/mtd", "r")) == NULL) {
		perror("/proc/mtd");
		exit(1);
	}

	while(fgets(line, sizeof(line), fp)) {
		struct partition *part;

		if(strncmp(line, "mtd", 3)) continue;
		if(strstr(line, "NAND")) {
			/* This marks the start of the NAND partitions. */
			is_nand = 1;
			continue;
		}
		part = new_partition(is_nand);

		if(sscanf(line, "mtd%d: %x %x \"%[^\"]\"",
				  &part->id, &part->size, &part->erasesize, name) != 4) {
			printf("Parse error: %s\n", line);
			++rc;
			continue;
		}

		if(!(part->name = strdup(name))) {
			puts("Out of memory!");
			exit(1);
		}
	}

	return rc;
}


static struct partition *get_part(char *name)
{
	struct partition *part;

	for(part = parts; part; part = part->next)
		if(strcasecmp(name, part->name) == 0)
			return part;

	printf("Invalid partition %s\n", name);
	exit(1);
}


static struct partition *get_partno(char *arg)
{
	struct partition *part;
	char *end;
	int id;

	id = strtol(arg, &end, 0);
	if(*end && *end != ' ') {
		printf("Not a parition number: %s\n", arg);
		exit(1);
	}

	for(part = parts; part; part = part->next)
		if(part->id == id)
			return part;

	printf("Invalid parition number %s\n", arg);
	exit(1);
}


void do_system(char *cmd)
{
	int rc;

	if(verbose) printf(">> %s\n", cmd);

	rc = system(cmd);
	if(rc == -1) {
		perror(cmd);
		exit(1);
	}

	rc = WEXITSTATUS(rc);
	if(rc) {
		printf("%s returned %d\n", cmd, rc);
		exit(1);
	}
}

static void set_force(char *arg)
{
	if(strcmp(arg, "orce")) {
		printf("Invalid force\n");
		exit(1);
	}

	force = 1;
}

/*
 * Local Variables:
 * compile-command: "${CROSS_COMPILE}gcc -O3 -Wall trogdor.c -o trogdor"
 * End:
 */