111

Author: Orange20000922

Filerestore_CLI/src/analysis/ImageTable.cpp

@@ -57,6 +57,7 @@ ULONGLONG ImageTableAnalyzer::GetFuncaddressByName(string name,string file)
 					PIMAGE_IMPORT_BY_NAME pImportByName = (PIMAGE_IMPORT_BY_NAME)((DWORD_PTR)lpBuffer + RVAtoFOA((DWORD)IAT->u1.AddressOfData,lpBuffer));
 					if (name.compare(string((char*)pImportByName->Name))==0) {
 						funcAddress = IAT->u1.Function+pOptionalHeader64->ImageBase;
+						CleanupMappedFile();
 						return funcAddress;
 					}
 				}
@@ -65,6 +66,7 @@ ULONGLONG ImageTableAnalyzer::GetFuncaddressByName(string name,string file)
 			pImportDescriptor++;
 		}
 	}
+	CleanupMappedFile();
 	return 0;
 }
 //遍历导入表中的DLL名称
@@ -128,6 +130,7 @@ vector<string> ImageTableAnalyzer::AnalyzeTableForDLL(string file)
 		dllList.push_back(string(dllName));
 		pImportDescriptor++;
 	}
+    CleanupMappedFile();
     return dllList;
 }
 //遍历导入表中的函数名称
@@ -226,6 +229,7 @@ map<string, vector<string>> ImageTableAnalyzer::AnalyzeTableForFunctions(string
 		funcList[string(dllName)] = funcNames;
 		pImportDescriptor++;
 	}
+	CleanupMappedFile();
 	return funcList;
 }
 bool ImageTableAnalyzer::IsImagineTable(LPVOID lpBuffer)
@@ -332,7 +336,19 @@ BOOL ImageTableAnalyzer::CheckIsAdmin()
 ImageTableAnalyzer::ImageTableAnalyzer() {
 }
 ImageTableAnalyzer::~ImageTableAnalyzer() {
-	CloseHandle(ImageTableAnalyzer::hFileMapping);
-	CloseHandle(ImageTableAnalyzer::hFile);
-	UnmapViewOfFile(lpBuffer);
+	CleanupMappedFile();
+}
+void ImageTableAnalyzer::CleanupMappedFile() {
+	if (lpBuffer) {
+		UnmapViewOfFile(lpBuffer);
+		lpBuffer = NULL;
+	}
+	if (hFileMapping) {
+		CloseHandle(hFileMapping);
+		hFileMapping = NULL;
+	}
+	if (hFile && hFile != INVALID_HANDLE_VALUE) {
+		CloseHandle(hFile);
+		hFile = NULL;
+	}
 }

Filerestore_CLI/src/analysis/ImageTable.h

@@ -33,6 +33,9 @@ class ImageTableAnalyzer
 	// - BOOL ElevatePrivilegeForSystem(wstring Privilege);
 	// - BOOL DLLImplantForSystemProcess(string dllfile, int PID);
 
+private:
+	void CleanupMappedFile();
+
 public:
 	ImageTableAnalyzer();
 	~ImageTableAnalyzer();

Filerestore_CLI/src/fileRestore/MFTReader.cpp

@@ -296,7 +296,47 @@ bool MFTReader::ReadMFTBatch(ULONGLONG startRecordNumber, ULONGLONG recordCount,
         return false;
     }
 
-    // 计算需要读取的簇范围
+    // 碎片化 MFT 路径：逐簇读取，缓存避免重复 I/O
+    if (mftDataRunsLoaded && mftDataRuns.size() > 1) {
+        LOG_DEBUG("Using fragmented MFT batch read path");
+
+        ULONGLONG lastLCN = (ULONGLONG)-1;
+        vector<BYTE> cachedCluster;
+
+        for (ULONGLONG i = 0; i < recordCount; i++) {
+            ULONGLONG recordNumber = startRecordNumber + i;
+            ULONGLONG lcn, offsetInCluster;
+
+            if (!GetMFTRecordLCN(recordNumber, lcn, offsetInCluster)) {
+                LOG_WARNING_FMT("Record #%llu: cannot resolve LCN, skipping", recordNumber);
+                continue;
+            }
+
+            // 同一个簇内的多条记录只读一次磁盘
+            if (lcn != lastLCN) {
+                if (!ReadClusters(lcn, 1, cachedCluster)) {
+                    LOG_WARNING_FMT("Failed to read cluster at LCN %llu for record #%llu", lcn, recordNumber);
+                    continue;
+                }
+                lastLCN = lcn;
+            }
+
+            if (offsetInCluster + bytesPerFileRecord > cachedCluster.size()) {
+                LOG_WARNING_FMT("Record #%llu offset out of bounds, skipping", recordNumber);
+                continue;
+            }
+
+            vector<BYTE> record(bytesPerFileRecord);
+            memcpy(record.data(), cachedCluster.data() + offsetInCluster, bytesPerFileRecord);
+            ApplyUSAFixup(record, bytesPerSector);
+            records.push_back(move(record));
+        }
+
+        LOG_DEBUG_FMT("ReadMFTBatch (fragmented) completed: extracted %zu valid records", records.size());
+        return true;
+    }
+
+    // 连续 MFT 路径：批量读取（原始代码）
     ULONGLONG startCluster = startRecordNumber / recordsPerCluster;
     ULONGLONG endRecord = startRecordNumber + recordCount - 1;
     ULONGLONG endCluster = endRecord / recordsPerCluster;

Filerestore_CLI/src/fileRestore/OverwriteDetector.cpp

@@ -16,7 +16,8 @@ using namespace std::chrono;
 OverwriteDetector::OverwriteDetector(MFTReader* mftReader)
     : reader(mftReader), detectionMode(MODE_BALANCED),
       threadingStrategy(THREADING_AUTO),
-      cachedStorageType(STORAGE_UNKNOWN), storageTypeDetected(false) {
+      cachedStorageType(STORAGE_UNKNOWN), storageTypeDetected(false),
+      bitmapLoaded(false), cachedBitmapLCN((ULONGLONG)-1) {
     bytesPerCluster = reader->GetBytesPerSector() * reader->GetSectorsPerCluster();
     threadPool = std::make_unique<OverwriteDetectionThreadPool>(this);
     LOG_DEBUG("OverwriteDetector 已初始化，支持多线程");
@@ -164,12 +165,79 @@ bool OverwriteDetector::HasValidFileStructure(const vector<BYTE>& data) {
     return false;
 }
 
+// 加载 $Bitmap (MFT记录#6) 的 data runs
+bool OverwriteDetector::LoadBitmapDataRuns() {
+    if (bitmapLoaded) return !bitmapRuns.empty();
+    bitmapLoaded = true;
+
+    vector<BYTE> bitmapRecord;
+    if (!reader->ReadMFT(6, bitmapRecord)) {
+        LOG_ERROR("Failed to read $Bitmap MFT record (#6)");
+        return false;
+    }
+
+    if (!ExtractDataRuns(bitmapRecord, bitmapRuns)) {
+        LOG_ERROR("Failed to extract $Bitmap data runs");
+        return false;
+    }
+
+    LOG_INFO_FMT("$Bitmap data runs loaded: %zu runs", bitmapRuns.size());
+    return true;
+}
+
 // 检查簇是否已分配给其他文件（通过读取$Bitmap）
 bool OverwriteDetector::CheckClusterAllocation(ULONGLONG clusterNumber) {
-    // TODO: 实现$Bitmap读取功能
-    // 这需要读取MFT记录#6（$Bitmap），检查对应位是否为1
-    // 暂时返回false（假设未分配）
-    return false;
+    if (!LoadBitmapDataRuns()) {
+        return false;  // 无法读取 $Bitmap，假设未分配
+    }
+
+    // 计算该簇号在 bitmap 中的字节偏移和位位置
+    ULONGLONG byteOffset = clusterNumber / 8;
+    BYTE bitMask = 1 << (clusterNumber % 8);
+
+    // 计算这个字节落在 $Bitmap 文件的哪个簇中
+    ULONGLONG bitmapBytesPerCluster = (ULONGLONG)reader->GetBytesPerSector() * reader->GetSectorsPerCluster();
+    ULONGLONG bitmapClusterIndex = byteOffset / bitmapBytesPerCluster;
+    ULONGLONG offsetInCluster = byteOffset % bitmapBytesPerCluster;
+
+    // 通过 $Bitmap 的 data runs 找到实际的物理 LCN
+    ULONGLONG currentClusterOffset = 0;
+    ULONGLONG targetLCN = 0;
+    bool found = false;
+
+    for (const auto& run : bitmapRuns) {
+        if (bitmapClusterIndex < currentClusterOffset + run.second) {
+            targetLCN = run.first + (bitmapClusterIndex - currentClusterOffset);
+            found = true;
+            break;
+        }
+        currentClusterOffset += run.second;
+    }
+
+    if (!found) {
+        return false;  // 簇号超出 bitmap 范围
+    }
+
+    // 检查缓存：如果上次读的就是这个 bitmap 簇，直接查表
+    if (cachedBitmapLCN == targetLCN && !cachedBitmapData.empty()) {
+        if (offsetInCluster >= cachedBitmapData.size()) return false;
+        return (cachedBitmapData[offsetInCluster] & bitMask) != 0;
+    }
+
+    // 读取 bitmap 簇并缓存
+    vector<BYTE> clusterData;
+    if (!reader->ReadClusters(targetLCN, 1, clusterData)) {
+        return false;
+    }
+
+    if (offsetInCluster >= clusterData.size()) {
+        return false;
+    }
+
+    cachedBitmapLCN = targetLCN;
+    cachedBitmapData = move(clusterData);
+
+    return (cachedBitmapData[offsetInCluster] & bitMask) != 0;
 }
 
 // 读取可变长度的无符号整数

Filerestore_CLI/src/fileRestore/OverwriteDetector.h

@@ -86,6 +86,13 @@ class OverwriteDetector
     bool HasValidFileStructureFromPointer(const BYTE* data, size_t size);
     bool CheckClusterAllocation(ULONGLONG clusterNumber);
 
+    // $Bitmap 簇分配检测
+    vector<pair<ULONGLONG, ULONGLONG>> bitmapRuns;  // $Bitmap 文件的 data runs
+    bool bitmapLoaded;
+    ULONGLONG cachedBitmapLCN;          // 缓存的 bitmap 簇 LCN
+    vector<BYTE> cachedBitmapData;      // 缓存的 bitmap 簇数据
+    bool LoadBitmapDataRuns();
+
     // 解析Data Runs的辅助函数
     ULONGLONG ReadVariableLength(const BYTE* data, int length);
     LONGLONG ReadVariableLengthSigned(const BYTE* data, int length);