#include <windows.h>
#include <stdio.h>
#include <string.h>

// Функция для парсинга PEB и поиска WinExec
FARPROC FindWinExec() {
    // Получаем PEB через FS регистр
    #ifdef _WIN64
        PPEB pPEB = (PPEB)__readgsqword(0x60);
    #else
        PPEB pPEB = (PPEB)__readfsdword(0x30);
    #endif
    
    // Получаем LDR (Loader Data)
    PPEB_LDR_DATA pLDR = pPEB->Ldr;
    
    // Проходим по списку загруженных модулей
    LIST_ENTRY* pModuleList = &pLDR->InMemoryOrderModuleList;
    LIST_ENTRY* pEntry = pModuleList->Flink;
    
    // Ищем kernel32.dll
    while (pEntry != pModuleList) {
        PLDR_DATA_TABLE_ENTRY pModule = CONTAINING_RECORD(pEntry, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks);
        
        // Проверяем имя модуля
        WCHAR* moduleName = pModule->BaseDllName.Buffer;
        if (moduleName && wcsstr(moduleName, L"kernel32.dll")) {
            HMODULE hKernel32 = (HMODULE)pModule->DllBase;
            
            // Ищем WinExec в kernel32.dll
            FARPROC pWinExec = GetProcAddress(hKernel32, "WinExec");
            if (pWinExec) {
                printf("[+] Found WinExec at: 0x%p\n", pWinExec);
                return pWinExec;
            }
        }
        pEntry = pEntry->Flink;
    }
    
    return NULL;
}

// Функция для открытия и парсинга .exe файла
void ParseExeFile(const char* filename) {
    HANDLE hFile = CreateFileA(
        filename,
        GENERIC_READ,
        FILE_SHARE_READ,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL
    );
    
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("[-] Failed to open file: %s\n", filename);
        return;
    }
    
    // Читаем DOS заголовок
    IMAGE_DOS_HEADER dosHeader;
    DWORD bytesRead;
    if (!ReadFile(hFile, &dosHeader, sizeof(dosHeader), &bytesRead, NULL)) {
        printf("[-] Failed to read DOS header\n");
        CloseHandle(hFile);
        return;
    }
    
    // Проверяем сигнатуру DOS
    if (dosHeader.e_magic != IMAGE_DOS_SIGNATURE) {
        printf("[-] Invalid DOS signature\n");
        CloseHandle(hFile);
        return;
    }
    
    // Переходим к PE заголовку
    SetFilePointer(hFile, dosHeader.e_lfanew, NULL, FILE_BEGIN);
    
    // Читаем PE сигнатуру
    DWORD peSignature;
    ReadFile(hFile, &peSignature, sizeof(peSignature), &bytesRead, NULL);
    
    if (peSignature != IMAGE_NT_SIGNATURE) {
        printf("[-] Invalid PE signature\n");
        CloseHandle(hFile);
        return;
    }
    
    // Читаем файловый заголовок
    IMAGE_FILE_HEADER fileHeader;
    ReadFile(hFile, &fileHeader, sizeof(fileHeader), &bytesRead, NULL);
    
    printf("[+] File is a valid PE executable\n");
    printf("[+] Number of sections: %d\n", fileHeader.NumberOfSections);
    printf("[+] Size of optional header: %d\n", fileHeader.SizeOfOptionalHeader);
    
    // Читаем опциональный заголовок
    IMAGE_OPTIONAL_HEADER32 optionalHeader;
    ReadFile(hFile, &optionalHeader, sizeof(optionalHeader), &bytesRead, NULL);
    
    printf("[+] Entry point: 0x%X\n", optionalHeader.AddressOfEntryPoint);
    printf("[+] Image base: 0x%X\n", optionalHeader.ImageBase);
    
    // Читаем секции
    printf("\n[+] Sections:\n");
    for (int i = 0; i < fileHeader.NumberOfSections; i++) {
        IMAGE_SECTION_HEADER sectionHeader;
        ReadFile(hFile, &sectionHeader, sizeof(sectionHeader), &bytesRead, NULL);
        
        printf("    %s - VA: 0x%X, Size: 0x%X\n", 
               sectionHeader.Name, 
               sectionHeader.VirtualAddress, 
               sectionHeader.SizeOfRawData);
    }
    
    CloseHandle(hFile);
}

// Функция-победитель (win)
void win(void) {
    printf("flag{ret2win_but_its_WINDOWS}\n");
    
    // Находим WinExec через PEB
    FARPROC pWinExec = FindWinExec();
    if (pWinExec) {
        // Запускаем калькулятор через WinExec
        typedef void (*WinExec_t)(LPCSTR, UINT);
        WinExec_t WinExec_func = (WinExec_t)pWinExec;
        WinExec_func("mspaint.exe", SW_SHOW);
        printf("[+] paint launched!\n");
    }
}

// Уязвимая функция
void vulnerable_function() {
    char buf[16];
    
    printf("enter your data:\n");
    scanf("%s", buf);
    
    printf("try again\n");
}

int main(int argc, char* argv[]) {
    printf("=== Windows Buffer Overflow CTF Challenge ===\n\n");
    
    // Если передан аргумент, парсим .exe файл
    if (argc > 1) {
        printf("[*] Parsing PE file: %s\n", argv[1]);
        ParseExeFile(argv[1]);
        printf("\n");
    }
    
    // Демонстрируем поиск WinExec через PEB
    printf("[*] Finding WinExec via PEB parsing...\n");
    FARPROC pWinExec = FindWinExec();
    if (pWinExec) {
        printf("[+] WinExec found at: 0x%p\n", pWinExec);
    }
    printf("\n");
    
    // Вызываем уязвимую функцию
    vulnerable_function();
    
    return 0;
}