// ai_anti_malware.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include "head.h"

auto getPeInfo(std::string inputFilePath) -> std::shared_ptr<BasicPeInfo> {
    auto sampleInfo = std::make_shared<BasicPeInfo>();
    sampleInfo->inputFilePath = inputFilePath;

    sampleInfo->peBuffer =
        peconv::load_pe_module((const char*)sampleInfo->inputFilePath.c_str(),
                               sampleInfo->peSize, false, false);
    sampleInfo->ntHead64 = peconv::get_nt_hdrs64((BYTE*)sampleInfo->peBuffer);
    sampleInfo->ntHead32 = peconv::get_nt_hdrs32((BYTE*)sampleInfo->peBuffer);
    sampleInfo->isX64 = peconv::is64bit((BYTE*)sampleInfo->peBuffer);
    sampleInfo->RecImageBase =
        sampleInfo->isX64
            ? (DWORD64)sampleInfo->ntHead64->OptionalHeader.ImageBase
            : (DWORD)sampleInfo->ntHead32->OptionalHeader.ImageBase;
    sampleInfo->isRelocated =
        peconv::relocate_module((BYTE*)sampleInfo->peBuffer, sampleInfo->peSize,
                                sampleInfo->RecImageBase);

    sampleInfo->entryPoint =
        sampleInfo->isX64
            ? sampleInfo->ntHead64->OptionalHeader.AddressOfEntryPoint
            : sampleInfo->ntHead32->OptionalHeader.AddressOfEntryPoint;
    sampleInfo->imageEnd =
        sampleInfo->RecImageBase +
        (sampleInfo->isX64 ? sampleInfo->ntHead64->OptionalHeader.SizeOfImage
                           : sampleInfo->ntHead32->OptionalHeader.SizeOfImage);
    printf("Debug - Memory mapping parameters:\n");
    printf("RecImageBase: 0x%llx\n", sampleInfo->RecImageBase);
    printf("peSize: 0x%llx\n", sampleInfo->peSize);
    printf("Page aligned base: 0x%llx\n", sampleInfo->RecImageBase & ~0xFFF);
    printf("Page aligned size: 0x%llx\n",
           (sampleInfo->peSize + 0xFFF) & ~0xFFF);
    sampleInfo->RecImageBase = sampleInfo->RecImageBase & ~0xFFF;
    sampleInfo->peSize = (sampleInfo->peSize + 0xFFF) & ~0xFFF;
    return sampleInfo;
}
int doMl(int argc, char* argv[]) {
    // 检查命令行参数
    if (argc < 3) {
        std::cout << "用法: " << argv[0] << " <样本目录路径> <输出CSV路径>"
                  << std::endl;
        std::cout << "或者: " << argv[0]
                  << " -single <单个文件路径> <输出CSV路径>" << std::endl;
        return 1;
    }
    MachineLearning ml;

    if (std::string(argv[1]) == "-single") {
        // 处理单个文件
        if (argc < 4) {
            std::cout << "处理单个文件时需要提供文件路径和输出CSV路径"
                      << std::endl;
            return 1;
        }

        std::string filePath = argv[2];
        std::string csvPath = argv[3];

        // 读取文件
        std::vector<uint8_t> buffer = ml.ReadFileToBuffer(filePath);
        if (buffer.empty()) {
            std::cerr << "无法读取文件: " << filePath << std::endl;
            return 1;
        }

        // 提取特征
        std::vector<double> features =
            ml.ExtractFeatures(buffer.data(), buffer.size());
        if (features.empty()) {
            std::cerr << "无法从文件提取特征: " << filePath << std::endl;
            return 1;
        }

        // 导出到CSV
        if (!ml.ExportToCSV(features, csvPath)) {
            std::cerr << "无法导出到CSV文件: " << csvPath << std::endl;
            return 1;
        }

        std::cout << "成功处理文件并导出特征到: " << csvPath << std::endl;
    } else {
        // 处理目录
        std::string dirPath = argv[1];
        std::string csvPath = argv[2];

        std::cout << "开始处理目录: " << dirPath << std::endl;
        std::cout << "特征将导出到: " << csvPath << std::endl;

        if (!ml.ProcessDirectory(dirPath, csvPath)) {
            std::cerr << "处理目录时发生错误" << std::endl;
            return 1;
        }
    }
    return 0;
};
int main(int argc, char* argv[]) {
    doMl(argc, argv);
    /*
 auto sampleInfo = getPeInfo(
     "E:\\对战平台\\CrowAntiCheat\\CrowAntiCheat\\client\\Console_"
     "Test\\Release\\Console_Test.exe");
 // auto sampleInfo = getPeInfo("C:\\ConsoleApplication1.exe");
 printf("input new file %s \n", sampleInfo->inputFilePath);
 printf("is x64: %d\n", sampleInfo->isX64);
 printf("is relocated: %d\n", sampleInfo->isRelocated);
 printf("RecImageBase: %llx\n", sampleInfo->RecImageBase);
 auto sandbox = std::make_shared<Sandbox>();
 sandbox->InitEnv(sampleInfo);
 sandbox->Run();
 auto [peBuffer, peSize] = sandbox->DumpPE();

 if (peBuffer) {
     printf("peBuffer: %p\n", peBuffer.get());
     printf("peSize: %d\n", peSize);
     // peconv::dump_to_file("z:\\dumped_main.exe", peBuffer.get(), peSize);
     MachineLearning ml;
     ml.ExtractFeatures(peBuffer.get(), peSize);
 }
 peBuffer.release();
 */
    system("pause");
    return 0;
}