#include <windows.h>
#include <winternl.h>
#include <iostream>
#include <iomanip>

#ifndef NT_SUCCESS
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
#endif

// Function pointer types for the NT functions
typedef NTSTATUS(NTAPI* PFN_NtCreateKey)(
    PHANDLE KeyHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES ObjectAttributes,
    ULONG TitleIndex, PUNICODE_STRING Class, ULONG CreateOptions, PULONG Disposition
    );
typedef NTSTATUS(NTAPI* PFN_NtSetValueKey)(
    HANDLE KeyHandle, PUNICODE_STRING ValueName, ULONG TitleIndex, ULONG Type,
    PVOID Data, ULONG DataSize
    );
typedef NTSTATUS(NTAPI* PFN_NtClose)(HANDLE Handle);
typedef VOID(NTAPI* PFN_RtlInitUnicodeString)(PUNICODE_STRING DestinationString, PCWSTR SourceString);

int main()
{
    std::wcout << L"[INFO] Starting NT service registry key creation process..." << std::endl;

    HMODULE hNtdll = LoadLibraryA("ntdll.dll");
    if (!hNtdll) {
        std::wcerr << L"[ERROR] Failed to get ntdll.dll handle." << std::endl;
        return 1;
    }
    std::wcout << L"[INFO] ntdll.dll handle obtained: " << hNtdll << std::endl;

    PFN_RtlInitUnicodeString pRtlInitUnicodeString =
        (PFN_RtlInitUnicodeString)GetProcAddress(hNtdll, "RtlInitUnicodeString");
    PFN_NtCreateKey pNtCreateKey =
        (PFN_NtCreateKey)GetProcAddress(hNtdll, "NtCreateKey");
    PFN_NtSetValueKey pNtSetValueKey =
        (PFN_NtSetValueKey)GetProcAddress(hNtdll, "NtSetValueKey");
    PFN_NtClose pNtClose =
        (PFN_NtClose)GetProcAddress(hNtdll, "NtClose");

    if (!pRtlInitUnicodeString || !pNtCreateKey || !pNtSetValueKey || !pNtClose) {
        std::wcerr << L"[ERROR] Failed to get one or more function pointers from ntdll.dll." << std::endl;
        return 1;
    }
    std::wcout << L"[INFO] Successfully loaded function pointers from ntdll.dll." << std::endl;

    NTSTATUS status;
    HANDLE hKey = NULL;
    ULONG disposition = 0;
    UNICODE_STRING regPath;
    pRtlInitUnicodeString(&regPath, L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\MyService"); // change name
    std::wcout << L"[INFO] Registry path initialized: " << regPath.Buffer << std::endl;

    OBJECT_ATTRIBUTES objAttr;
    InitializeObjectAttributes(&objAttr, &regPath, OBJ_CASE_INSENSITIVE, NULL, NULL);
    std::wcout << L"[INFO] OBJECT_ATTRIBUTES initialized." << std::endl;

    std::wcout << L"[INFO] Attempting to create/open registry key..." << std::endl;
    status = pNtCreateKey(&hKey, KEY_ALL_ACCESS, &objAttr, 0, NULL, REG_OPTION_NON_VOLATILE, &disposition);
    if (!NT_SUCCESS(status)) {
        std::wcerr << L"[ERROR] NtCreateKey failed: 0x" << std::hex << status << std::endl;
        return 1;
    }
    std::wcout << L"[INFO] Registry key created successfully. Disposition: 0x"
        << std::hex << disposition << std::endl;

    auto setValue = [&](const wchar_t* name, ULONG type, const void* data, ULONG dataSize) -> NTSTATUS {
        std::wcout << L"[INFO] Setting registry value: " << name << std::endl;
        UNICODE_STRING uName;
        pRtlInitUnicodeString(&uName, name);
        NTSTATUS st = pNtSetValueKey(hKey, &uName, 0, type, (PVOID)data, dataSize);
        if (!NT_SUCCESS(st))
            std::wcerr << L"[ERROR] NtSetValueKey (" << name << L") failed: 0x"
            << std::hex << st << std::endl;
        else
            std::wcout << L"[INFO] " << name << L" set successfully." << std::endl;
        return st;
        };

    const wchar_t* imagePath = L"%SystemRoot%\\System32\\myservice.exe"; // change name
    DWORD serviceType = 0x10, startType = 2, errorControl = 1;
    if (!NT_SUCCESS(setValue(L"ImagePath", REG_EXPAND_SZ, imagePath, (wcslen(imagePath) + 1) * sizeof(wchar_t))) ||
        !NT_SUCCESS(setValue(L"Type", REG_DWORD, &serviceType, sizeof(serviceType))) ||
        !NT_SUCCESS(setValue(L"Start", REG_DWORD, &startType, sizeof(startType))) ||
        !NT_SUCCESS(setValue(L"ErrorControl", REG_DWORD, &errorControl, sizeof(errorControl))) ||
        !NT_SUCCESS(setValue(L"ObjectName", REG_SZ, L"LocalSystem", (wcslen(L"LocalSystem") + 1) * sizeof(wchar_t))))
    {
        pNtClose(hKey);
        std::wcerr << L"[ERROR] One or more registry values failed to setc." << std::endl;
        return 1;
    }

    std::wcout << L"[INFO] All registry values set successfully." << std::endl;
    pNtClose(hKey);
    std::wcout << L"[INFO] Registry key handle closed." << std::endl;
    return 0;
}