/*
 * extract-jpeg 1.0, by Mike Hearn <mike@plan99.net>
 *
 * Given a binary file or disk image extract-jpeg will dump any
 * EXIF/JPEG files it can locate by brute force, ie it will scan the
 * entire input.  It was originally written to recover deleted
 * pictures from a digital camera.
 *
 * 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 <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>

typedef enum
{
    FALSE = 0,
    TRUE  = 1
} boolean;
    
#define chat(x...) printf("extract-jpeg: " x);

int signature_count = 0;

int dump(unsigned char *soi, unsigned char *cursor, boolean corrupt)
{
    char fn[256];
    memset(fn, 0, sizeof(fn));
    snprintf(fn, sizeof(fn), "recovered-%d.jpg%s", signature_count, corrupt ? ".corrupt" : "");

    int recovered_fd;
    if ((recovered_fd = open(fn, O_CREAT | O_RDWR | O_TRUNC, S_IREAD | S_IWRITE)) == -1)
    {
        chat("error creating file: %s\n", strerror(errno));
        return 1;
    }
    
    /* write it out */
    write(recovered_fd, soi, cursor - soi);
    close(recovered_fd);

    return 0;
}


int main(int argc, char *argv[])
{
    char *filename = argv[1];
    
    if (!filename || !strcmp(filename, "--help") || !strcmp(filename, "-h") ||
        !strcmp(filename, "--version") || !strcmp(filename, "-v"))
    {
        printf("Usage: extract-jpeg <FILENAME>\n\n");
        printf("Will extract all JPEG/EXIF files it can find by brute force\nand put them in the current directory\n");
        printf("FILENAME can be a device path, for instance to a FAT partition on a digital camera\n");
        printf("\n(C) 2005 Mike Hearn <mike@navi.cx>\n");
        return 0;
    }
    
    chat("extracting from %s\n", filename);

    int fd = open(filename, O_RDONLY);
    struct stat buf;
    stat(filename, &buf);
    chat("size is 0x%lx bytes\n", buf.st_size);
    
    unsigned char *start = mmap(NULL, buf.st_size, PROT_READ, MAP_SHARED, fd, 0);
    unsigned char *end = start + buf.st_size;
    
    if (!start)
    {
        chat("failed to map file: %s\n", strerror(errno));
        return 1;
    }

    unsigned char *cursor = start;
    double last_percentage = 0;
    boolean corrupted = FALSE;
    
    while (cursor < end - 4)
    {
        double completed = (double) (cursor - start) / (double) buf.st_size;

        if (completed - last_percentage > 0.1)
        {
            last_percentage = completed;
            chat("scanned %0.1f%% of file\n", completed * 100);
        }
        
        /* search for a JFIF+EXIF signature */
        if (*((unsigned int *)cursor) != 0xe1ffd8ff /* 0xffd8ffe1 */)
        {
            cursor += 1;
            continue;
        }

        signature_count++;
        chat("found JFIF/EXIF signature %d at offset 0x%x\n", signature_count, cursor - start);

        unsigned char *soi = cursor;
        cursor += 2;
                
        /* JFIF is really simple. there is the start of image marker
           0xffd8, then a series of marker:size:data triples. size includes
           its own count, ie data size is size - 2, as the size field is 16bit.

           once the SOS (start of stream) marker is encountered we just continue
           until we find EOI (end of image) FFD9.

           one catch is that it's all big-endian (motorola format)
         */


        /* now skip each marker as it's found */
        while (1)
        {
            unsigned short marker = *(unsigned short *) cursor;
            
            if ((marker & 0x00FF) != 0x00FF)
            {
                chat("PANIC: marker check failed at 0x%x: got 0x%x but expected 0xFF??\n", cursor - start, *(unsigned short *)cursor);
                chat("PANIC: dumping and restarting SOI scan\n");
                dump(soi, cursor, TRUE);
                cursor += 4;
                corrupted = TRUE;
                break;
            }            

            cursor += 2;            /* skip the marker */

            unsigned short size = (*cursor << 8) + *(cursor+1);
            cursor += size;         /* skip the size+data fields */

            if (marker == 0xDAFF)
            {
                /* start of stream marker */
                break;
            }
        }

        if (corrupted)
        {
            corrupted = FALSE;
            chat("PANIC: image %d could not be scanned, fragmentation? let's carry on ...\n", signature_count);
            continue;
        }
        
        chat("located start of image stream for signature %d at offsets 0x%x ... ", signature_count, cursor - start);
        
        while (1)
        {
            unsigned short val = *(unsigned short *) cursor;
            cursor += 1;
            if (val != 0xD9FF) continue;
            break;
        }

        cursor += 1;
        printf("0x%x\n", cursor - start);

        if (dump(soi, cursor, FALSE)) return 1;
    }
    
    return 0;
}