/* zImage_util.c
 *
 * Copyright 2003 Robert Whaley <rwhaley@applieddata.net>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.

 */

#include <sys/types.h>
#include <sys/stat.h>

#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <time.h>

#include <sys/ioctl.h>
#include <sys/mount.h>

#define CMDLINE_LIMIT 1024
#define REGISTERS_LIMIT 8192

typedef	struct	_file_header {
	unsigned long	Size;
	unsigned long	CheckSum;
	unsigned long	Date;
	unsigned long	Version;
	unsigned long   Count;
	unsigned long   BootStatus;
	char cmdline[CMDLINE_LIMIT];
	char registers[REGISTERS_LIMIT];
}	FileHeader;

#define BUF_SIZE (2*1024*1024)

char buf[BUF_SIZE];

int fix_cmdline_end(char *cmdline, int end)
{
	while (end && (cmdline[end] == 0 || cmdline[end] == ' ' || cmdline[end] == '\n' ||
		       cmdline[end] == '\r' || cmdline[end] == '\t')) {

		cmdline[end--] = 0;
	}

	return end;
}

unsigned long checksum(char *data, unsigned long len)
{
	unsigned long result;
	int i;

	result = 0;

	for (i=0; i<len; i++)
		result += ( i ^ data[i] );

	return result;
}

// histerical reasons for using this format!
unsigned long convert_to_dos_date_time(time_t mtime)
{
	struct tm *dt;
	unsigned long result = 0;

	dt = gmtime(&mtime);

	result |= (0x0000001f & dt->tm_sec) / 2;
	result |= 0x000007e0 & (dt->tm_min << 5);
	result |= 0x0000f100 & (dt->tm_hour << 11);

	result |= 0x001f0000 & (dt->tm_mday << 16);
	result |= 0x01e00000 & ((dt->tm_mon + 1) << 21);
	result |= 0xfe000000 & ((dt->tm_year - 80) << 25);

	return result;
}

struct tm * convert_from_dos_date_time(unsigned long dos_dt)
{
	static struct tm result;

	result.tm_sec =   (0x0000001f & dos_dt) * 2;
	result.tm_min =   (0x000007e0 & dos_dt) >> 5;
	result.tm_hour =  (0x0000f100 & dos_dt) >> 11;
	result.tm_mday =  (0x001f0000 & dos_dt) >> 16;
	result.tm_mon  = ((0x01e00000 & dos_dt) >> 21) - 1;
	result.tm_year = ((0xfe000000 & dos_dt) >> 25) + 80;

	return &result;
}


void list_zImage_partition(int partition_fd)
{
	FileHeader fh;
	int zImage_size;
	int res;
	int computed_checksum = 0;

	res = read(partition_fd, (char *)&fh, sizeof(fh));
	if (res != sizeof(fh)) {
		perror("Error reading partition\n");
		exit(-1); 
	}

	res = read(partition_fd, buf, fh.Size);
	if (res != fh.Size) {
		perror("Error reading partition\n");	
		printf("Error expected zImage of %d bytes but only read %d bytes\n", fh.Size, res);
		exit(-1); 
	}
       
	computed_checksum = checksum(buf, fh.Size);

	printf("zImage size: %d\n", fh.Size);
	printf("header checksum: %#x, zImage checksum: %#x, %s\n", fh.CheckSum, computed_checksum,
	       fh.CheckSum == computed_checksum ? "Match" : "No Match!");
	printf("header date: %s\n", asctime(convert_from_dos_date_time(fh.Date)));
}

void write_zImage_partition(int partition_fd, int zImage_fd, int cmdline_fd, int register_fd)
{
	FileHeader fh;
	int zImage_size;
	int res;
	struct stat st;

	zImage_size = read(zImage_fd, buf, sizeof(buf));
	if (zImage_size < 0) {
		perror("Error reading zImage file.\n");
		exit(-1);
	}

	fstat(zImage_fd, &st);

	bzero(&fh, sizeof(fh));

	fh.Size = zImage_size;
	fh.CheckSum = checksum(buf, zImage_size);
	fh.Version = 0xFFFFFFFF;
	fh.Date = convert_to_dos_date_time(st.st_mtime);
  
	res = read(cmdline_fd, fh.cmdline, sizeof(fh.cmdline));
	if (res < 0) {
		perror("Error reading cmdline file.\n");
		exit(-1);
	}

	fix_cmdline_end(fh.cmdline, res);

	if (register_fd != -1) {

		if (res < 0) {
			perror("Error reading register file\n");
			exit(-1);
		} 
		res = read(register_fd, fh.registers, sizeof(fh.registers));
	}

	res = write(partition_fd, (char *)&fh, sizeof(fh));
	if (res != sizeof(fh)) {
		perror("Error writing partition\n");
		exit(-1); 
	}

	res = write(partition_fd, buf, zImage_size);
	if (res != zImage_size) {
		perror("Error writing partition\n");
		exit(-1); 
	}

}

void read_zImage_partition(int partition_fd, int zImage_fd, int cmdline_fd, int register_fd)
{
	FileHeader fh;
	int zImage_size;
	int res;

	res = read(partition_fd, (char *)&fh, sizeof(fh));
	if (res != sizeof(fh)) {
		perror("Error reading partition\n");
		exit(-1); 
	}

	res = read(partition_fd, buf, fh.Size);
	if (res != fh.Size) {
		perror("Error reading partition\n");	
		printf("Error expected zImage of %d bytes but only read %d bytes\n", fh.Size, res);
		exit(-1); 
	}

	zImage_size = write(zImage_fd, buf, fh.Size);
	if (zImage_size < 0) {
		perror("Error writing zImage file.\n");
		exit(-1);
	}

	res = write(cmdline_fd, fh.cmdline, strlen(fh.cmdline));
	if (res < 0) {
		perror("Error writing cmdline file.\n");
		exit(-1);
	}
	
	res = write(register_fd, fh.registers, strlen(fh.registers));
	if (res < 0) {
		perror("Error writing register file.\n");
		exit(-1);
	}
}

#define USAGE "Usage: %s [-l] [-r] [-w] [-p <mtd_partition #>] [-d <mtd_partition_name>] [-z <zImage>] [-c <cmdline>] [-R <register>]\n"

int main(int argc, char *argv[])
{
	char erase_command[100] = "eraseall ";
	char *zImage_filename;
	int zImage_fd;
	char *cmdline_filename;
	int cmdline_fd;
	char *register_filename;
	int register_fd;
	char *partition_filename;
	int partition_number;
	int partition_fd;
	int list_it=0, read_it=0, write_it=0;
	int partition_flag = O_RDONLY;
	int file_flag = O_RDONLY;
	int k;

	partition_filename = "/dev/mtdc2";
	zImage_filename = "zImage";
	cmdline_filename = "cmdline.txt";
	register_filename = "register.txt";

	for (k=1; k<argc; k++) {

		if (!strcmp(argv[k], "-l"))
			list_it = 1;
		else if (!strcmp(argv[k], "-r")) {
			read_it = 1;
			file_flag = O_RDWR | O_CREAT;
		}
		else if (!strcmp(argv[k], "-w")) {
			write_it = 1;
			partition_flag = O_RDWR;
		}
		else if (!strcmp(argv[k], "-p")) {
			sscanf(argv[++k], "%d", &partition_number);
 sprintf(partition_filename, "/dev/mtdc%1d", partition_number * 2);
		}
		else if (!strcmp(argv[k], "-d"))
			partition_filename = argv[++k];
		else if (!strcmp(argv[k], "-z"))
			zImage_filename = argv[++k]; 
		else if (!strcmp(argv[k], "-c"))
			cmdline_filename = argv[++k]; 
		else if (!strcmp(argv[k], "-R"))
			register_filename = argv[++k];
		else {
			printf(USAGE);
			exit(-1);
		}
	}
	if (!read_it && !write_it)
		list_it = 1;

	if (read_it && write_it) {
		printf("Can't use both -r and -w in same run\n");
		exit(-1);
	}

	partition_fd = open(partition_filename, partition_flag);
	if (partition_fd == -1) {
		printf("Error: Could not open partition special file: %s\n", partition_filename);
		exit(-1);
	}

	if (list_it) {
		list_zImage_partition(partition_fd);
		lseek(partition_fd, 0, SEEK_SET);
	}

	if (read_it || write_it) {
		zImage_fd = open(zImage_filename, file_flag);
		if (zImage_fd == -1) {
			printf("Error: Could not open zImage file: %s\n", zImage_filename);
			exit(-1);
		}
		
		cmdline_fd = open(cmdline_filename, file_flag);
		if (cmdline_fd == -1) {
			printf("Error: could not open cmdline.txt file: %s\n", cmdline_filename);
			exit(-1);
		}
   
		register_fd = open(register_filename, file_flag);
		if (register_fd == -1) {
			printf("Warning: could not open register.txt file: %s\n", register_filename);
		}
	}

	if (read_it) {
		read_zImage_partition(partition_fd, zImage_fd, cmdline_fd, register_fd);
		lseek(partition_fd, 0, SEEK_SET);
		lseek(zImage_fd, 0, SEEK_SET);
		lseek(cmdline_fd, 0, SEEK_SET);
		lseek(register_fd, 0, SEEK_SET);

	}

	if (write_it) {
		printf("Erasing zImage partition %s ... ", partition_filename);
		fflush(stdout);
		strncat(erase_command, partition_filename, sizeof(erase_command));
		strncat(erase_command, "| grep complete > /dev/null 2>&1", sizeof(erase_command));
		if (system(erase_command)) {
                        printf("\nErase failed.\n");
			exit(-1);
		}

		printf("done.\nWriting zImage partition %s ... ", partition_filename);
		fflush(stdout);
		write_zImage_partition(partition_fd, zImage_fd, cmdline_fd, register_fd);

		printf("done.\n");
	}

	exit(0);
}