#include "../MemoryModule/LoadDllMemoryApi.h"
#include <cstdio>
#pragma warning(disable:4996)

static PVOID ReadDllFile(LPCSTR FileName) {
    LPVOID buffer;
    size_t size;
    FILE* f = fopen(FileName, "rb");
    if (!f)return 0;
    _fseeki64(f, 0, SEEK_END);
    if (!(size = _ftelli64(f))) {
        fclose(f);
        return 0;
    }
    _fseeki64(f, 0, SEEK_SET);
    fread(buffer = new char[size], 1, size, f);
    fclose(f);
    return buffer;
}

int test_default() {
    LPVOID buffer = ReadDllFile("a.dll");
    
    HMEMORYMODULE m1 = nullptr, m2 = m1;
    HMODULE hModule = nullptr;
    FARPROC pfn = nullptr;
    DWORD MemoryModuleFeatures = 0;

    typedef int(* _exception)(int code);
    _exception exception = nullptr;
    HRSRC hRsrc;
    DWORD SizeofRes;
    HGLOBAL gRes;
    char str[10];

    NtQuerySystemMemoryModuleFeatures(&MemoryModuleFeatures);
    if (MemoryModuleFeatures != MEMORY_FEATURE_ALL) {
        printf("not support all features on this version of windows.\n");
    }
    
    if (!NT_SUCCESS(NtLoadDllMemoryExW(&m1, nullptr, 0, buffer, 0, L"kernel64", nullptr))) goto end;
    LoadLibraryW(L"wininet.dll");
    if (!NT_SUCCESS(NtLoadDllMemoryExW(&m2, nullptr, 0, buffer, 0, L"kernel128", nullptr))) goto end;

    //forward export
    hModule = (HMODULE)m1;
    pfn = (decltype(pfn))(GetProcAddress(hModule, "Socket")); //ws2_32.WSASocketW
    pfn = (decltype(pfn))(GetProcAddress(hModule, "VerifyTruse")); //wintrust.WinVerifyTrust
    hModule = (HMODULE)m2;
    pfn = (decltype(pfn))(GetProcAddress(hModule, "Socket"));
    pfn = (decltype(pfn))(GetProcAddress(hModule, "VerifyTruse"));

    //exception
    hModule = (HMODULE)m1;
    exception = (_exception)GetProcAddress(hModule, "exception");
    if (exception) {
        for (int i = 0; i < 4; ++i)exception(i);
    }
    
    //tls
    pfn = GetProcAddress(hModule, "thread");
    if (pfn && pfn()) {
        printf("thread test failed.\n");
    }

    //resource
    if (!LoadStringA(hModule, 101, str, 10)) {
        printf("load string failed.\n");
    }
    else {
        printf("%s\n", str);
    }
    if (!(hRsrc = FindResourceA(hModule, MAKEINTRESOURCEA(102), "BINARY"))) {
        printf("find binary resource failed.\n");
    }
    else {
        if ((SizeofRes = SizeofResource(hModule, hRsrc)) != 0x10) {
            printf("invalid res size.\n");
        }
        else {
            if (!(gRes = LoadResource(hModule, hRsrc))) {
                printf("load res failed.\n");
            }
            else {
                if (!LockResource(gRes))printf("lock res failed.\n");
                else {
                    printf("resource test success.\n");
                }
            }
        }
    }

end:
    delete[]buffer;
    if (m1)NtUnloadDllMemory(m1);
    FreeLibrary(LoadLibraryW(L"wininet.dll"));
    FreeLibrary(GetModuleHandleW(L"wininet.dll"));
    if (m2)NtUnloadDllMemory(m2);

    return 0;
}

#define WSADESCRIPTION_LEN      256
#define WSASYS_STATUS_LEN       128
typedef USHORT ADDRESS_FAMILY;
typedef int (PASCAL* WSAStartup_t)(WORD wVersionRequired, LPWSADATA lpWSAData);
typedef int (PASCAL* WSACleanup_t)(void);
typedef SOCKET (PASCAL* socket_t)(int af, int type, int protocol);
typedef int (PASCAL* closesocket_t)(SOCKET s);
typedef int (PASCAL* connect_t)(SOCKET s, const struct sockaddr FAR* name, int namelen);
typedef unsigned long (PASCAL* inet_addr_t)(const char FAR* cp);
typedef u_short (PASCAL* htons_t)(u_short hostshort);

int test_ws2_32() {
    PVOID buffer = ReadDllFile("C:\\Windows\\system32\\ws2_32.dll");

    HMEMORYMODULE hMemoryModule = nullptr;
    HMODULE hModule = nullptr;
    NTSTATUS status;

    WSAData data{};
    SOCKET sock = INVALID_SOCKET;
    sockaddr_in addr{};
    WSAStartup_t _WSAStartup = nullptr;
    WSACleanup_t _WSACleanup = nullptr;
    socket_t _socket = nullptr;
    closesocket_t _closesocket = nullptr;
    connect_t _connect = nullptr;
    inet_addr_t _inet_addr = nullptr;
    htons_t _htons = nullptr;

    hMemoryModule = LoadLibraryMemoryExW(buffer, 0, L"ws2.dll", nullptr, LOAD_FLAGS_NOT_FAIL_IF_HANDLE_TLS);
    hModule = MemoryModuleToModule(hMemoryModule);
    if (buffer)delete[]buffer;
    if (!hModule)return 0;

    _WSAStartup = (decltype(_WSAStartup)(GetProcAddress(hModule, "WSAStartup")));
    _WSACleanup = (decltype(_WSACleanup)(GetProcAddress(hModule, "WSACleanup")));
    _socket = (decltype(_socket)(GetProcAddress(hModule, "socket")));
    _closesocket = (decltype(_closesocket)(GetProcAddress(hModule, "closesocket")));
    _connect = (decltype(_connect)(GetProcAddress(hModule, "connect")));
    _inet_addr = (decltype(_inet_addr)(GetProcAddress(hModule, "inet_addr")));
    _htons = (decltype(_htons)(GetProcAddress(hModule, "htons")));
    if (!_WSAStartup || !_WSACleanup || !_socket || !_closesocket || !_connect || !_inet_addr || !_htons)goto end;

    if (_WSAStartup(MAKEWORD(2, 2), &data) != 0)goto end;
    if ((sock = _socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) == INVALID_SOCKET)goto end;
    addr.sin_family = AF_INET;
    addr.sin_port = _htons(80);
    addr.sin_addr.S_un.S_addr = _inet_addr("1.1.1.1");
    if (_connect(sock, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)goto end;

    //success
    printf("ws2_32 completed successfully.\n");

end:
    if (sock != INVALID_SOCKET && _closesocket)_closesocket(sock);
    if (_WSACleanup)_WSACleanup();
    FreeLibraryMemory(hMemoryModule);
    return 0;
}

int main() {
    DWORD t;
    NtQuerySystemMemoryModuleFeatures(&t);
    LoadLibraryMemory(ReadDllFile("C:\\Windows\\explorer.exe"));
    return 0;

    test_default();
    test_ws2_32();
    
    return 0;
}