#include "sandbox.h"
#include "sandbox_callbacks.h"
#include <fstream>

// 文件句柄信息结构体
struct FileHandleInfo {
    uint64_t handle;               // 文件句柄
    std::string path;              // 文件路径
    std::vector<char> content;     // 文件内容
    size_t currentPosition;        // 当前读取位置
    bool isValid;                  // 是否有效
};

// 在Sandbox类中添加文件句柄表
std::map<uint64_t, FileHandleInfo*> Sandbox::m_fileHandles;

// 生成唯一的文件句柄
uint64_t Sandbox::GenerateFileHandle() {
    static uint64_t nextHandle = 0x2000; // 起始值，避免与特殊句柄冲突
    return nextHandle++;
}

// 获取文件句柄信息
FileHandleInfo* Sandbox::GetFileHandleInfo(uint64_t handle) {
    auto it = m_fileHandles.find(handle);
    if (it != m_fileHandles.end()) {
        return it->second;
    }
    return nullptr;
}

// CreateFileW API实现
auto Api_CreateFileW(void* sandbox, uc_engine* uc, uint64_t address) -> void {
    auto context = static_cast<Sandbox*>(sandbox);
    uint64_t lpFileName = 0;
    uint64_t dwDesiredAccess = 0;
    uint64_t dwShareMode = 0;
    uint64_t lpSecurityAttributes = 0;
    uint64_t dwCreationDisposition = 0;
    uint64_t dwFlagsAndAttributes = 0;
    uint64_t hTemplateFile = 0;

    // 获取参数
    if (context->GetPeInfo()->isX64) {
        // x64: rcx = lpFileName, rdx = dwDesiredAccess, r8 = dwShareMode, r9 = lpSecurityAttributes
        uc_reg_read(uc, UC_X86_REG_RCX, &lpFileName);
        uc_reg_read(uc, UC_X86_REG_RDX, &dwDesiredAccess);
        uint64_t temp_share_mode;
        uc_reg_read(uc, UC_X86_REG_R8, &temp_share_mode);
        dwShareMode = static_cast<uint32_t>(temp_share_mode);
        uc_reg_read(uc, UC_X86_REG_R9, &lpSecurityAttributes);

        // 从栈上读取其他参数
        uint64_t rsp;
        uc_reg_read(uc, UC_X86_REG_RSP, &rsp);
        uc_mem_read(uc, rsp + 0x28, &dwCreationDisposition, sizeof(uint32_t));
        uc_mem_read(uc, rsp + 0x30, &dwFlagsAndAttributes, sizeof(uint32_t));
        uc_mem_read(uc, rsp + 0x38, &hTemplateFile, sizeof(uint64_t));
    } else {
        // x86: 从栈上读取参数
        uint32_t esp;
        uc_reg_read(uc, UC_X86_REG_ESP, &esp);
        esp += 0x4;  // 跳过返回地址

        uint32_t temp_values[7];
        for (int i = 0; i < 7; i++) {
            uc_mem_read(uc, esp + (i * 4), &temp_values[i], sizeof(uint32_t));
        }

        lpFileName = temp_values[0];
        dwDesiredAccess = temp_values[1];
        dwShareMode = temp_values[2];
        lpSecurityAttributes = temp_values[3];
        dwCreationDisposition = temp_values[4];
        dwFlagsAndAttributes = temp_values[5];
        hTemplateFile = temp_values[6];
    }

    // 读取文件名
    std::wstring filename;
    if (lpFileName != 0) {
        wchar_t buffer[MAX_PATH] = {0};
        size_t i = 0;
        do {
            wchar_t wch;
            uc_mem_read(uc, lpFileName + (i * 2), &wch, sizeof(wchar_t));
            buffer[i] = wch;
            i++;
        } while (buffer[i - 1] != 0 && i < MAX_PATH);
        
        if (i > 0 && i < MAX_PATH) {
            filename = buffer;
        }
    }

    // 将宽字符串转换为窄字符串，用于日志和内部存储
    std::string ansiFilename(filename.begin(), filename.end());

    // 创建文件句柄信息
    auto fileHandleInfo = new FileHandleInfo();
    fileHandleInfo->path = ansiFilename;
    fileHandleInfo->currentPosition = 0;
    fileHandleInfo->isValid = true;

    // 读取实际文件内容(如果存在)
    std::ifstream file(ansiFilename, std::ios::binary | std::ios::ate);
    if (file.is_open()) {
        std::streamsize size = file.tellg();
        file.seekg(0, std::ios::beg);

        fileHandleInfo->content.resize(size);
        if (size > 0) {
            file.read(fileHandleInfo->content.data(), size);
        }
        file.close();
    } else {
        // 文件不存在或无法打开
        if (dwCreationDisposition == CREATE_NEW || 
            dwCreationDisposition == CREATE_ALWAYS || 
            dwCreationDisposition == OPEN_ALWAYS) {
            // 创建空文件
            fileHandleInfo->content.clear();
        } else {
            // 返回无效句柄
            delete fileHandleInfo;
            uint64_t invalidHandle = (uint64_t)INVALID_HANDLE_VALUE;
            uc_reg_write(uc, 
                context->GetPeInfo()->isX64 ? UC_X86_REG_RAX : UC_X86_REG_EAX, 
                &invalidHandle);

            // 设置错误码
            DWORD error = ERROR_FILE_NOT_FOUND;
            if (context->GetPeInfo()->isX64) {
                context->GetTeb64()->LastErrorValue = error;
            } else {
                context->GetTeb32()->LastErrorValue = error;
            }
            return;
        }
    }

    // 生成唯一的文件句柄
    uint64_t fileHandle = context->GenerateFileHandle();
    fileHandleInfo->handle = fileHandle;

    // 添加到文件句柄表
    context->m_fileHandles[fileHandle] = fileHandleInfo;

    printf("[*] CreateFileW: File=%ls, Access=0x%x, Share=0x%x, Create=0x%x, Handle=0x%llx\n",
           filename.c_str(), dwDesiredAccess, dwShareMode, dwCreationDisposition, fileHandle);

    // 返回文件句柄
    uc_reg_write(uc, 
        context->GetPeInfo()->isX64 ? UC_X86_REG_RAX : UC_X86_REG_EAX, 
        &fileHandle);
}

// ReadFileA API实现
auto Api_ReadFileA(void* sandbox, uc_engine* uc, uint64_t address) -> void {
    Api_ReadFile(sandbox, uc, address, false);
}

// ReadFileW API实现 
auto Api_ReadFileW(void* sandbox, uc_engine* uc, uint64_t address) -> void {
    Api_ReadFile(sandbox, uc, address, true);
}

// 统一的ReadFile实现
auto Api_ReadFile(void* sandbox, uc_engine* uc, uint64_t address, bool isWideChar) -> void {
    auto context = static_cast<Sandbox*>(sandbox);
    uint64_t hFile = 0;
    uint64_t lpBuffer = 0;
    uint32_t nNumberOfBytesToRead = 0;
    uint64_t lpNumberOfBytesRead = 0;
    uint64_t lpOverlapped = 0;

    // 获取参数
    if (context->GetPeInfo()->isX64) {
        // x64: rcx = hFile, rdx = lpBuffer, r8 = nNumberOfBytesToRead, r9 = lpNumberOfBytesRead
        uc_reg_read(uc, UC_X86_REG_RCX, &hFile);
        uc_reg_read(uc, UC_X86_REG_RDX, &lpBuffer);
        uint64_t temp_bytes_to_read;
        uc_reg_read(uc, UC_X86_REG_R8, &temp_bytes_to_read);
        nNumberOfBytesToRead = static_cast<uint32_t>(temp_bytes_to_read);
        uc_reg_read(uc, UC_X86_REG_R9, &lpNumberOfBytesRead);

        // 从栈上读取lpOverlapped参数
        uint64_t rsp;
        uc_reg_read(uc, UC_X86_REG_RSP, &rsp);
        uc_mem_read(uc, rsp + 0x28, &lpOverlapped, sizeof(uint64_t));
    } else {
        // x86: 从栈上读取参数
        uint32_t esp;
        uc_reg_read(uc, UC_X86_REG_ESP, &esp);
        esp += 0x4;  // 跳过返回地址

        uint32_t temp_handle, temp_buffer, temp_bytes_read, temp_overlapped;
        uc_mem_read(uc, esp, &temp_handle, sizeof(uint32_t));
        uc_mem_read(uc, esp + 0x4, &temp_buffer, sizeof(uint32_t));
        uc_mem_read(uc, esp + 0x8, &nNumberOfBytesToRead, sizeof(uint32_t));
        uc_mem_read(uc, esp + 0xC, &temp_bytes_read, sizeof(uint32_t));
        uc_mem_read(uc, esp + 0x10, &temp_overlapped, sizeof(uint32_t));

        hFile = temp_handle;
        lpBuffer = temp_buffer;
        lpNumberOfBytesRead = temp_bytes_read;
        lpOverlapped = temp_overlapped;
    }

    // 从文件句柄表中查找句柄信息
    auto fileInfo = context->GetFileHandleInfo(hFile);
    bool success = false;
    uint32_t bytesRead = 0;

    if (fileInfo && fileInfo->isValid) {
        // 计算实际可读取的字节数
        size_t availableBytes = fileInfo->content.size() - fileInfo->currentPosition;
        bytesRead = (nNumberOfBytesToRead < availableBytes) ? nNumberOfBytesToRead : (uint32_t)availableBytes;

        if (bytesRead > 0) {
            // 读取数据
            uc_mem_write(uc, lpBuffer, 
                fileInfo->content.data() + fileInfo->currentPosition, 
                bytesRead);
            
            // 更新当前位置
            fileInfo->currentPosition += bytesRead;
            success = true;
        } else {
            // 没有更多数据可读
            success = true; // 返回TRUE，但是bytesRead为0
        }
    } else {
        // 设置错误码
        DWORD error = ERROR_INVALID_HANDLE;
        if (context->GetPeInfo()->isX64) {
            context->GetTeb64()->LastErrorValue = error;
        } else {
            context->GetTeb32()->LastErrorValue = error;
        }
    }

    // 写入实际读取的字节数
    if (lpNumberOfBytesRead != 0) {
        uc_mem_write(uc, lpNumberOfBytesRead, &bytesRead, sizeof(uint32_t));
    }

    // 设置返回值
    uint64_t result = success ? TRUE : FALSE;
    uc_reg_write(uc, 
        context->GetPeInfo()->isX64 ? UC_X86_REG_RAX : UC_X86_REG_EAX, 
        &result);

    const char* apiName = isWideChar ? "ReadFileW" : "ReadFileA";
    printf("[*] %s: Handle=0x%llx, Buffer=0x%llx, BytesToRead=%u, BytesRead=%u, Result=%s\n",
        apiName, hFile, lpBuffer, nNumberOfBytesToRead, bytesRead, 
        success ? "TRUE" : "FALSE");
}

// CloseFile API实现
auto Api_CloseFile(void* sandbox, uc_engine* uc, uint64_t address) -> void {
    auto context = static_cast<Sandbox*>(sandbox);
    uint64_t hFile = 0;

    // 获取参数
    if (context->GetPeInfo()->isX64) {
        uc_reg_read(uc, UC_X86_REG_RCX, &hFile);
    } else {
        uint32_t esp;
        uc_reg_read(uc, UC_X86_REG_ESP, &esp);
        esp += 0x4;  // 跳过返回地址
        uint32_t temp_handle;
        uc_mem_read(uc, esp, &temp_handle, sizeof(uint32_t));
        hFile = temp_handle;
    }

    // 在文件句柄表中查找并释放资源
    auto it = context->m_fileHandles.find(hFile);
    bool success = false;

    if (it != context->m_fileHandles.end() && it->second) {
        delete it->second;
        context->m_fileHandles.erase(it);
        success = true;
    }

    // 设置返回值
    uint64_t result = success ? TRUE : FALSE;
    uc_reg_write(uc, 
        context->GetPeInfo()->isX64 ? UC_X86_REG_RAX : UC_X86_REG_EAX, 
        &result);

    printf("[*] CloseFile: Handle=0x%llx, Result=%s\n", 
        hFile, success ? "TRUE" : "FALSE");
}