#include "ntfsmftentryreader.hh"

NTFSMFTEntryReader::NTFSMFTEntryReader () {}

NTFSMFTEntryReader::~NTFSMFTEntryReader () {}

void NTFSMFTEntryReader::set_MFTentry (int64_t offset, DataSource& ds) {
  int count = 0;

  for (int i = 0; i < 1024; i++) {
    MFTentry[i] = ds.read (offset + i);
    count++;
    if (count % 16 == 0) {
    }
  }
  
}

int64_t NTFSMFTEntryReader::get_MFTentryAt (int64_t position) {
  return MFTentry[position];
}

uint64_t NTFSMFTEntryReader::readBytes (int64_t offset, int quantity) {
  uint64_t bytes = 0x00;
  for (int i = 0; i < quantity; i++) {
    bytes = bytes | (MFTentry[offset + i] << i*8);
  }
  return bytes;
}

void NTFSMFTEntryReader::fixup_array () {
  uint64_t fixupOffset = readBytes (4, 2);
  uint64_t numFixupArray = readBytes (6, 2);
    
  unsigned short fixupArray[numFixupArray];

  if (numFixupArray >= 10) {
    throw 1;
  }

  for (unsigned int i = 0; i < numFixupArray; i++) {
    fixupArray[i] = readBytes (fixupOffset + (i*2), 2);
  }

  fixupSignature = fixupArray[0];
  int signatureValue = fixupArray[0];
    
  int signatureSlot;
  for (unsigned int i = 0; i < numFixupArray; i++) {
    if ((i*512+510) < 1024) {
      signatureSlot = readBytes(i*512+510, 2);
      if (signatureValue == signatureSlot) {
        MFTentry[i*512+510] = fixupArray[i+1];
        MFTentry[i*512+511] = fixupArray[i+1] >> 8;
      } else {
        throw 2;
      }
    }
  }
}

std::string NTFSMFTEntryReader::read_filename (uint64_t attrOffset, unsigned int length ) {
  unsigned int offsetToContent = readBytes (attrOffset + length + 20, 2);
  unsigned int nameLength = readBytes (attrOffset + length + offsetToContent + 64, 1);
  std::string name;
  for (unsigned int i = 0; i < nameLength; i++) {
    char c = readBytes(attrOffset + length + offsetToContent + 66 + i*2, 1);
    name.push_back(c);
  }
  return name;
}

std::vector<int64_t> NTFSMFTEntryReader::read_runlist (uint64_t attrOffset, unsigned int length, unsigned int attrLength ) {
  std::vector<int64_t> clusters;
  seg segment;

  unsigned int attributeContentLength = readBytes (attrOffset + length + 48, 8);
  unsigned int initializedSize = readBytes (attrOffset + length + 56, 8);

  unsigned int offsetToRunlist = readBytes (attrOffset + length + 32, 2);

    
  unsigned int attrLeft = attrLength - offsetToRunlist;

    
  unsigned int offsetInRunlist = 0x00;
  unsigned int offsetSinceLast = 0x00;
    
  unsigned char runlistByte = readBytes ( attrOffset + length + offsetToRunlist + offsetInRunlist, 1 );
    
  unsigned int runLength, offsetField, numClusters;
  int relativeOffset;
  unsigned int off = 0x00;

  unsigned int fileLength = 0;

  while ( (attrLeft > 0) && (runlistByte != 0)) {
        
    runLength = 0xF & runlistByte;
    offsetField = 0xF & (runlistByte >> 4);

    numClusters = readBytes ( attrOffset + length + offsetToRunlist + 1 + offsetInRunlist, runLength);
        
    relativeOffset = readBytes ( attrOffset + length + offsetToRunlist + 1 + runLength + offsetInRunlist, offsetField);
    relativeOffset = (relativeOffset << (8 * (8 - offsetField))) >> (8 * (8 - offsetField));
    off = off + relativeOffset;                                       // add to overall offset

    if ( offsetField != 0 ) {                                           // if not sparse
      // fill in segment info
      segment.file_offset = fileLength;
      segment.fs_offset = off * BLOCK_SIZE;
      segment.img_offset = segment.fs_offset + 63*BLOCK_SIZE;
      segment.compressed = 0;

      // filesize < initalized size
      int len = numClusters * BLOCK_SIZE;
      if (((fileLength + len) > initializedSize) && (initializedSize < attributeContentLength)) {

	int diff = attributeContentLength - initializedSize;
	segment.len = initializedSize - segment.file_offset;

	seg segment2;
	segment2.file_offset = fileLength + segment.len;
	segment2.fs_offset = 0;
	segment2.img_offset = 0;
	segment2.len = diff;
	segment2.compressed = 0;
	
	segs.push_back (segment);
	segs.push_back (segment2);
	
      } else {
       	segment.len = len;
	segs.push_back ( segment );
      }

      fileLength += len;

      for (unsigned int i = 0; i < numClusters; i++) {
        clusters.push_back ( off + i );
      }
    } else {
      // compressed byte runs are printed out in runs of 512 in fiwalk
      unsigned int i = 0;
      while ((i < numClusters) && (fileLength < attributeContentLength)) {
	segment.file_offset = fileLength;
 	segment.fs_offset = 0;
	segment.img_offset = 0;
	segment.len = BLOCK_SIZE;
	segment.compressed = 1;
	segs.push_back ( segment );

	fileLength += BLOCK_SIZE;
	i++;
      }
    }
        
    offsetSinceLast = 1 + runLength + offsetField;
    offsetInRunlist = offsetInRunlist + offsetSinceLast;

        
    runlistByte = readBytes (attrOffset + length + offsetToRunlist + offsetInRunlist, 1);
    if (runlistByte == 0x00) {
      attrLeft = 0;
    } else {
      // when would attributeContentLength be 0?
      // once the fileLength read in has become bigger than the actual size of the attribute content
      if ( (0 != attributeContentLength) && (attributeContentLength <= fileLength)) { // why do i have this check in here?
	  attrLeft = 0;
      } else {
	attrLeft = attrLeft - offsetSinceLast;
      }
    }
  }
  return clusters;
}

NTFSMFTEntry NTFSMFTEntryReader::read_entry ( int64_t offset, DataSource& ds ) {
  NTFSMFTEntry entry;
  set_MFTentry ( offset, ds );
  fixup_array ();
  entry.set_fixupSignature (fixupSignature);
  if (readBytes (0, 4) == 0x454c4946) {                                           // if starts with $FILE

    int flags = readBytes (22, 2);
    if ((flags & 0x02) == 0x02) {
      entry.set_name_type ('d');      
    } else if ((flags & 0x02) == 0) {
      entry.set_name_type ('r');
    }

    if (flags & 0x01) {
      entry.set_inUse ( 1 );
    } else {
      entry.set_inUse ( 0 );
    }

    if (flags == 1 || flags == 3 || flags == 0 || flags == 2) { 

      uint64_t attrOffset = readBytes (20, 2);                                // offset to first attribute
      //      std::cout << std::hex << "attrOffset: 0x" << attrOffset << std::endl;
            
      unsigned int attribute = 0xFF;
      unsigned int length = 0x00;
      unsigned int attrLength = 0x00;
      unsigned char nonResidentFlag = 0xFF;

      std::string name = "";
      int64_t parentMFTentry = 0;
      uint64_t ctime = 0;
      uint64_t atime = 0;
      uint64_t crtime = 0;
      uint64_t mtime = 0;

      while (attribute != 0x80 && attribute != 0xffffffff) {
        attribute = readBytes( attrOffset + length, 4 );                    // attribute type
        attrLength = readBytes( attrOffset + length + 4, 4);                // attribute length
                
        // std::cout << "attribute: " << attribute << std::endl;
	//        std::cout << "attrLength: 0x" << attrLength << std::endl;

	if (attrLength > (BLOCK_SIZE*2 - length)) {
	  throw 11;
	}

	if (attribute == 0x10) {
	  int offsetContent = readBytes (attrOffset + length + 20, 2);
	  entry.set_ownerID ( readBytes (attrOffset + length + offsetContent + 48, 4) );

	} else if (attribute == 0x30) {
	  entry.set_offsetFileNameAttr ( attrOffset + length );
          name = read_filename ( attrOffset, length );                    // get filename
          // std::cout << "name: " << name << std::endl;

	  int offsetContent =  readBytes(attrOffset + length + 20, 2);
	  parentMFTentry = readBytes (attrOffset + length + offsetContent, 6);
	  ctime = readBytes (attrOffset + length + offsetContent + 24, 8);
	  atime = readBytes (attrOffset + length + offsetContent + 32, 8);
	  crtime = readBytes (attrOffset + length + offsetContent + 8, 8);
	  mtime = readBytes (attrOffset + length + offsetContent + 16, 8);

	} else if (attribute == 0x80) {
          nonResidentFlag = readBytes (attrOffset + length + 8, 1);
          entry.set_resVnonres ( nonResidentFlag );
	  entry.set_offsetDATAattr ( attrOffset + length );
	  entry.set_compressed ( readBytes (attrOffset + length + 12, 2) & 0x01); // flag 0x0001 set if compressed
	  //	  std::cout << "Location of $DATA: " << attrOffset + length << std::endl;
        }
                
        if (attribute != 0x80) {
          length = length + attrLength;
        }
      }
            
      entry.set_filename ( name );
      entry.set_parentMFTentry ( parentMFTentry );
      entry.set_nlink ( readBytes (18, 2) );
      entry.set_sequenceValue ( readBytes (16, 2) );
      entry.set_ctime ( ctime );
      entry.set_atime ( atime );
      entry.set_crtime ( crtime );
      entry.set_mtime ( mtime );

      if (nonResidentFlag == 0x00) {                                          // if resident
	//         std::cout << "$DATA is resident" << std::endl;
                
        entry.set_length( readBytes ( attrOffset + length + 16, 4 ));
                
        entry.set_entryLocation ( offset );
                
        int64_t offInEntry = readBytes ( attrOffset + length + 20, 2 ) + attrOffset + length;
        entry.set_offsetInEntry( offInEntry );
                
      }
      else if (nonResidentFlag == 0x01) {                                     // if non-resident
	//         std::cout << "$DATA is non-resident" << std::endl;
                
        uint64_t sizeOfContent = readBytes (attrOffset + length + 48, 8);
        entry.set_length ( sizeOfContent );
	uint64_t initializedSize = readBytes (attrOffset + length + 56, 8);
	entry.set_initializedSize ( initializedSize );
	entry.set_runlistOffset ( readBytes (attrOffset + length + 32, 2) + (attrOffset + length));
        entry.set_clusters( read_runlist (attrOffset, length, attrLength) );
	// has to be done after read_runlist has been called! 
	// otherwise segs isn't filled
	entry.set_seglist ( segs );
      }         
    }
  }
  return entry;
}
