/*
 * The Sleuth Kit
 *
 * Brian Carrier [carrier <at> sleuthkit [dot] org]
 * Copyright (c) 2004 Brian Carrier.  All rights reserved
 *
 * sigfind
 *
 * This software is distributed under the Common Public License 1.0
 */

#include "tsk/tsk_tools_i.h"

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

static char *progname;

int     opterr = 1,             /* if error message should be printed */
optind = 1,             /* index into parent argv vector */
optopt,                 /* character checked for validity */
optreset;               /* reset getopt */
char    *optarg;                /* argument associated with option */

#define BADCH   (int)'?'
#define BADARG  (int)':'
#define EMSG    ""


void printSectorMd5(TSK_MD5_CTX *md){
	
	unsigned char hash[16];

	TSK_MD5_Final(hash, md);
	printf("MD5: ");
	for (int x = 0; x < 16; x++) {
		printf("%x%x", (hash[x] >> 4) & 0xf, hash[x] & 0xf);
	}
	printf("\n");

}


uint8_t* changeEndian(uint8_t* sig, int sig_size){
	uint8_t tmp;

	if (sig_size == 2) {
		tmp = sig[1];
		sig[1] = sig[0];
		sig[0] = tmp;
	}
	else if (sig_size == 3) {
		tmp = sig[2];
		sig[2] = sig[0];
		sig[0] = tmp;
	}
	else if (sig_size == 4) {
		tmp = sig[3];
		sig[3] = sig[0];
		sig[0] = tmp;

		tmp = sig[2];
		sig[2] = sig[1];
		sig[1] = tmp;
	}
	return sig;
}

/*
* getopt --
*      Parse argc/argv argument vector.
*/
int
getopt(int nargc, char * const nargv[], const char *ostr)
{
	static char *place = EMSG;              /* option letter processing */
	const char *oli;                        /* option letter list index */

	if (optreset || !*place) {              /* update scanning pointer */
		optreset = 0;
		if (optind >= nargc || *(place = nargv[optind]) != '-') {
			place = EMSG;
			return (-1);
		}
		if (place[1] && *++place == '-') {      /* found "--" */
			++optind;
			place = EMSG;
			return (-1);
		}
	}                                       /* option letter okay? */
	if ((optopt = (int)*place++) == (int)':' ||
		!(oli = strchr(ostr, optopt))) {
		/*
		* if the user didn't specify '-' as an option,
		* assume it means -1.
		*/
		if (optopt == (int)'-')
			return (-1);
		if (!*place)
			++optind;
		if (opterr && *ostr != ':')
			(void)printf("illegal option -- %c\n", optopt);
		return (BADCH);
	}
	if (*++oli != ':') {                    /* don't need argument */
		optarg = NULL;
		if (!*place)
			++optind;
	}
	else {                                  /* need an argument */
		if (*place)                     /* no white space */
			optarg = place;
		else if (nargc <= ++optind) {   /* no arg */
			place = EMSG;
			if (*ostr == ':')
				return (BADARG);
			if (opterr)
				(void)printf("option requires an argument -- %c\n", optopt);
			return (BADCH);
		}
		else                            /* white space */
			optarg = nargv[optind];
		place = EMSG;
		++optind;
	}
	return (optopt);                        /* dump back option letter */
}

void
usage()
{
    fprintf(stderr,
            "%s [-b bsize] [-o offset] [-t template] [-lV] [hex_signature] file\n",
            progname);
    fprintf(stderr, "\t-b bsize: Give block size (default 512)\n");
    fprintf(stderr,
            "\t-o offset: Give offset into block where signature should exist (default 0)\n");
    fprintf(stderr, "\t-l: Signature will be little endian in image\n");
    fprintf(stderr, "\t-V: Version\n");
	fprintf(stderr, "\t-r: rsize: Give the block size to read (default 512)\n");
    fprintf(stderr,
            "\t-t template: The name of a data structure template:\n");
    fprintf(stderr,
            "\t\tdospart, ext2, ext3, ext4, fat, hfs, hfs+, ufs1, ufs2\n");
	fprintf(stderr,
		"\t\t * New templates: gpt, fat32, fat12or16, exfat, ntfs, refs, hfsx, apfs\n\n");
	fprintf(stderr,
		"Modified version running on Windows, also included more advanced Windows FS templates\n");
	fprintf(stderr,
		"Modified by Rune Nordvik - 2015/12/17\n");
    exit(1);
}

// @@@ Should have a big endian flag as well
int
main(int argc, char **argv)
{
    int ch;
    uint8_t sig[4] = { 0, 0, 0, 0 };
    uint8_t sig2[4] = { 0, 0, 0, 0 }; // Another signature
    uint8_t block[4096];

    char **err = NULL;
    TSK_IMG_INFO *img_info;
    TSK_OFF_T cur_offset;
    int sig_offset = 0, rel_offset = 0;
    int read_size=512, bs = 512;
    int sig2_offset = 0, rel2_offset = 0; // Added offsets for sig 2
	bool hasHits = false; // Since prec_hit is unsigned we can not use -1 for it
	TSK_DADDR_T i, prev_hit; // Must be unsigned because of the large number of sectors
	TSK_DADDR_T * sectors64bitInVolume = NULL;
    int sig_size = 0;
    int sig2_size = 0;
    uint8_t lit_end = 0;
    int sig_print = 0, sig2_print = 0;
	TSK_MD5_CTX md;
	


    progname = argv[0];

    while ((ch = getopt(argc, argv, "b:r:lo:t:V")) > 0) {
        switch (ch) {
        case 'b':
            bs = strtol(optarg, err, 10);
            if ((bs == 0) || (errno == EINVAL)) {
                fprintf(stderr, "Error converting block size: %s\n",
                        optarg);
                exit(1);
            }

            if (bs % 512) {
                fprintf(stderr, "Invalid block size\n");
                exit(1);
            }
            break;
        case 'l':
            lit_end = 1;
            break;

        case 'o':

            /* Get the sig_offset in the sector */
            sig_offset = strtol(optarg, err, 10);
            if ((sig_offset == 0) || (errno == EINVAL)) {
                fprintf(stderr, "Error converting offset value: %s\n",
                        optarg);
                exit(1);
            }
            break;

        case 't':
            if ((strcmp(optarg, "ext2") == 0) ||
                (strcmp(optarg, "ext3") == 0) ||
                (strcmp(optarg, "ext4") == 0)) {
                lit_end = 1;
                sig[0] = 0x53;
                sig[1] = 0xef;
                sig_size = 2;
                sig_offset = 56;
                bs = 512;
            }
            else if (strcmp(optarg, "apfs") == 0) {
                lit_end = 1;
                sig[0] = 0x4E; // N
                sig[1] = 0x58; // X
                sig[2] = 0x53; // S
                sig[3] = 0x42; // B
                sig_size = 4;
                sig_offset = 32;
                bs = 512;
            }
            else if ((strcmp(optarg, "dospart") == 0) ||
                     (strcmp(optarg, "fat") == 0)) {
                lit_end = 1;
                sig[0] = 0x55;
                sig[1] = 0xaa;
                sig_size = 2;
                sig_offset = 510;
                bs = 512;
            }
            else if (strcmp(optarg, "gpt") == 0) {
                lit_end = 1;
                sig[0] = 0x55;
                sig[1] = 0xaa;
                sig2[0] = 0x45; // E
                sig2[1] = 0x46; // F
                sig2[2] = 0x49; // I
                sig2[3] = 0x20; //
                sig_size = 2;
                sig2_size = 4;
                sig_offset = 510;
                sig2_offset = 0;
                bs = 512;

            }
            else if (strcmp(optarg, "ntfs") == 0) {
                lit_end = 1;
                sig[0] = 0x55;
                sig[1] = 0xaa;
                sig2[0] = 0x4e; // N
                sig2[1] = 0x54; // T
                sig2[2] = 0x46; // F
                sig2[3] = 0x53; // S
                sig_size = 2;
                sig2_size = 4;
                sig_offset = 510;
                sig2_offset = 3;
                bs = 512;

            }
			else if (strcmp(optarg, "fat32") == 0) {
				lit_end = 1;
				sig[0] = 0x55;
				sig[1] = 0xaa;
				sig2[0] = 0x46; // F
				sig2[1] = 0x41; // A
				sig2[2] = 0x54; // T
				sig2[3] = 0x33; // 3
				sig_size = 2;
				sig2_size = 4;
				sig_offset = 510;
				sig2_offset = 82;
				bs = 512;
			}
			else if (strcmp(optarg, "fat12or16") == 0) {
				lit_end = 1;
				sig[0] = 0x55;
				sig[1] = 0xaa;
				sig2[0] = 0x46; // F
				sig2[1] = 0x41; // A
				sig2[2] = 0x54; // T
				sig2[3] = 0x31; // 1
				sig_size = 2;
				sig2_size = 4;
				sig_offset = 510;
				sig2_offset = 54;
				bs = 512;
			}
			else if (strcmp(optarg, "exfat") == 0) {
				lit_end = 1;
				sig[0] = 0x55;
				sig[1] = 0xaa;
				sig2[0] = 0x45; // E
				sig2[1] = 0x58; // X
				sig2[2] = 0x46; // F
				sig2[3] = 0x41; // A
				sig_size = 2;
				sig2_size = 4;
				sig_offset = 510;
				sig2_offset = 3;
				bs = 512;
			}
			else if (strcmp(optarg, "refs") == 0) {
				lit_end = 1;
				sig[0] = 0x52; // R
				sig[1] = 0x65; // e
				sig[2] = 0x46; // F
				sig[3] = 0x53; // S
				sig2[0] = 0x46; // F(ile)
				sig2[1] = 0x53; // S(ystem)
				sig2[2] = 0x52; // R(ecognition)
				sig2[3] = 0x53; // S(tructure)
				sig_size = 4;
				sig2_size = 4;
				sig_offset = 3;
				sig2_offset = 16;
				bs = 512;
			}
            else if (strcmp(optarg, "ufs1") == 0) {
                lit_end = 1;
                sig[0] = 0x54;
                sig[1] = 0x19;
                sig[2] = 0x01;
                sig[3] = 0x00;
                sig_size = 4;
                /* Located 1372 into SB */
                sig_offset = 348;
                bs = 512;
            }
            else if (strcmp(optarg, "ufs2") == 0) {
                lit_end = 1;
                sig[0] = 0x19;
                sig[1] = 0x01;
                sig[2] = 0x54;
                sig[3] = 0x19;
                sig_size = 4;
                /* Located 1372 into SB */
                sig_offset = 348;
                bs = 512;
            }
            else if (strcmp(optarg, "hfs+") == 0) {
                lit_end = 1;
                sig[0] = 0x48;
                sig[1] = 0x2b;
                sig[2] = 0x00;
                sig[3] = 0x04;
                sig_size = 4;
                /* Located 1024 into image */
                sig_offset = 0;
                bs = 512;
            }
			else if (strcmp(optarg, "hfsx") == 0) {
				lit_end = 1;
				sig[0] = 0x48; // H
				sig[1] = 0x58; // X
				sig[2] = 0x00;
				sig[3] = 0x05; // 5 - Version for HFSX
				sig_size = 4;
				/* Located 1024 into image */
				sig_offset = 0;
				bs = 512;
			}
            else if (strcmp(optarg, "hfs") == 0) {
                lit_end = 1;
                sig[0] = 0x42;
                sig[1] = 0x44;
                sig_size = 2;
                /* Located 1024 into image */
                sig_offset = 0;
                bs = 512;
            }
            else {
                fprintf(stderr, "Invalid template\n");
                exit(1);
            }
            break;

        case 'V':
            tsk_version_print(stdout);
            exit(0);
		case  'r':
			// Size of block to read
			read_size = strtol(optarg, err, 10);
			if ((read_size == 0) || (errno == EINVAL)) {
				fprintf(stderr, "Error converting block size: %s\n",
					optarg);
				exit(1);
			}
			break;
        default:
            usage();
        }
    }


    /* If we didn't get a template then check the cmd line */
    if (sig_size == 0) {
        if (optind + 1 > argc) {
            usage();
        }
        /* Get the hex value */
        sig_size = 0;
        for (i = 0; i < 9; i++) {
            uint8_t tmp;
            tmp = argv[optind][i];

            if (tmp == 0) {
                if (i % 2) {
                    fprintf(stderr,
                            "Invaild signature - full bytes only\n");
                    exit(1);
                }
                break;
            }

            /* Digit */
            if ((tmp >= 0x30) && (tmp <= 0x39)) {
                tmp -= 0x30;
            }
            /* lowercase a-f */
            else if ((tmp >= 0x61) && (tmp <= 0x66)) {
                tmp -= 0x57;
            }
            else if ((tmp >= 0x41) && (tmp <= 0x46)) {
                tmp -= 0x37;
            }
            else {
                fprintf(stderr, "Invalid signature value: %c\n", tmp);
                exit(1);
            }

            /* big nibble */
            if (0 == (i % 2)) {
                sig[sig_size] = 16 * tmp;
            }
            else {
                sig[sig_size] += tmp;
                sig_size++;
            }
        }
        optind++;

        /* Check the signature length */
        if (i == 9) {
            fprintf(stderr,
                    "Error: Maximum supported signature size is 4 bytes\n");
            exit(1);
        }


        /* Need to switch order */
        if (lit_end) {
            // call change endian order
			uint8_t* pointerToSig = 0;
			pointerToSig = changeEndian(sig, sig_size);
			for (int i = 0; i < sig_size; i++)
				sig[i] = pointerToSig[i];
			if (0 != sig2[0]){
				pointerToSig = changeEndian(sig2, sig2_size);
				for (int i = 0; i < sig2_size; i++)
					sig2[i] = pointerToSig[i];
			}
        }
    }

    if (sig_offset < 0) {
        fprintf(stderr, "Error: negative signature offset\n");
        exit(1);
    }


    /* Check that the signature and offset are not larger than a block */
    if ((sig_offset + sig_size) > bs) {
        fprintf(stderr,
                "Error: The offset and signature sizes are greater than the block size\n");
        exit(1);
    }

//    read_size = 512;
    /* If our signature crosses the 512 boundary, then read 1k at a time */
    if ((sig_offset / 512) != ((sig_offset + sig_size - 1) / 512)) {
        read_size = 4096;
    }

    /* Get the image */
    if (optind + 1 != argc) {
        usage();
    }

    if ((img_info =
         tsk_img_open_utf8_sing(argv[optind],
                      TSK_IMG_TYPE_DETECT, 0)) == NULL) {
        tsk_error_print(stderr);
        exit(1);
    }

    /* Make a version that can be more easily printed */
    for (i = 0; i < sig_size; i++) {
        sig_print |= (sig[i] << ((sig_size - 1 - i) * 8));
    }

	for (i = 0; i < sig2_size; i++) {
		sig2_print |= (sig2[i] << ((sig2_size - 1 - i) * 8));
	}

	if (0 != sig2[0]){
		printf("Block size: %d\nOffset1: %d  Signature1: %X\nOffset2: %d  Signature2: %X\n\n", bs, sig_offset,
			sig_print, sig2_offset, sig2_print);

	}
	else{
		printf("Block size: %d\nOffset: %d  Signature: %X\n\n", bs, sig_offset,
			sig_print);
	}
    /* Loop through by blocks  - we will read in block sized chunks
     * so that we can be used on raw devices 
     */
    cur_offset = (sig_offset / 512) * 512;
    rel_offset = sig_offset % 512;
    rel2_offset = sig2_offset % 512;
    //prev_hit = -1;
    for (i = 0;; i++) {
        ssize_t retval;

        /* Read the signature area */
        retval = tsk_img_read(img_info, cur_offset,
                                    (char *)block, read_size);
        if (retval == 0) {
            break;
        }
        else if (retval == -1) {
            fprintf(stderr, "error reading bytes %lu\n",
                    (unsigned long) i);
            exit(1);
        }

        /* Check the sig */
        if ((block[rel_offset] == sig[0]) &&
            ((sig_size < 2) || (block[rel_offset + 1] == sig[1])) &&
            ((sig_size < 3) || (block[rel_offset + 2] == sig[2])) &&
            ((sig_size < 4) || (block[rel_offset + 3] == sig[3]))) {
            if( 0 == sig2[0]){ // Only one signature
				if (! hasHits){
					printf("Block: %lu (-)\t", (unsigned long)i);
					hasHits = !hasHits;
				}
				else{
					printf("Block: %lu (+%lu)\t", (unsigned long)i,
						(unsigned long)(i - prev_hit));
				} // inner if else
				
				prev_hit = i; // only update if a hit
				TSK_MD5_Init(&md);
				TSK_MD5_Update(&md, block, bs);
				printSectorMd5(&md); // Prints the md5 hash value of current sector
				
				
            }else{ // We have a signature number two
                if ((block[rel2_offset] == sig2[0]) &&
                    ((sig2_size < 2) || (block[rel2_offset + 1] == sig2[1])) &&
                    ((sig2_size < 3) || (block[rel2_offset + 2] == sig2[2])) &&
                    ((sig2_size < 4) || (block[rel2_offset + 3] == sig2[3]))) {
					if (!hasHits){
						printf("Block: %lu (-)\t", (unsigned long)i);
						hasHits = !hasHits;
					}
					else{
						printf("Block: %lu (+%lu)\t", (unsigned long)i,
							(unsigned long)(i - prev_hit));
					} // inner else if
					
					prev_hit = i; // only update if a hit
					TSK_MD5_Init(&md);
					TSK_MD5_Update(&md, block, bs);
					printSectorMd5(&md); // Prints the md5 hash value of current sector	
					// Perhaps count the sectors in volume
					//sectors64bitInVolume = (TSK_DADDR_T*)&block[40];
					//printf("Sectors in Volume: %lu\t", (unsigned long)*sectors64bitInVolume);
					//sectors64bitInVolume = NULL; // Best practice
                } // inner if
				
				
            } // if else		
            
        }
        cur_offset += bs;
    }

    tsk_img_close(img_info);
    exit(0);
}