/*
 * nwlilo.c: a small LILO clone for the NetWinder
 * Copyright 2000  Rebel.com
 *
 * History
 *	9.Mar.2000 - Woody Suwalski <woody@netwinder.org>
 *	Original writing.
 *
 *	27.Mar.2000 - woody - v.0.6
 *	since BIOS 2.2.1 can interpret a zero chain entry to generate
 *	NULL block w/o reading one from the hard disk - stop adding 
 *	empty blocks at the end of chain table
 *	Added nwlilo -uninstall option
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <linux/fs.h>


#define NWLILO_VERSION	6

#define debug_printf printf

#define SECTORSIZE 512
#define PART_TABLE 0x1BE
typedef struct partition {
	char C_boot;		/* Boot Indicator */
	char C_starthead;	/* Starting head */
	char C_startsec;	/* Starting sector  */
	/*
	 * bits 0-5
	 * bits 6-7 are bits 8 and 9 
	 * of cylynder
	 */
	char C_lowstartcyl;	/* Starting cylynder (low 8 bits) */
	char C_ID;		/* System ID */
	char C_endhead;		/* Ending Head */
	char C_endsec;		/* Ending Sector (see startsec) */
	char C_lowendcyl;	/* Ending Cylynder (low 8 bits) */
#if 0
	unsigned int UI_firstpsector;	// First partition sector (MISALIGNED!!!!)
#endif
	char C_firstpsector[4];
#if 0
	unsigned int UI_psize;	// Sectors in partition (MISALIGNED!!!)
#endif
	char C_psize[4];
} __attribute__ ((packed)) PARTITION;


/*
* This is the signature and params struc for NWLILO, written to sector 2...
*/
#define NWLILOSIGN "*** NW  lilo ***"
#define	NWLILOSIZE 16
typedef struct nwlilo {
char	signature[NWLILOSIZE];
unsigned int	version;		/* version */
unsigned int	sectors;		/* number of sectors in the table */
unsigned int	words_per_block;	/* straight table or sector+count */
unsigned int	sectors_per_block;	/* how many sectors per table entry */
char	filename[128];			/* filename of the kernel image */
char	cmdline[256];			/* to be copied to the param struc */
} __attribute__ ((packed)) NWLILO;

char	buff[SECTORSIZE];
int	chain[4096];


int main(int argc, char **argv)
{
int	fd;
int	partition_start;
int	i, j;
int	fsize;
int	curblock;
NWLILO* nwlilo;
int	show_prompt = 1;
int	compact = 0;
int	verbose = 0;
int	verify = 0;
int	uninstall = 0;
char*	our_name = NULL;
char*	fname = NULL;
char*	cmdline = NULL;
int	partitions[4];
int	blocksize;
int	sectorsperblock;
int	nwlilosectors;
struct	stat imgstat;

	/* Get our executable's name (should be "nwlilo") */
	our_name = (our_name = strrchr(argv[0], '/')) ? our_name + 1 : argv[0];

	/* Parse the command line args */
	for (i = 1; i < argc; i++) {
		if (!strcmp(argv[i], "-compact")) {
			compact = 1;
			continue;
		}
		if (!strcmp(argv[i], "-yes")) {
			show_prompt = 0;
			continue;
		}
		if (!strcmp(argv[i], "-verbose")) {
			verbose = 1;
			continue;
		}
		if (!strcmp(argv[i], "-verify")) {
			verify = 1;
			continue;
		}
		if (!strcmp(argv[i], "-uninstall")) {
			uninstall = 1;
			continue;
		}
		if (!fname) {
			fname = argv[i];
			continue;
		}
		if (!cmdline) {
			cmdline = argv[i];
			continue;
		}
	}

	if (!fname && !verify && !uninstall) {
		fprintf(
			stderr, "\n"
			"NetWinder LILO clone program V.0.%02d, Rebel.com 2000.\n\n"
			"Usage: %s [flags] kernelfile [cmdline]\n"
			"where optional flags are:\n"
			"  -yes       does not prompt for confirmation\n"
			"  -compact   compact the table for faster boot\n"
			"  -verbose   print more detailed info\n"
			"  -verify    read the imag back from the map\n"
			"  -uninstall erase the NWLILO table from the disk\n"
			"If cmdline has spaces, surround it by double quotes.\n"
			"\n",
			NWLILO_VERSION, our_name
		);
		exit(1);
	}

    if (getuid() != 0)
    {
	fprintf(stderr, "NWLILO: this program needs to be run with root priviliges. Aborting...\n");
	exit(2);
    }

    if (uninstall)
    {
    
	fd = open("/dev/hda", O_RDWR);
	i = lseek(fd, 1*SECTORSIZE, SEEK_SET); /* skip the partition sector */
	if (i != 1*SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: can not seek to the NWLILO space! Uninstall aborting...\n");
	    close(fd);
	    exit(3);
	}
	i = read(fd, &buff, SECTORSIZE);
	if (i != SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: can not read the hard disk NWLILO sector! Aborting...\n");
	    close(fd);
	    exit(4);
	}

	nwlilo = (NWLILO*)&buff;

	if ( strncmp(nwlilo->signature, NWLILOSIGN, NWLILOSIZE))
	{
	    fprintf(stderr, "NWLILO not installed. Aborting uninstall...\n");
	    close(fd);
	    exit(5);
	}

	printf ("NWLILO: Uninstalling image %s written by nwlilo version 0.%02X with cmdline \"%s\".\n",
	    nwlilo->filename,
	    nwlilo->version,
	    nwlilo->cmdline);

	i = lseek(fd, 1*SECTORSIZE, SEEK_SET); /* skip the partition sector */
	if (i != 1*SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: can not seek to the NWLILO space! Uninstall aborting...\n");
	    close(fd);
	    exit(3);
	}

	/* clear the table buff to make sure we write zeros... */
	for (i = 0; i < 4096; i++)
	    chain[i] = 0;

	i = write(fd, &chain, (nwlilo->sectors + 1) * SECTORSIZE); /* write over table */
	if (i != (nwlilo->sectors + 1) * SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: error erasing the NWLILO table from the hard disk! Aborting...\n");
	    close(fd);
	    exit(6);
	}
    
	close(fd);
	printf ("NWLILO table erased OK.\n");

	return (0);
    }	/* uninstall done here... */


    if (verify)
    {
    int	fdout;
    char* vbuf;
    
	fd = open("/dev/hda", O_RDONLY);
	i = lseek(fd, 1*SECTORSIZE, SEEK_SET); /* skip the partition sector */
	if (i != 1*SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: can not seek to the NWLILO space! Verify aborting...\n");
	    close(fd);
	    exit(3);
	}
	i = read(fd, &buff, SECTORSIZE);
	if (i != SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: can not read the hard disk NWLILO sector! Aborting...\n");
	    close(fd);
	    exit(4);
	}

	nwlilo = (NWLILO*)&buff;

	if ( strncmp(nwlilo->signature, NWLILOSIGN, NWLILOSIZE))
	{
	    fprintf(stderr, "NWLILO not installed. Aborting verify...\n");
	    close(fd);
	    exit(5);
	}

	printf ("NWLILO: Image %s written by nwlilo version 0.%02X with cmdline \"%s\".\n",
	    nwlilo->filename,
	    nwlilo->version,
	    nwlilo->cmdline);

	sectorsperblock = nwlilo->sectors_per_block;
	i = read(fd, &chain, nwlilo->sectors * SECTORSIZE); /* read table */
	if (i != nwlilo->sectors * SECTORSIZE)
	{
	    fprintf(stderr, "NWLILO: error reading the table from the hard disk! Aborting...\n");
	    close(fd);
	    exit(6);
	}
    
	if (!nwlilo->filename)
	    strcpy(nwlilo->filename, "nwlilo_verify");

	strcat(nwlilo->filename, ".verify");
	fdout = open (nwlilo->filename, O_WRONLY | O_CREAT);
	if (fd < 0) {
	    fprintf(stderr, "Error creating file %s.... Aborting.\n",nwlilo->filename);
	    exit(7);
	}

	vbuf = (char*)malloc(sectorsperblock * SECTORSIZE);
    
	j = 0;
	while (chain[j] != 0xFFFFFFFF)
	{
	    if (!chain[j])	/* NULL entry == img hole */
	    {
		for (i = 0; i < sectorsperblock * SECTORSIZE; i++)
		    *(char*)(vbuf + i) = 0;
	    }
	    else
	    {
		i = lseek(fd, chain[j]*SECTORSIZE, SEEK_SET);
		i = read(fd, vbuf, sectorsperblock * SECTORSIZE); /* read block */
	    }
	    i = write(fdout, vbuf, sectorsperblock * SECTORSIZE);
	    j++;
	}
    
	free (vbuf);    
	close(fdout);
	close(fd);
	printf ("Created verify file %s.\n", nwlilo->filename);

	return (0);
    }	/* verify done here... */

    /*
    * snoop out from parition table offsets to the partitions...
    */
    fd = open("/dev/hda", O_RDONLY);
    i = read(fd, &buff, SECTORSIZE);
    close(fd);
    if (i < 0)
    {
	fprintf(stderr, "Could not read the partition table of /dev/hda (error=%d).\n", i);
	exit(8);
    }
    
    for (i=0; i<4; i++)
    {
    PARTITION* part;

	part = (PARTITION *) (buff + PART_TABLE + i * sizeof(PARTITION));

	/* the integer is misaligned - do the reading by bytes... */
	partition_start = part->C_firstpsector[0];
	partition_start += (part->C_firstpsector[1] << 8);
	partition_start += (part->C_firstpsector[2] << 16);
	partition_start += (part->C_firstpsector[3] << 24);

	partitions[i] = partition_start;
	if (verbose)
    	    debug_printf("Partition %d starts at sector %X.\n", i+1, partitions[i]);
    }

    /*
    * open the kernel file, collect the sector map table...
    */
    fd = open(fname, O_RDONLY);
    if (fd < 0) {
		fprintf(stderr, "Error opening file %s.... Aborting.\n",fname);
		exit(9);
    }

    i = (ioctl (fd, FIGETBSZ, &blocksize));
    if (i < 0)
    {
	blocksize = 1024;
    }
    else
    {
	if (!blocksize || (blocksize & (SECTORSIZE-1)))
	    blocksize = 1024;
    }
    sectorsperblock = blocksize / SECTORSIZE;

    i = fstat(fd, &imgstat);
    if (i < 0)
    {
	fprintf(stderr, "Could not get the file info for %s... (error=%d).\n", fname, i);
	close(fd);
	exit(10);
    }

    if (imgstat.st_dev < 0x301 || imgstat.st_dev > 0x304)
    {
	fprintf(stderr, "Error: use kernel image on a local ide hard disk (/dev/hda1.../dev/hda4).\n");
	close(fd);
	exit(11);
    }

    i = imgstat.st_dev - 0x301;
    partition_start = partitions[i];

    if (show_prompt) {
	int ch;
	fprintf(
		stderr, "\n%s V.0.%02d : ready to map\n"
			"  file    : \"%s\"\n"
			"  size    : %d bytes\n"
			"  cmdline : \"%s\"\n"
			"  compact : %s\n"
			"\nEnter Y to continue, or anything else to abort...\n",
			our_name, NWLILO_VERSION, fname,
			(int)imgstat.st_size,
			cmdline?cmdline:"none",
			compact?"yes":"no"
		);
		ch = getchar();
		if (!((ch == 'y') || (ch == 'Y'))) {
			fprintf(stderr, "Not this time... aborted.\n");
			exit(12);
		}
	}


    i = read(fd, buff, 32); /* read ELF header, if there */
    if (*(int*)buff == 0x464C457F)	/* ~ELF? */
    {
	fprintf(stderr, "NWLILO: ELF header detected in %s: use the Image or zImage kernel files... Aborting...\n", fname);
	exit(13);
    }

    /*
    * convert file sector size to 'blocksize' blocks, rounded up...
    */
    fsize = (imgstat.st_size + blocksize - 1) / blocksize;
    if (verbose)
	debug_printf("nwlilo: block size = %d bytes, img size = %d blocks.\n",
	blocksize, fsize);   

    if (!compact)
    {
	i = SECTORSIZE / sizeof(int);
	nwlilosectors = (fsize + i -1)/i;	/* 128 entries per sector */
    }
    else
    {
	fprintf(stderr, "NWLILO: compact option not supported yet... Aborting...\n");
	exit(14);
    }

    /*
    * run the mapping routine, collect output in the chain table...
    */
    for (curblock = 0; curblock < fsize; curblock++)
    {
	j = curblock;
	i = ioctl(fd,FIBMAP,&j);
	if (i<0)
	{
	    fprintf(stderr, "Could not get the mapping of block %d (error=%d). Aborting...\n",curblock, i);
	    close(fd);
	    exit(15);
	}

	if (!j)
	{
	    if (verbose)
		debug_printf("Covering hole at block %d...\n",curblock);
	}
	else
	{
	    j *= sectorsperblock;
	    j += partition_start;
	    if (verbose)
		debug_printf("Block %d sector %d mapped at %d.\n",
		    curblock, curblock * sectorsperblock, j);
	}

	chain[curblock]=j;

    }
    /* mark end-of-chain */
    chain[curblock] = 0xFFFFFFFF;
    close(fd);

    /*
    * prepare the NWLILO header sector structure...
    */

    for(i = 0; i < SECTORSIZE; i++)
	*(buff+i) = 0; 

    nwlilo = (NWLILO*)&buff;
    strcpy(nwlilo->signature, NWLILOSIGN);
    nwlilo->version = NWLILO_VERSION;

    if (!compact)
    {
	nwlilo->words_per_block = 1;
    }
    else
    {
	nwlilo->words_per_block = 2;
    }

    nwlilo->sectors = nwlilosectors;
    nwlilo->sectors_per_block = sectorsperblock;
    if (fname) strcpy(nwlilo->filename, fname);
    if (cmdline) strcpy(nwlilo->cmdline, cmdline);

    /* we start at setcor 2, then nwlilo sectors */
    if (partitions[0] < (2 + nwlilo->sectors))
    {
	fprintf(stderr, "NWLILO: not enough free space at the head of the disk. Aborting...\n");
	exit(16);
    }
#if 1
    fd = open("/dev/hda", O_WRONLY);
    i = lseek(fd, 1*SECTORSIZE, SEEK_SET); /* skip the partition sector */
    if (i != 1*SECTORSIZE)
    {
	fprintf(stderr, "NWLILO: can not seek to the NWLILO space! Aborting...\n");
	close(fd);
	exit(18);
    }
    
    i = write(fd, &buff, SECTORSIZE);	/* write our NWLILO signature */
    if (i != SECTORSIZE)
    {
	fprintf(stderr, "NWLILO: can not write to the hard disk! Aborting...\n");
	close(fd);
	exit(19);
    }

    i = write(fd, &chain, nwlilo->sectors * SECTORSIZE); /* write table */
    if (i != nwlilo->sectors * SECTORSIZE)
    {
	fprintf(stderr, "NWLILO: error writing to the hard disk! Aborting...\n");
	close(fd);
	exit(20);
    }

    close(fd);

#endif
    fprintf(stderr, "NWLILO: file %s mapped OK.\n",fname);
    return(0);
}