#include "ntfsbitmap.hh"


using namespace Gecode;



NTFSBitmapExtractor::NTFSBitmapExtractor (DataSource& ds) :
  bitmap (2 * 1024 * 1024)
{
  std::cout << "Start NTFS $Bitmap scan" << std::endl;

  int64_t filesystemOffset = 0; // It is assumed that the DataSource aligns the NTFS image at offset 0 now.

  int64_t bps2 = int64_t (ds.read (filesystemOffset + 12)) << 8;
  int64_t bps1 = int64_t (ds.read (filesystemOffset + 11));
  int64_t bytesPerSector = bps1 | bps2;
	
  int64_t startingClusterOfMFT = 0;
  for (int s = 0; s < 8; s++) {
    startingClusterOfMFT = startingClusterOfMFT | (int64_t (ds.read (filesystemOffset + 48 + s)) << 8 * s);
  }

  int64_t sectorsPerCluster = 0xff && int64_t (ds.read (filesystemOffset + 13));
  int64_t start = filesystemOffset + (startingClusterOfMFT * bytesPerSector * sectorsPerCluster) + 6 * 1024;
  int length = 1024; //1 KB (size of entry)
  int attributeType = 0; //want 128 for $DATA

  int fa1 = ds.read (start + 21) << 8;
  int fa2 = ds.read (start + 20);
  int firstAttribute = fa1 | fa2;

  int64_t currentAttributeStart = start + firstAttribute; //start of $Bitmap + go to start of attributes
  int64_t attributeLength;
  int64_t dataStart;
  int64_t dataEnd;

  int64_t allocated = 0;
  int64_t unallocated = 0;

  while (currentAttributeStart < start + length && attributeType != -1) {
    //Is $DATA non-resident?
    int nonResidentFlag = 0xff & ds.read (currentAttributeStart + 8);

    //Attribute type
    attributeType = 0;
    for(int a = 0; a < 4; a++) {
      attributeType = attributeType |(ds.read (currentAttributeStart + a) << (8 * a));
    }

    // std::cout << attributeType << std::endl;

    //Attribute length
    int64_t al1 = ds.read (currentAttributeStart + 7) << 24;
    int64_t al2 = ds.read (currentAttributeStart + 6) << 16;
    int64_t al3 = ds.read (currentAttributeStart + 5) << 8;
    int64_t al4 = ds.read (currentAttributeStart + 4);
    attributeLength = al1 | al2 | al3 | al4;

    if (attributeType == 128) {
      if ( (nonResidentFlag & 0x1) == 0) { //resident
        error ("$DATA attribute of NTFS $Bitmap is resident");
      }
      else { //non-resident
        int64_t startingVCNOfRunlist = 0;
        int64_t endingVCNOfRunlist = 0;
        int offsetToRunlist = 0;
        // int compressionUnitSize = 0;
        int64_t allocatedSizeOfAttributeContent = 0;
        int64_t actualSizeOfAttributeContent = 0;
        int64_t initializedSizeOfAttributeContent = 0;
				
        for (int v = 0; v < 8; v++) {
          startingVCNOfRunlist = startingVCNOfRunlist | (ds.read (currentAttributeStart + 16 + v) << (8 * v));
        }
				
        for (int v2 = 0; v2 < 8; v2++) {
          endingVCNOfRunlist = endingVCNOfRunlist |(ds.read (currentAttributeStart + 24 + v2) << (8 * v2));
        }
				
        offsetToRunlist = ds.read (currentAttributeStart + 32) | (ds.read (currentAttributeStart + 33) << 8);
				
        for (int v3 = 0; v3 < 8; v3++)
          {
            allocatedSizeOfAttributeContent=allocatedSizeOfAttributeContent|(ds.read(currentAttributeStart + 40+v3)<<(8*v3));
          }

        for(int v4=0;v4<8;v4++)
          {
            actualSizeOfAttributeContent=actualSizeOfAttributeContent|(ds.read(currentAttributeStart + 48+v4)<<(8*v4));
          }
				
        for(int v5=0;v5<8;v5++)
          {
            initializedSizeOfAttributeContent=initializedSizeOfAttributeContent|(ds.read(currentAttributeStart + 56+v5)<<(8*v5));
          }
				
        int currentRunlistOffset = currentAttributeStart + offsetToRunlist;
        int64_t previousRunLocation = startingVCNOfRunlist;
        int64_t cumulativeRunLength = 0;

        while (cumulativeRunLength <= endingVCNOfRunlist) {
          int sizeOfRunLength = 0xf & ds.read (currentRunlistOffset);

          int sizeOfRunOffset = 0xf & (ds.read(currentRunlistOffset) >> 4);

          int64_t runLength = 0;
          for(int s = 0; s < sizeOfRunLength; s++) {
            runLength = runLength | (ds.read (currentRunlistOffset + 1 + s) << (8 * s));
          }
					
          int64_t runOffset = 0;
          for (int r = 0; r < sizeOfRunOffset; r++) {
            runOffset = runOffset | (ds.read (currentRunlistOffset + 1 + sizeOfRunLength + r) << 8 * r);
          }
          //sign extension
          runOffset = (runOffset << (8 * (8 - sizeOfRunOffset))) >> (8 * (8 - sizeOfRunOffset));
          
          cumulativeRunLength += runLength;
          previousRunLocation += runOffset;
          currentRunlistOffset += 1 + sizeOfRunLength + sizeOfRunOffset;

          //run.add (runLength * bytesPerSector * sectorsPerCluster, previousRunLocation * bytesPerSector * sectorsPerCluster);
          dataStart =previousRunLocation * bytesPerSector * sectorsPerCluster;
          dataEnd = dataStart + runLength * bytesPerSector * sectorsPerCluster;

          char ch;
          unsigned int b;
          unsigned int bitmask;

          std::cout << "About to write out; dataStart = " << dataStart << "; dataEnd = " << dataEnd << std::endl;
          for (int64_t x = dataStart; x < dataEnd; x++) { //Getting all the bits out of the run
            ch = char (ds.read (filesystemOffset + x));
            bitmap[x-dataStart] = ch;

            for (int i = 0; i < 8; i++) { //going through bits in the chosen byte
              bitmask = 1 << i;
              b = (ch & bitmask) ? 1 : 0;

              if (b == 0) {
                unallocated++;
              } else {
                allocated++;
              }
            }
          }
        }
      }			
      break;
    }
    currentAttributeStart += attributeLength;
  }

  std::cout << "End NTFS $Bitmap scan" << std::endl;
  std::cout << "Allocated: " << allocated << std::endl;
  std::cout << "Unallocated: "<< unallocated << std::endl;
}


bool NTFSBitmapExtractor::isClusterAllocated (int cluster)
{
  unsigned char ch = bitmap[cluster / 8];
  unsigned int bitmask = 1 << (cluster % 8);
  return (ch & bitmask);
}
