/*
 * main.c
 *
 *  Created on: 2009-11-10
 *      Author: wodz
 */

#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <argp.h>
#include <stdio.h>
#include <string.h>

struct arguments {
	char *filename;
	char *patchname;
	char *outfilename;
	unsigned long int offset;
};

/* argp strings */
const char *argp_program_version = "binpatch 0.1";
const char *argp_program_bug_address = "marcin.bukat@gmail.com";
static char doc[] = "Implant binary blob into firmware at specyfied offset";
static char args_doc[] = "filenames";

/* valid arguments */
static struct argp_option options[] = {
		{"file", 'i', "FIRMWARE FILE", 0, "Firmware file"},
		{"patch", 'p', "PATCH FILE", 0, "Patch file (binary blob"},
		{"offset", 'o', "OFFSET", 0, "Offset at which patch should be implanted"},
		{"outfile", 'O', "OUT FILE", 0, "Output filename"},
		{0}
};

/* argp parsing backend function */
static error_t parse_opt(int key, char *arg, struct argp_state *state)
{
	struct arguments *arguments = state->input;

	switch (key)
	{
		case 'i':
			arguments->filename = arg;
			break;
		case 'p':
			arguments->patchname = arg;
			break;
		case 'o':
			sscanf(arg,"%li",&arguments->offset);
			break;
		case 'O':
			arguments->outfilename = arg;
			break;
		default:
			return ARGP_ERR_UNKNOWN;
	}
return 0;
}

static struct argp argp = {options, parse_opt, args_doc, doc};

int main (int argc, char **argv)
{
		struct arguments arguments;

		int filefd,patchfd,outfilefd;
        struct stat filesb, patchsb;
        char *fileaddr,*patchaddr;

        /* parse arguments */
        argp_parse (&argp, argc, argv, 0, 0, &arguments);

        filefd = open(arguments.filename,O_RDWR);
        patchfd = open(arguments.patchname,O_RDONLY);
        outfilefd = creat(arguments.outfilename,S_IRUSR|S_IWUSR|S_IRGRP|S_IROTH);

        /* check validity */
        if (filefd == -1)
        {
        	printf("invalid firmware filename %s\n",arguments.filename);
        	return -1;
        }

        if (patchfd == -1)
        {
        	printf("invalid patch filename %s\n",arguments.patchname);
        	return -1;
        }

        if (outfilefd == -1)
        {
        	printf("can't create output file %s\n",arguments.outfilename);
        	return -1;

        }

        if ( fstat(filefd, &filesb) == -1 )
        {
        	printf("unable to get firmware file size\n");
        	return -1;
        }

        if ( fstat(patchfd, &patchsb) == -1 )
        {
        	printf("unable to get patch file size\n");
        	return -1;
        }

        if (arguments.offset + patchsb.st_size > filesb.st_size)
        {
        	printf("unable to patch at this offset - end of file");
        	return -1;
        }

        fileaddr = mmap(NULL,filesb.st_size,PROT_READ|PROT_WRITE,MAP_PRIVATE,filefd,0);
        patchaddr = mmap(NULL,patchsb.st_size,PROT_READ,MAP_PRIVATE,patchfd,0);

        memcpy(fileaddr + arguments.offset,patchaddr,patchsb.st_size);

        if ( write(outfilefd,fileaddr,filesb.st_size) != filesb.st_size )
        {
        	printf("error writing output file %s\n",arguments.outfilename);
        }
        else
        {
        	printf("firmware file: %s size 0x%x bytes\npatch file: %s size 0x%x bytes\noffset: 0x%x\n",
        			arguments.filename,
        			(unsigned int)filesb.st_size,
        			arguments.patchname,
        			(unsigned int)patchsb.st_size,
        			(unsigned int)arguments.offset);
        	printf("writing patched firmware to %s\n",arguments.outfilename);
        }

        munmap(fileaddr,filesb.st_size);
        munmap(patchaddr,patchsb.st_size);

        close(filefd);
        close(patchfd);
        close(outfilefd);

return 0;
}