// 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);
    return sampleInfo;
}
int main() {
    auto sampleInfo = getPeInfo("z:\\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, "z:\\features.txt");
    }
    peBuffer.release();
    system("pause");
    return 0;
}
