/* helpers.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"

/* On non-glibc systems we need to compute our name. While this can
   ostensibly be found in argv[0], we clean up this value to make
   it look nicer. The argument s should be argv[0] passed from main(). */
#ifndef __GLIBC__
void setProgramName(char *s) {

  /* 's' should be argv[0]. Our goal is get the real filename for this
     value and then strip out everything before the last DIR_TRAIL_CHAR */
  int pos;
  char *fn = (char *)malloc(sizeof(char) * PATH_MAX);
  if (realpath(s,fn) == NULL) {
    __progname = strdup("foremost");
    return;
  }

  __progname = basename(fn);
}

#endif /* ifndef __GLIBC__ */



/* Write an entry to both the screen and the audit file */
void foremostLog(struct foremostState *state, char *format, ...) {

  va_list argp;

  va_start(argp,format);
  vfprintf (stdout,format,argp);
  vfprintf (state->auditFile,format,argp);
  va_end(argp);
}
  



int charactersMatch(char a, char b, int caseSensitive){

  if (a == wildcard || a == b) return 1;
  if (caseSensitive || (a < 'A' || a > 'z' || b < 'A' || b > 'z')) return 0;

  /* This line is equivalent to (abs(a-b)) == 'a' - 'A' */
  return (abs(a-b) == 32);
}


/*
 memwildcardcmp is exactly like memcmp, except a wilcard 
 (usually '?' but can be redefined in the configuration file)  
 in s1 will match any single character in s2.
 We added it to support wildcards in the search specification.
 Be *very* careful about the order of s1 and s2... We want for a 
 wildcard in s1 to match anything, but not the other way around.
*/

int memwildcardcmp(const void *s1, const void *s2, 
		   size_t n,int caseSensitive){
  if (n!=0){
    register const unsigned char *p1 = s1, *p2 = s2;
    do{
      if(!charactersMatch(*p1++, *p2++, caseSensitive))
	return (*--p1 - *--p2);
    } while (--n !=0);
  }
  return(0);
}

/*
  init_search initializes a Boyer-Moore "jump table" for whatever "needle" you pass in
  After running, table will be loaded with the appropriate values for use by the boyer-moore search algorithm
  
*/

void init_bm_table(char *needle, size_t table[UCHAR_MAX + 1], size_t len, int casesensitive,int searchtype)
{
      size_t i=0,j=0,currentindex=0;

      for (i = 0; i <= UCHAR_MAX; i++)                   
	table[i] = len;
      for (i = 0; i < len; i++){
	if(searchtype == SEARCHTYPE_REVERSE){
	  currentindex = i; //If we are running our searches backwards 
	                      //we count from the beginning of the string
	}else{
	  currentindex = len-i-1; //Count from the back of string
	}

		  if(needle[i] == wildcard){           //No skip entry can advance us past the last wildcard in the string
		    for(j=0; j<=UCHAR_MAX; j++)
		      table[j] = currentindex; 
		  }	
	      table[(unsigned char)needle[i]] = currentindex;
	      if (!casesensitive){
		//RBF - this is a little kludgy but it works and this isn't the part 
                //of the code we really need to worry about optimizing...
		//If we aren't case sensitive we just set both the upper and lower case 
		//entries in the jump table.
		table[tolower(needle[i])] = currentindex;
                table[toupper(needle[i])] = currentindex;
	      }
      }	  
}

/*
   Perform a modified boyer-moore string search (w/ support for wildcards and case-insensitive searches)
   and allows the starting position in the buffer to be manually set, which allows data to be skipped
*/

char *bm_needleinhaystack_skipnchars(char *needle, size_t needle_len,
                          char *haystack, size_t haystack_len,
                          size_t table[UCHAR_MAX + 1], int casesensitive,
                          int searchtype,
			  int start_pos)		// the starting position (usually needle_len - 1)
{
      register size_t shift=0;
      register size_t pos = start_pos;
      char *here; 
      
      if(needle_len == 0)
        return haystack;
      
      if(searchtype == SEARCHTYPE_REVERSE){ //Run our search backwards
	while (pos < haystack_len){
	  while( pos < haystack_len && (shift = table[(unsigned char)haystack[haystack_len-pos-1]]) > 0)
	    {
	      pos += shift;
            }
	  if (0 == shift)
	    {
	      if (0 == memwildcardcmp(needle,here = (char *)&haystack[haystack_len-pos-1], needle_len, casesensitive))
		{
		  return(here);
		}
	      else  pos++;
            }
	}
	return NULL;
      }
      else if (searchtype == SEARCHTYPE_FORWARD || searchtype == SEARCHTYPE_FORWARD_NEXT ) { 
	//Run the search forwards ...
	while (pos < haystack_len){
	  while( pos < haystack_len && (shift = table[(unsigned char)haystack[pos]]) > 0)
	    {
	      pos += shift;
            }
	  if (0 == shift)
	    {
	      if (0 == memwildcardcmp(needle,here = (char *)&haystack[pos-needle_len+1], needle_len, casesensitive))
		{
		  return(here);
		}
	      else  pos++;
            }
	}
	return NULL;
      }
      return NULL;	
}

/*
   Perform a modified boyer-moore string search (w/ support for wildcards and case-insensitive searches)
   and allows the starting position in the buffer to be manually set, which allows data to be skipped
*/

char *bm_needleinhaystack(char *needle, size_t needle_len,
                          char *haystack, size_t haystack_len,
                          size_t table[UCHAR_MAX + 1], int casesensitive,
                          int searchtype)
{
      return bm_needleinhaystack_skipnchars(needle, 
				 needle_len, 
				 haystack, 
				 haystack_len, 
				 table, 
				 casesensitive, 
				 searchtype,
				 needle_len - 1);

}



/* extractString() - Given that a header has been found in the buffer
                     recover that file from the image file. Read from
                     the disk only if necessary. 
     
   extractbuf   - Where our result goes 
   fileoffset   - offset for start of found header in the image file
   infile       - the file on disk, in case we need it
   buf          - the buffer of the image file in memory
   needle       - the needle we've matched
*/
struct CharBucket* extractString (struct CharBucket* extractbuf,
				  unsigned long long  fileoffset,
				  FILE* infile,
				  char *foundat,
				  char *buf,
				  unsigned long long lengthofbuf,
				  struct SearchSpecLine needle) {

  unsigned long long offsetInsideBuffer = foundat-buf;
  unsigned long long count, currentfileposition = 0;
  unsigned long long bytesread = 0;
  unsigned long long filelength = 0;
  int startpos = 0;
  /* We need to read needle.length bytes. If we have that much already
     in the buffer, we don't need to go back to the disk. */
  
  if (offsetInsideBuffer + needle.length < SIZE_OF_BUFFER) {
    
    /* We can't use regular strncpy because the buffer may
       contain '\0' characters. If strncpy encounters a
       '\0' it stops immediately. We have to be sure we get the
       whole buffer. I learned this the hard way. (JK)
       
       I wonder if we could squeeze even more speed by adding
       the search for the ending needle here too. That may break
       the abstraction barrier we made by writing needleinhaystack() */
    
    filelength = ((lengthofbuf-offsetInsideBuffer)<needle.length)
      ?lengthofbuf-offsetInsideBuffer:needle.length;
    for (count = 0 ; count <= filelength ; count++) {
      extractbuf->str[count] = foundat[count];
    }
    extractbuf->length = filelength;
    bytesread = filelength;  
  }
  else { 
    
    currentfileposition = ftello(infile);
    fseeko(infile,fileoffset,SEEK_SET);
    bytesread = fread(extractbuf->str,1,needle.length,infile);
    fseeko(infile,currentfileposition,SEEK_SET);
    extractbuf->length = (bytesread<needle.length)?bytesread:needle.length;
  }
  
  /* Now the buffer has the region we want to search for the end marker in.
     If we don't have an endmarker we just want to return,
     else look for the endmarker..... 
     
     The function needleInHaystack returns the *beginning* of the
     needle, not the end. Because we want to return the whole file,
     we have to manually add the length of the needle to the result */
  
  if (needle.endlength > 0){
    

    if(needle.searchtype == SEARCHTYPE_FORWARD_NEXT) // don't allow the string to match twice
	    startpos = needle.endlength - 1 + needle.endlength - 1;
    else
	    startpos = needle.endlength - 1;

    // now call needleinhaystack to find end
    foundat = (char*) bm_needleinhaystack_skipnchars(needle.end,
					    needle.endlength,
					    extractbuf->str,
					    bytesread,
					    needle.end_bm_table,
					    needle.casesensitive,
					    needle.searchtype,
					    startpos);			// pass the start pos so we can override and 
									// not match the same string twice when doing 									        // FORWARD_NEXT searches

    if(foundat){
	if(needle.searchtype == SEARCHTYPE_FORWARD_NEXT) {
	      // don't include the length of the end search tag as the buffer should not include this data
	      extractbuf->length = (foundat - extractbuf->str);
	}
	else {  // add the length of the end search tag so it is included in the results buffer
	      extractbuf->length = (foundat + needle.endlength - extractbuf->str);
	}
    }
  }
  return(extractbuf);
}

int removeWhitespace(char* str){
  int i,j;
  for(i=j=0;str[i];i++) if (!isspace(str[i])) str[j++] = str[i];
  str[j]='\0';
  return i-j;
}


/* These two functions could be combined into one to save some time.
   But because the time saved would be minimal (two loops through the
   SeachSpecLine array), it's not worth the added complication)   */


  int findLongestNeedle(struct SearchSpecLine* SearchSpec){
    int longest = 0;
    int i=0;
    for(i=0; SearchSpec[i].suffix != NULL; i++)
      if(SearchSpec[i].beginlength > longest)
	longest = SearchSpec[i].beginlength;
    return(longest);  
  }     


int findLongestLength(struct SearchSpecLine* SearchSpec){
  int i, longest = 0;
  for(i=0; SearchSpec[i].suffix != NULL; i++)
    if(SearchSpec[i].length > longest)
      longest = SearchSpec[i].length;
  return(longest);    
}     


/* translate() decodes strings with escape sequences in them and
   returns the total length of the string it translated (we 
   can't use strlen because we would like to have null's in 
   our "strings") and has the side effect of truncating the string */

int translate(char *str) {
  char next;
  char *rd=str,*wr=str,*bad;
  char temp[1+3+1];
  char ch;
  
  if(!*rd){  //If it's a null string just return
    return 0;
  }

  while (*rd) {
  /* Is it an escaped character ? */
    if (*rd=='\\') {
      rd++;
      switch(*rd) {
        case '\\': *rd++;*wr++='\\'; break;
        case 'a':  *rd++;*wr++='\a'; break;
        case 's':  *rd++;*wr++=' '; break;
        case 'n':  *rd++;*wr++='\n'; break;
        case 'r':  *rd++;*wr++='\r'; break;
        case 't':  *rd++;*wr++='\t'; break;
        case 'v':  *rd++;*wr++='\v'; break;
  /* Hexadecimal/Octal values are treated in one place using strtoul() */
        case 'x':
        case '0': case '1': case '2': case '3':
         next = *(rd+1); 
	 if (next < 48 || (57 < next && next < 65) || 
	     (70 < next && next < 97) || next > 102) 
	   break;  //break if not a digit or a-f, A-F 
	 next = *(rd+2); 
	 if (next < 48 || (57 < next && next < 65) || 
	     (70 < next && next < 97) || next > 102) 
	   break;  //break if not a digit or a-f, A-F 
		 
	 temp[0]='0'; bad=temp;
         strncpy(temp+1,rd,3);
         temp[4] = '\0';
		 ch=strtoul(temp,&bad,0);
         if (*bad=='\0') {
           *wr++=ch;
           rd+=3;
         } /* else INVALID CHARACTER IN INPUT ('\\' followed by *rd) */
         break;
        default: /* INVALID CHARACTER IN INPUT (*rd)*/
	  *wr++='\\';
	  break;
      }
    }
/* Unescaped characters go directly to the output */
    else *wr++=*rd++;
  }
  *wr = '\0';  //Null terminate the string that we just created... 
  return wr-str;
}

char* skipWhiteSpace(char* str){
  while (isspace(str[0])) 
    str++;
  return str;
}

void handleError(struct foremostState *state, int error) {

  switch(error) {

  case FOREMOST_ERROR_FILE_OPEN:
    /* This is not a fatal error so we can return */
    foremostLog(state,"Foremost was unable to open the image file: %s\n"
		"Skipping...\n\n",
		state->imagefile);
    break;

  case FOREMOST_ERROR_FILE_READ:
    /* This is not a fatal error so we can return */
    foremostLog(state,"Foremost was unable to read the image file: %s\n"
		"Skipping...\n\n",
		state->imagefile);
    break;

  case FOREMOST_ERROR_FILE_WRITE:
    /* We can't write out files, which is very bad. We can't use the 
       formostLog command here because we can't write files ! */
    fprintf(stderr,
	    "Foremost was unable to write output files and will abort.\n");
    closeFile(state->auditFile);
    exit (-1);

  case FOREMOST_ERROR_NO_SEARCH_SPEC:
    foremostLog (state,
		 "ERROR: The configuration file didn't specify anything to search for.\n");
    closeFile(state->auditFile);
    exit (-1);
  
  default:
    foremostLog(state,
		"Foremost has encountered an error it doesn't know"
		"how to handle.\nError code: %d\n", error);
    closeFile(state->auditFile);
    
    /* If we're returning an unknown error code, chances are that 
       something is very wrong. We don't want to process any more 
       files and will exit now. */
    exit (-1);
  }
}


void setttywidth(){
  /* RBF - How do we find the TTY size on Win32? */
  /* Here's the answer... pretty isn't it ;-) */
#ifdef __WIN32
  CONSOLE_SCREEN_BUFFER_INFO csbi;  
  HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
  if(! GetConsoleScreenBufferInfo(hConsole, &csbi)){
    ttywidth = 80;
  }else{
    ttywidth = csbi.dwSize.X;
  }
#else
  struct winsize winsize;
  if(ioctl(fileno(stdout),TIOCGWINSZ, &winsize) != -1)
    ttywidth = winsize.ws_col;
  else
    ttywidth = 80; 
#endif
}


int skipInFile(struct foremostState *state, FILE *infile) {

  int retries = 0;
  while(TRUE) {
    if ((fseeko(infile,state->skip,SEEK_SET))) {
      
      fprintf(stderr,
	      "ERROR: Couldn't skip %lld bytes at the start of image file %s\n",
	      state->skip,state->imagefile);
	
      if (retries++ > 3) {
	fprintf(stderr,"Sorry, maximum retries exceeded...\n");
	return FALSE;
      } else {
	fprintf (stderr,"Waiting to try again... \n");
	sleep(3);
      }
    }
    else {
      fprintf(stderr,"Skipped the first %lld bytes of %s...\n",state->skip,
	      state->imagefile);
      return TRUE;
    }
  }
}