#include "utils.h"

namespace utils
{
    // 读取二进制文件到内存
    BYTE* ReadBinaryFile(const char* filePath, DWORD* outSize) {
        HANDLE hFile = CreateFileA(
            filePath,
            GENERIC_READ,
            FILE_SHARE_READ,
            NULL,
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            NULL
        );

        if (hFile == INVALID_HANDLE_VALUE) {
            printf("打开文件失败: %s\n", filePath);
            return NULL;
        }

        DWORD fileSize = GetFileSize(hFile, NULL);
        if (fileSize == INVALID_FILE_SIZE) {
            printf("获取文件大小失败\n");
            CloseHandle(hFile);
            return NULL;
        }

        BYTE* buffer = (BYTE*)malloc(fileSize);
        if (!buffer) {
            printf("内存分配失败\n");
            CloseHandle(hFile);
            return NULL;
        }

        DWORD bytesRead = 0;
        BOOL result = ReadFile(hFile, buffer, fileSize, &bytesRead, NULL);
        CloseHandle(hFile);

        if (!result || bytesRead != fileSize) {
            printf("读取文件失败\n");
            free(buffer);
            return NULL;
        }

        *outSize = bytesRead;
        return buffer;
    }
    BOOL WriteShellcodeToFile(const char* filePath, const BYTE* shellcode, SIZE_T shellcodeSize) {
        // 1. 创建文件
        HANDLE hFile = CreateFileA(
            filePath,
            GENERIC_WRITE,
            0,                      // 独占访问
            NULL,
            CREATE_ALWAYS,          // 总是创建新文件
            FILE_ATTRIBUTE_NORMAL,  // 普通文件
            NULL
        );

        if (hFile == INVALID_HANDLE_VALUE) {
            //printf("CreateFile 失败: %d\n", GetLastError());
            return FALSE;
        }

        // 2. 写入 shellcode
        DWORD bytesWritten = 0;
        BOOL result = WriteFile(hFile, shellcode, (DWORD)shellcodeSize, &bytesWritten, NULL);

        if (!result || bytesWritten != shellcodeSize) {
            //printf("WriteFile 失败: %d\n", GetLastError());
            CloseHandle(hFile);
            return FALSE;
        }

        // 3. 刷新缓冲区并关闭
        FlushFileBuffers(hFile);
        CloseHandle(hFile);

        //printf("Shellcode 写入成功: %s (大小: %d 字节)\n", filePath, bytesWritten);
        return TRUE;
    }
	// 创建多级目录（递归创建）
	BOOL CreateDirectoryRecursive(const char* path) {
		char tempPath[MAX_PATH];
		strcpy_s(tempPath, path);

		// 去掉末尾的反斜杠
		size_t len = strlen(tempPath);
		if (len > 0 && tempPath[len - 1] == '\\') {
			tempPath[len - 1] = '\0';
		}

		// 检查目录是否已存在
		DWORD attrs = GetFileAttributesA(tempPath);
		if (attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY)) {
			return TRUE;  // 目录已存在
		}

		// 递归创建父目录
		char* lastSlash = strrchr(tempPath, '\\');
		if (lastSlash) {
			*lastSlash = '\0';
			if (!CreateDirectoryRecursive(tempPath)) {
				return FALSE;
			}
			*lastSlash = '\\';
		}

		// 创建当前目录
		if (!CreateDirectoryA(tempPath, NULL)) {
			if (GetLastError() == ERROR_ALREADY_EXISTS) {
				return TRUE;
			}
			//printf("创建目录失败: %s (错误: %d)\n", tempPath, GetLastError());
			return FALSE;
		}

		//printf("创建目录成功: %s\n", tempPath);
		return TRUE;
	}
    auto GetProcessByName(CONST wchar_t* Name) -> DWORD
    {
        DWORD Pid = 0;

       auto SnapshotHandle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
       if (SnapshotHandle == INVALID_HANDLE_VALUE)
           return Pid;

       PROCESSENTRY32 pe32{ 0 };
       pe32.dwSize = sizeof(pe32);

       if (Process32First(SnapshotHandle, &pe32))
       {
           do
           {
               if (!wcscmp(pe32.szExeFile, Name)) 
               {
                   Pid = pe32.th32ProcessID;
                   break;
               }
 
           } while (Process32Next(SnapshotHandle, &pe32));
       }

       CloseHandle(SnapshotHandle);

       return Pid;

    }

    // 不区分大小写查找子串
    bool StrContainsI(const char* src, const char* sub)
    {
        char* p = strstr(_strlwr((char*)src), _strlwr((char*)sub));
        return p != nullptr;
    }

    // 不区分大小写宽字符串查找
    bool StrContainsI(LPCWSTR Source, LPCWSTR Sub)
    {
        if (!Source || !Sub)
            return FALSE;
        WCHAR srcBuf[1024] = { 0 };
        WCHAR subBuf[1024] = { 0 };
        lstrcpyW(srcBuf, Source);
        lstrcpyW(subBuf, Sub);
        _wcslwr(srcBuf);
        _wcslwr(subBuf);
        return wcsstr(srcBuf, subBuf) != nullptr;
    }


    // 读取注册表字符串值
    BOOL RegReadString(HKEY hRoot, LPCSTR szSubKey, LPCSTR szValueName, char* outBuf, DWORD bufSize)
    {
        HKEY hKey;
        LONG ret = RegOpenKeyExA(hRoot, szSubKey, 0, KEY_READ, &hKey);
        if (ret != ERROR_SUCCESS)
            return FALSE;

        DWORD dataType = REG_SZ;
        DWORD dataLen = bufSize;
        ret = RegQueryValueExA(hKey, szValueName, nullptr, &dataType, (LPBYTE)outBuf, &dataLen);
        RegCloseKey(hKey);

        if (ret != ERROR_SUCCESS)
            return FALSE;
        outBuf[dataLen] = '\0';
        return TRUE;
    }

    void GetProcessName(
        DWORD ProcessId,
        char* ProcessName,
        DWORD ProcessNameSize
    )
    {
        ProcessName[0] = '\0';

        HANDLE ProcessHandle =
            OpenProcess(
                PROCESS_QUERY_LIMITED_INFORMATION,
                FALSE,
                ProcessId
            );

        if (!ProcessHandle)
        {
            strcpy_s(
                ProcessName,
                ProcessNameSize,
                "Unknown"
            );

            return;
        }

        DWORD Size = ProcessNameSize;

        if (!QueryFullProcessImageNameA(
            ProcessHandle,
            0,
            ProcessName,
            &Size))
        {
            strcpy_s(
                ProcessName,
                ProcessNameSize,
                "Unknown"
            );
        }
        else
        {
            char* FileName =
                strrchr(ProcessName, '\\');

            if (FileName)
            {
                memmove(
                    ProcessName,
                    FileName + 1,
                    strlen(FileName)
                );
            }
        }

        CloseHandle(ProcessHandle);
    }
   
    
}