/*
 * jpeg-recover2.c (2008-09-04)
 *
 * by Robert Koehler <firstname.lastname@gimi.name>
 * http://gimi.name/go/qiwRE
 *
 * A slightly extended and extremely faster implementation of Adam Glass'
 * jpg-recover (http://www.clarity.net/~adam/jpg-recover/) written in pure c.
 *
 * compiled using
 * gcc -o jpeg-recover2 -s -O3 jpeg-recover2.c
 *
 *
 * 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.
 *
 * 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.
 *
 * See http://www.gnu.org/licenses/gpl-2.0.txt for details.
 *
 */

#include <stdio.h>
#include <unistd.h>
#include <string.h>


char be_verbose = 0;
#define VERBOSE if(be_verbose)


/* \x00 added to act as a single byte wildcard  */
char *header[] = {
	"\xff\xd8\xff\xe1\x00\x00\x45\x78\x69\x66",	
	"\xff\xd8\xff\xe1\x00\x00\x45\x79\x69\x66",
	"\xff\xd8\xff\xe0\x00\x10\x4a\x46\x49\x46",	
	NULL
};


/* return true (non-zero) if there is a matching header */
int header_match(FILE *f) {
	long pos = ftell(f);
	char buf[10];
	size_t len;m
	int i, j;

	len = fread(buf, sizeof(buf), 1, f);
	fseek(f, pos + 1, SEEK_SET);

	if(!len)
		return 0;

	for(i = 0; header[i] != NULL && j != 10; i ++)
		for(j = 0; j < 10; j ++)
			if(header[i][j] != buf[j] && header[i][j] != 0)
				break;

	if(j == 10 && be_verbose) {
		printf("HEADER ");
		for(i = 0; i < 10; i++)
			printf("%02x ", buf[i] & 0xff);
		printf("\n");
	}

	return (j == 10);
}


/* return true (non-zero) if there is an end marker */
int end_match(FILE *f) {
	if(fgetc(f) != 0xFF)
		return 0;

	if(fgetc(f) == 0xD9)
		return 1;

	fseek(f, -1, SEEK_CUR);
	return 0;
}


/* copy data from f to file given by filename */
int image_copy(FILE *f, long start, long end, char *filename) {
	FILE *f2;
	long left;
	long i;
	char buf[1024 * 10];

	/* prepare */
	left = end - start;
	printf("SAVING to %s (%ld bytes)\r\n", filename, left);

	f2 = fopen(filename, "wb");
	if(f2 == NULL) {
		perror("Error: Cannot open file");
		return 0;
	}

	/* copy part */
	fseek(f, start, SEEK_SET);
	for(; left > 0; left -= sizeof(buf)) {
		fread(buf, (left > sizeof(buf)) ? sizeof(buf) : left, 1, f);
		fwrite(buf, (left > sizeof(buf)) ? sizeof(buf) : left, 1, f2);
	}

	/* finish */
	fclose(f2);
	return 1;
}


int main(int argc, char *argv[]) {
	FILE *f;
	long start = 0;
	long end;
	int count = 0;
	char buf[128];
	int i;

	/* parse parameter */
	int opt_all = 0;
	int opt_seek = 0;
	int opt_lower = 100 * 1024;
	int opt_upper = 5000 * 1024;
	char *opt_pre = "RIMG_";
	char *arg_img = NULL;

	for(i = 1; i < argc; i ++)
		if(argv[i][0] == '-')
			switch(argv[i][1]) {
				case 'a':
					opt_all = 1;
					break;
				case 'l':
					opt_lower = atoi(argv[++ i]);
					break;
				case 'u':
					opt_upper = atoi(argv[++ i]);
					break;
				case 'o':
					opt_pre = argv[++ i];
					if(strlen(opt_pre) > 100) {
						fprintf(stderr, "Error: prefix given by '-o' too long\r\n", argv[0]);
						return 1;
					}
					break;
				case 's':
					opt_seek = atol(argv[++ i]);
					break;
				case 'v':
					be_verbose = 1;
					break;
				case 'h':
					printf("Usage: %s [parameter] disk-image\r\n", argv[0]);
					printf("Parameter:\r\n");
					printf("    -o p  prepend given p to every written image, defaults to %s\r\n", opt_pre);
					printf("    -s b  skip b bytes before searching\r\n");
					printf("    -a    save all matching JPEG header and endings\r\n");
					printf("          (!) requires much space and manual selection\r\n");
					printf("              default is to save first matching\r\n");
					printf("    -l b  match images greater than b bytes, defaults to %d\r\n", opt_lower);
					printf("    -u b  match images smaller than b bytes, defaults to %d\r\n", opt_upper);
					printf("    -v    be verbose\r\n");
					printf("    -h    display this help\r\n");
					printf("\r\n");
					printf("Example:\r\n");
					printf("    %s -o out/RIMG_ -s 10000000 /tmp/defect.dd\r\n", argv[0]);
					printf("\r\n");
					printf("For more info, please visit: http://gimi.name/go/qiwRE\r\n");
					printf("\r\n");
					return 0;
			}
		else
			arg_img = argv[i];

	/* check args */
	if(arg_img == NULL) {
		fprintf(stderr, "Error: missing argument. for help, check '%s -h'\r\n", argv[0]);
		return 1;
	}

	/* open file */
	f = fopen(arg_img, "rb");
	if(f == NULL) {
		perror("Error: Cannot open file");
		return 1;
	}

	/* seek */
	fseek(f, opt_seek, SEEK_SET);

	/* search */
	while(!feof(f)) {
		/* search header */
		if(header_match(f)) {
			start = ftell(f) - 1;
			VERBOSE printf("HEADER at %ld\r\n", start);
			count = 0;
		}

		/* search end, save file */
		if(start && end_match(f)) {
			end = ftell(f);
			VERBOSE printf("END at %ld\r\n", end);

			if(end - start > opt_upper || end - start < opt_lower)
				/* size not matching */
				continue;

			if(!opt_all && count)
				/* already saved */
				continue;

			sprintf(buf, "%s%012ld_%02d.JPG", opt_pre, start, ++ count);
			if(!image_copy(f, start, end, buf))
				return 1;
			
			/* TODO: test this; prevent unneccessary checks for END */
			start = 0;
		}
	}

	fclose(f);
	return 0;
}