/* dig.c
 *
 * This is a work of the US Government. In accordance with 17 USC 105,
 * copyright protection is not available for any work of the US Government.
 *
 * 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.
 *
 */

#include "foremost.h"

void clean_up (struct foremostState* state, int signum) {

  foremostLog (state,"Cleaning up...\n");
  foremostLog (state,
               "\nCaught signal: %s. Program is terminating early\n",
               (char*) strsignal(signum));
  if (closeFile(state->auditFile)) {
    foremostLog(state,"Error closing %s/audit.txt -- %s",
                state->outputdirectory,
                (char*) strerror(ferror(state->auditFile)));
  }
  exit(1);
}


/* display Position: Tell the user how far through the infile we are */
int displayPosition(int *units,
		    unsigned long long pos,
		    unsigned long long size, 
		    char *fn) {
  
  double percentDone = (((double)pos)/(double)(size) * 100);
  double position = (double) pos;
  int count;
  int barlength,i,len;
  double elapsed;
  long remaining;

  /* We don't use MAX_STRING_LENGTH because we're just printing the
     units of how much of the file we've read. At worst case, this is
     the string "bytes\0" == 6 characters. */

  char buf[7];
  char line[256];

#ifdef __WIN32
  static LARGE_INTEGER start;
  LARGE_INTEGER now;
  static LARGE_INTEGER freq;
  QueryPerformanceFrequency(&freq);
#else
  static struct timeval start;
  struct timeval now, td;
#endif



  /* First we need to know what time it is. If we're processing the first
     chunk of the file, we should also set the  starting time too. */

  if(pos <= SIZE_OF_BUFFER){ 
    (void)gettimeofday(&start, (struct timezone *)0);
  }
  (void)gettimeofday(&now, (struct timezone *)0);

  /* First, reduce the position to the right units */
  for (count = 0; count < *units; count++) {
    position = position / 1024;
  }
  
  /* Now check if we've hit the next type of units */
  while (position > 1023) {
    position = position / 1024;
    (*units)++;
  }
		       
  switch (*units) {

  case UNITS_BYTES:
    sprintf(buf,"bytes");    break;
  case UNITS_KILOB:
    sprintf(buf,"KB");       break;
  case UNITS_MEGAB:
    sprintf(buf,"MB");       break;
  case UNITS_GIGAB:
    sprintf(buf,"GB");       break;
  case UNITS_TERAB:
    sprintf(buf,"TB");       break;
  case UNITS_PETAB:
    sprintf(buf,"PB");       break;
  case UNITS_EXAB:
    sprintf(buf,"EB");       break;

  default:
    
    /* Steinbach's Guideline for Systems Programming:
         Never test for an error condition you don't know how to handle.
	 
       Because we're going to hit this time and time again, we should
       humbly admit our failure at this point. We might want to be silly
       and chastise the user for trying to process an 1025+ Exabyte file 
       with this program, but we'll take the moral high ground. */
    
    fprintf (stdout, "Unable to compute progress.\n");
    return FOREMOST_OK;
  }
  len = 0;
  len += snprintf(line+len,sizeof(line)-len,"\r%s: %5.1f%% ",fn, percentDone);
  barlength = ttywidth - strlen(fn) - strlen(buf) - 32;
   if(barlength>0){
     i = barlength * (int) percentDone / 100;
     len += snprintf(line+len, sizeof(line)-len,
    		    "|%.*s%*s|", i,
    "****************************************************************************************************************************************************************",
		    barlength-i, "");
  }
    
  len += snprintf(line+len,sizeof(line)-len," %6.1f %s",position,buf);

#ifdef __WIN32
  elapsed = ((double)now.QuadPart - (double)start.QuadPart)/((double)freq.QuadPart);
  //printf("elapsed: %f\n",elapsed);
#else
  timersub(&now, &start, &td);
  elapsed = td.tv_sec + (td.tv_usec / 1000000.0);
#endif
  remaining = (100-percentDone)/percentDone*elapsed;
  //printf("Ratio remaining: %f\n",(100-percentDone)/percentDone);
  //printf("Elapsed time: %f\n",elapsed);
  if(remaining >= 100*(60*60)){  //60*60 is seconds per hour
    len +=snprintf(line+len, sizeof(line)-len," --:--ETA");
  }else{
    i = remaining / (60*60);
    if(i)
      len += snprintf(line+len,sizeof(line)-len," %2d:",i);
    else
      len += snprintf(line+len,sizeof(line)-len,"    ");
    i = remaining%(60*60);
    len += snprintf(line+len,sizeof(line)-len,"%02d:%02d ETA",i/60, i%60);
  }
  

  fprintf(stdout,"%s",line);
  fflush(stdout);

  return FOREMOST_OK;
}


/* Creates initial entries in the audit file for each image file */
void setupAuditFile(struct foremostState* state) {
  
  char imageFile[MAX_STRING_LENGTH];  

  realpath(state->imagefile,imageFile);
  
  foremostLog(state,"\n\nOpening %s\n", imageFile);

  if (state->skip) {
    fprintf(state->auditFile,"Skipped the first %lld bytes of %s...\n",
	    state->skip, state->imagefile);
    if (state->modeVerbose) {
      fprintf(stdout,"Skipped the first %lld bytes of %s...\n",
	      state->skip, state->imagefile);
    }
  }
  
  fprintf(state->auditFile,
	  "File            Found at Byte  Int  Chop   Length     Extracted From\n");
}




/*
  digBuffer: looks for any of our tags within a buffer and extracts and writes 
             them to disk...
  Args:
  state           -- The foremost state object
  buf             -- the buffer to dig
  infile          -- the file we are reading
  maxchars        -- size of the largest extraction we would make
  lengthofbuf     -- length of buf
  offset          -- where does buf start within the infile
*/

int bm_digBuffer(struct foremostState* state, char* buf, FILE* infile, 
		 int maxchars, unsigned long long lengthofbuf, 
		 unsigned long long offset) {
  
  unsigned long long startLocation = 0;
  int needlenum;
  char* foundat;
  char* current;
  struct CharBucket extractbuf = {0,NULL};
  struct SearchSpecLine currentneedle;
    
  extractbuf.str = (char*) malloc(maxchars*sizeof(char));
  
  /* Loop through each needle, calling needleinhaystack to find all of those needles */
  for(needlenum=0;state->SearchSpec[needlenum].suffix != NULL; needlenum++){
    foundat = buf;
    currentneedle = state->SearchSpec[needlenum];
    while(foundat){
      /*  First check to make sure we haven't caught a signal */
      if (signal_caught == SIGTERM || signal_caught == SIGINT){
	clean_up(state,signal_caught);
      }

      if(state->modeQuick){
	//Only search every 512 bytes in the buffer
	current=foundat;
	foundat = NULL;
	while(current<=buf+lengthofbuf-FOREMOST_BLOCK_SIZE){
	  if(!memwildcardcmp(currentneedle.begin,
			     current,
			     currentneedle.beginlength,
			     currentneedle.casesensitive)){
	    //We have a match for our quick search
	    foundat = current;
	    break;
	  }else{	
	    current += FOREMOST_BLOCK_SIZE;
	  }
	}
      }
      else{
	//Not in quick mode, Search the whole buffer
	foundat = bm_needleinhaystack(currentneedle.begin, 
				      currentneedle.beginlength,
				      foundat,
				      lengthofbuf-(foundat-buf),
				      currentneedle.begin_bm_table,
				      currentneedle.casesensitive,
				      currentneedle.searchtype);
	
      }
      
      if(foundat != NULL && foundat >= 0){
		
	/*  We've found a header!
	    now we extract it into extractbuf and write to a file. */
	
	startLocation = offset + (foundat-buf);	

	/* Zero out the extraction buffer before we load new data into it */
	memset(extractbuf.str,'\x00',maxchars*sizeof(char));
	extractbuf.length = 0;
	
	/* The "startLocation" variable represents the position of the 'found' 
	   file inside of the image we're digging. Originally this
	   function had a "startLocation + 1" in the function call, which 
	   skewed the on-screen output. I've removed the +1  (JK) */
	
	if (state->modeVerbose) {
	  printf("A %s was found at: %lld\n",
		 currentneedle.suffix,startLocation);
	}
	
	extractString(&extractbuf,startLocation,infile,foundat,buf,lengthofbuf,currentneedle);
	
	if (writeToDisk(currentneedle.suffix,currentneedle.length,state,&extractbuf,startLocation)) {
	  return FOREMOST_ERROR_FILE_WRITE;	  
	}
	
	/* Skip forward past the header we just found and move on...
	   If we are in quick mode we skip to the next sector boundary */
	if (state->modeQuick){
	  foundat = foundat + FOREMOST_BLOCK_SIZE;
	}else{
	  foundat = foundat + currentneedle.beginlength;
	}	
      }
      
      
    }
  }      
  free(extractbuf.str);  
  return FOREMOST_OK;
}


/* digImageFile is the engine for the program. It uses the foremostState
   variable passed to it to find the image file to use, looks for the
   specified headers to find, and writes any files it finds to the disk
   Return values are defined in foremost.h                                */

int digImageFile(struct foremostState* state) {
  
  FILE* infile;
  unsigned long long filesize = 0, bytesread = 0, 
    fileposition = 0, filebegin = 0, beginreadpos = 0;
  long maxchars=0, err=0;
  int status, displayUnits = UNITS_BYTES;
  char *buf = (char*) malloc(SIZE_OF_BUFFER*sizeof(char));
  int retries=0,success=0;
  int longestneedle;
  setupAuditFile(state);
  
  if (state->SearchSpec[0].suffix == NULL) {
    free(buf);
    return FOREMOST_ERROR_NO_SEARCH_SPEC;
  }

  /* RBF - It would be nice if digImageFile was broken into multiple,
     RBF - smaller functions, but it's okay the way it is. */
  
  /* We need to know which needle snarfs the most data, so we can set 
     extractbuf.str to be that size */
  
  maxchars       = findLongestLength  (state->SearchSpec);
  longestneedle  = findLongestNeedle  (state->SearchSpec);

    
  /*  Now we're ready to open the image file and start digging */
  if ((infile = fopen(state->imagefile,"r")) == NULL) {
    fprintf(stderr, "ERROR: Couldn't open input file: %s -- %s\n", 
	    (*(state->imagefile)=='\0')?"<blank>":state->imagefile,
	    strerror(errno));
    free(buf);
    return FOREMOST_ERROR_FILE_OPEN;
  }

#ifdef __WIN32
  /*We need to EXPLICITLY open the file in binary mode for Win32
    this was very annoying to find out ;-)... */
   setmode(fileno(infile),O_BINARY);
#endif
#ifdef __LINUX
  fcntl(fileno(infile),F_SETFL, O_LARGEFILE);
#endif
  
  /* If we're skipping part of the input file, do it now. This function
     lives in helpers.c */
  if(state->skip > 0){
    if (!skipInFile(state,infile)) {
      free(buf);
      return FOREMOST_ERROR_FILE_READ;
    }
  }


  filebegin = ftello(infile);
  if ((filesize = measureOpenFile(infile)) == -1) {
    fprintf (stderr,
	     "ERROR: Couldn't measure size of image file %s\n", 
	     state->imagefile);
    free(buf);
    return FOREMOST_ERROR_FILE_READ;
  }
  
  if (state->modeVerbose) {
    fprintf (stdout, "Total file size is %lld bytes\n", filesize);
  }
  
  /* Now we can get started reading the image file.
     We read SIZE_OF_BUFFER bytes into memory and then
     look for headers    */   

  retries=0;  
  success=1;
  while((bytesread = fread(buf,1,
			   SIZE_OF_BUFFER,infile)) > longestneedle-1 || success==0){
    /* Check for read errors */
    if ((err = ferror(infile))) {
      if(retries<3){
	fprintf (stderr,
		 "Failed to read from position %lld -- %s\n",
		 (long long int)ftello(infile),strerror(err));
	fprintf(stderr,"Pausing 3 seconds to try again...\n"); 
	sleep(3);
	success=0;
	retries++;
	continue;
      }else{
	foremostLog (state,
		     "Maximum retries exceeded reading from position %lld -- %s.\nAborting.\n",
		     (long long int)ftello(infile),strerror(err));
	free(buf);
	return FOREMOST_ERROR_FILE_READ;      
      }
    }
    success=1;
    
    /* Find out where we are in the file and tell the user */
    fileposition = ftello(infile);
    displayPosition(&displayUnits,fileposition-filebegin,
		    filesize,state->imagefile);
    beginreadpos = fileposition - bytesread;
    
    /* Check to see if the program has been interrupted by the user */
    if (signal_caught == SIGTERM || signal_caught == SIGINT)
      clean_up(state,signal_caught);
    
    /* Let's dig through the current buffer and write the results to disk */
    if ((status = bm_digBuffer(state,buf,infile,
			    maxchars,bytesread,beginreadpos)) != FOREMOST_OK) {
    
      /* Uh oh. Something bad happened. I'm outta here. */
      free(buf);
      return status;
    }
    
    /* We seek back a little bit in the file so that we don't miss
       any headers that bridge the gap between buffers */
    fseeko(infile, ftello(infile) - (longestneedle - 1), SEEK_SET);
  }
  
  closeFile(infile);
  
  /* Clean up the memory we're using here */
  free(buf);
  return FOREMOST_OK;
}