/* trimexe.c
 *
 * A small program for extracting PE files from small chunks
 * of raw data.  This is designed to fish executables out of the
 * reassembled raw data you can get from "follow TCP stream" in
 * wireshark, for example. 
 *
 * At the moment it finds the first PE file in the stream and uses
 * the PE coff and section header data to compute the file size
 * of the executable.  Output of various header values and offsets
 * is in hex and decimal to save a trip to the calculator.
 *
 * It then opens a temporary file in . and carves the file out 
 * of the image based on the offset of MZ and
 * the size. 
 *
 * Code is ugly but commented, probably not enough error checking,
 * but if fseek fails, you have bigger problems than I. 
 * 
 * Written for SANS network forensics puzzle contest #6
 * Copyright 2010 Candice Quates, license GPL.
 */

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

int
main(int argc, char *argv[]) {
    long mzoffset;
    long peoffset;
    long filesize;
    unsigned int sections;
    FILE *source;
    FILE *dest;
    unsigned int i,j,k,l;
    int destfd;
    char tmpname[14] = "./pefileXXXXX";
   
    if (argc==2) {
	printf("filename:  %s\n",argv[1]);
    } else {
	fprintf(stderr, "usage: %s filename", argv[0]);
	exit (1);
    }

    // open file
    source=fopen(argv[1],"r");
    if (source == NULL)  {
	fprintf(stderr,"error: %s, %s\n", argv[0],strerror(errno));
	exit(1);
    }  
    // search for MZ
    j=0;
    while (j!=EOF) {
	k=j;
	j=fgetc(source);
	if (j=='Z' && k=='M') {
	    mzoffset=ftell(source);
            mzoffset=mzoffset-2;
	    break;
	}
    }
    if (j==EOF) {
	fprintf(stderr,"I cannot find an executable file here.\n");
 	exit(1);
    }
     
    fprintf(stdout,"MZ located %d %x\n",mzoffset,mzoffset);
    // fseek(source,0x3c,SEEK_CUR); // offset to PEoffset
    // I should be chasing offsets here.  but this is easier.
    // search for PE 
    while (j!=EOF) {
	k=j;
	j=fgetc(source);
	if (j=='E' && k=='P') {
	    peoffset=ftell(source);
            peoffset=peoffset-2;
	    break;
	}
    }
    if (j==EOF) {
	fprintf(stderr,"I cannot find an executable file here.\n");
 	exit(1);
    }
    fprintf(stdout,"PE located %d %x\n",peoffset,peoffset);
    // the section count is a few bytes further. +2, but need +2 for 00
    fseek(source,4,SEEK_CUR);
    sections=fgetc(source); 
    fprintf(stdout,"Number of sections %d %x\n",sections,sections);
    // optional header size
    fseek(source,13,SEEK_CUR);
    i=fgetc(source); 
    fprintf(stdout,"Opt header size %d %x\n",i,i);
    // 2 more bytes in coff header, plus size of optional, should
    // put us in the section table.
    fseek(source,i+3,SEEK_CUR);
    // go down N sections. 
    fseek(source,40*(sections-1),SEEK_CUR);
    // size of last section
    fseek(source,16,SEEK_CUR);
    j=getw(source);
    fprintf(stdout,"Last section raw size %d %x\n",j,j);
    l=getw(source);
    fprintf(stdout,"Last section raw offset %d %x\n",l,l);
    // add together.  voila.
    filesize=l+j;   
    fprintf(stdout,"Size of file from MZ, %d %x\n",filesize,filesize);

    // now, to carve it, open new temp file
    destfd=mkstemp(tmpname);
    if (destfd==-1) {
	fprintf(stderr,"Error creating temp file %s: %s\n", 
		tmpname, strerror(errno));
	exit(1);
    }
    fprintf(stdout,"Output file: %s\n",tmpname);
    dest=fdopen(destfd,"w");
    // go back up to MZ
    fseek(source,mzoffset,SEEK_SET);
    // this is ugly.
    if ((filesize % 2) != 0) filesize--;
    for (i=0; i<filesize; i++) {
        j=fgetc(source);
        fputc(j,dest); 
    }
    fclose(dest); 
    fclose(source); 
    return 0;
}