/*
  Copyright 2005 Richard C. Perrin <rcp@sentientmeat.ca>

  jpegextr: extract JPEG files from a byte stream

  Useful for recovering deleted files from a disk image. Writes
  numbered files to the current working directory.  If you have
  deleted files from a disk (or more likely, a flash card) and are
  unable to recover them by more simple means, then use dd and this
  program to attempt to recover them:

  dd < /dev/disk-dev | jpegextr

  or:

  dd < /dev/disk-dev | bzip2 -c > disk_image.bz2
  jpegextr disk_image.bz2
  
  Features:
  - supports reading directly from .gz and .bz2 files

  Bugs:
  - only Exif format JPEG files are currently supported
  - we make no attempt to read the proper file length - writing will
  continue until the start of a new JPEG file is matched.
  - no command line options available (for help or for output filename
  prefix)
  - no status output available
  - horribly inefficient; but really, do you need to run this often?
  - more..., see FIXME(s) below

  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., 59 Temple Place, Suite 330, Boston, MA
  02111-1307 USA
*/

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <stdio.h>
#include <limits.h>
#include <string.h>
#include <assert.h>

#define NELEMS(array) (sizeof(array) / sizeof(array[0]))
#define PREFIX "FJPG"		/* output filename prefix */

void fatal(char * s)
{
     fprintf(stderr, "%s\n", s);
     exit(-1);
}
void fatalerr(char * s)
{
     fprintf(stderr, "%s: %s\n", s, strerror(errno));
     exit(-1);
}

enum states { Lk_1st, Lk_2nd, Lk_3rd, MidFile };
// Exif
// shorts: d8ff e1ff fc8b 7845 6669
// bytes: ff d8 ff e1 8b fc 45 78 69 66
// JFIF is not handled

#define NUM_SHORTS 5
#define S1st 0xd8ff
#define S2nd 0x7845
#define S3rd 0x6669

static int outfd = -1;

void newfile(int n)
{
     char name[PATH_MAX];
     snprintf(name, sizeof(name), "%s%04d.jpg", PREFIX, n);
     if (outfd >= 0)
	  assert(!close(outfd));

     outfd = open(name, O_WRONLY|O_CREAT|O_EXCL, 0666);
     if (outfd < 0)
	  fatalerr(name);
}

void write_bytes(const void *buf, size_t count)
{
     if (outfd >= 0)
	  if (write(outfd, buf, count) < count)
	       fatalerr("write");
}

void purge_buf(unsigned short buf[], int *n)
{
     unsigned short x;
     int i;
	
     for (i = 0; i < *n; i++) {
	  x = buf[i];
	  write_bytes(&x, sizeof(x));
     }
     *n = 0;
}

void extract(int fd, char *name)
{
     unsigned short x;
     unsigned char c, c2;
     enum states state = Lk_1st, prev = Lk_1st;
     int cnt = 0;
     unsigned short buf[NUM_SHORTS];
     int buf_n = 0;
     int skip;
     ssize_t sz = 0;

     do {
	  assert(buf_n < NELEMS(buf));

	  switch (state) {
	  case Lk_1st:
	  case MidFile:
	       /* byte I/O */
	       if ((sz = read(fd, &c, sizeof(c))) > 0) {
		    if (c == 0xff &&
			(sz = read(fd, &c2, sizeof(c2))) == sizeof(c2)) {
			 if (c2 == 0xd8) { /* potential file header */
			      x = S1st;
			 } else {
			      write_bytes(&c, sizeof(c));
			      write_bytes(&c2, sizeof(c2));
			      continue;
			 }
		    } else {
			 write_bytes(&c, sizeof(c));
			 continue;
		    }
	       }
	       break;
	  case Lk_2nd:
	  case Lk_3rd:
	       /* short I/O */
	       if ((sz = read(fd, &x, sizeof(x))) > 0) {
		    if (sz < sizeof(x)) { /* short read */
			 purge_buf(buf, &buf_n);
			 write_bytes(&x, sz);
			 break;
		    }
	       }
	       break;
	  default:
	       fatal("unknown state\n");
	  }

	  if (sz <= 0)
	       continue; /* FIXME: crufty, the two separate switches
			  * should go back into one */

	  switch (state) {
	  case Lk_1st:
	       if (x == S1st) {
		    // buffer and go to next state
		    buf[buf_n++] = x;
		    prev = state;
		    skip = 2;
		    state = Lk_2nd;
	       }
	       break;
	  case Lk_2nd:
	       buf[buf_n++] = x;
	       if (skip) {
		    skip--;
	       } else {
		    if (x == S2nd) {
			 // buffer and go to next state
			 state = Lk_3rd;
		    } else {
			 // deal with the buffered bits and return to
			 // previous state
			 purge_buf(buf, &buf_n);
			 state = prev;
		    }
	       }
	       break;
	  case Lk_3rd:
	       buf[buf_n++] = x;
	       if (x == S3rd) {
		    // start new outfile, deal with the
		    // buffered bits
		    newfile(++cnt);
		    state = MidFile;
	       } else {
		    // deal with the buffered bits and return to
		    // previous state
		    state = prev;
	       }
	       purge_buf(buf, &buf_n);
	       break;
	  case MidFile:
	       buf[buf_n++] = x;
	       if (x == S1st) {
		    // buffer and go to next state
		    prev = state;
		    skip = 2;
		    state = Lk_2nd;
	       } else {
		    // continue output
		    purge_buf(buf, &buf_n);
	       }
	       break;
	  default:
	       fatal("unknown state\n");

	  }
     } while (sz > 0);
}

__inline__ int matchext(const char *file, const char *ext) {
     assert(ext);
     return !strcmp(file + (strlen(file) - strlen(ext)), ext);
}

int main(int argc, char * argv[])
{
     int i, fd;

     /* args are expected to be filenames, none means use stdin */
     if (argc == 1) {
	  extract(STDIN_FILENO, "-");
	  return 0;
     }
     for (i = 1; i < argc; i++) {
	  FILE *fp = NULL;
	  char *prog = NULL;
	  if (matchext(argv[i], ".bz2"))
	       prog = "bzip2";
	  else if (matchext(argv[i], ".gz"))
	       prog = "gzip";

	  if (prog) {
	       char cmd[PATH_MAX];
	       snprintf(cmd, sizeof(cmd), "%s -dc %s", prog, argv[i]);
	       fp = popen(cmd, "r");
	       if (!fp) {
		    perror(cmd);
		    continue;
	       }
	       fd = fileno(fp);
	  } else {
	       fd = open(argv[i], O_RDONLY);
	  }
	  if (fd < 0) {
	       perror(argv[i]);
	       continue;
	  }

	  extract(fd, argv[i]);

	  if (fp)
	       pclose(fp);
	  else
	       close(fd);
     }
     return 0;
}