#include "target_process.h"
#include <winternl.h>
#include <string.h>

typedef NTSTATUS(NTAPI* NtGetNextProcess_t)(
    HANDLE ProcessHandle,
    ACCESS_MASK DesiredAccess,
    ULONG HandleAttributes,
    ULONG Flags,
    PHANDLE NewProcessHandle
    );

typedef NTSTATUS(NTAPI* NtQueryInformationProcess_t)(
    HANDLE ProcessHandle,
    PROCESSINFOCLASS ProcessInformationClass,
    PVOID ProcessInformation,
    ULONG ProcessInformationLength,
    PULONG ReturnLength
    );

typedef NTSTATUS(NTAPI* NtReadVirtualMemory_t)(
    HANDLE ProcessHandle,
    PVOID BaseAddress,
    PVOID Buffer,
    SIZE_T BufferSize,
    PSIZE_T NumberOfBytesRead
    );

typedef NTSTATUS(NTAPI* NtTerminateProcess_t)(
    HANDLE ProcessHandle,
    NTSTATUS ExitStatus
    );

typedef BOOL(WINAPI* IsWow64Process2_t)(HANDLE, USHORT*, USHORT*);

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

static bool check_architecture(TargetProcess* target) {
    if (!target || !target->process) return false;

    HMODULE kernel32 = GetModuleHandleW(L"kernel32.dll");
    if (!kernel32) return false;

    IsWow64Process2_t fnIsWow64Process2 =
        (IsWow64Process2_t)GetProcAddress(kernel32, "IsWow64Process2");

    if (!fnIsWow64Process2) {
        BOOL is_wow64 = FALSE;
        if (IsWow64Process(target->process, &is_wow64)) {
            target->arch = is_wow64 ? IMAGE_FILE_MACHINE_I386 : IMAGE_FILE_MACHINE_AMD64;
            return true;
        }
        return false;
    }

    USHORT process_arch = 0, native_machine = 0;
    if (!fnIsWow64Process2(target->process, &process_arch, &native_machine)) {
        return false;
    }

    target->arch = (process_arch == IMAGE_FILE_MACHINE_UNKNOWN) ?
        native_machine : process_arch;

    USHORT injector_arch =
#if defined(_M_X64)
        IMAGE_FILE_MACHINE_AMD64;
#elif defined(_M_ARM64)
        IMAGE_FILE_MACHINE_ARM64;
#elif defined(_M_IX86)
        IMAGE_FILE_MACHINE_I386;
#else
        IMAGE_FILE_MACHINE_UNKNOWN;
#endif

    return target->arch == injector_arch;
}

bool target_create_suspended(TargetProcess* target, const wchar_t* exe_path) {
    if (!target || !exe_path) return false;

    memset(target, 0, sizeof(TargetProcess));

    STARTUPINFOW si = { 0 };
    PROCESS_INFORMATION pi = { 0 };
    si.cb = sizeof(si);

    if (!CreateProcessW(exe_path, NULL, NULL, NULL, FALSE,
        CREATE_SUSPENDED, NULL, NULL, &si, &pi)) {
        return false;
    }

    target->process = pi.hProcess;
    target->thread = pi.hThread;
    target->pid = pi.dwProcessId;

    if (!check_architecture(target)) {
        target_cleanup(target);
        return false;
    }

    return true;
}

void target_terminate(TargetProcess* target) {
    if (!target || !target->process) return;

    HMODULE ntdll = GetModuleHandleW(L"ntdll.dll");
    if (ntdll) {
        NtTerminateProcess_t pNtTerminateProcess =
            (NtTerminateProcess_t)GetProcAddress(ntdll, "NtTerminateProcess");
        if (pNtTerminateProcess) {
            pNtTerminateProcess(target->process, 0);
        }
    }
    else {
        TerminateProcess(target->process, 0);
    }

    WaitForSingleObject(target->process, 2000);
}

void target_cleanup(TargetProcess* target) {
    if (!target) return;

    if (target->thread) {
        CloseHandle(target->thread);
        target->thread = NULL;
    }

    if (target->process) {
        CloseHandle(target->process);
        target->process = NULL;
    }
}

bool target_kill_network_service(const wchar_t* process_name) {
    if (!process_name) return false;

    HMODULE ntdll = GetModuleHandleW(L"ntdll.dll");
    if (!ntdll) return false;

    NtGetNextProcess_t pNtGetNextProcess = (NtGetNextProcess_t)GetProcAddress(ntdll, "NtGetNextProcess");
    NtQueryInformationProcess_t pNtQueryInformationProcess = (NtQueryInformationProcess_t)GetProcAddress(ntdll, "NtQueryInformationProcess");
    NtReadVirtualMemory_t pNtReadVirtualMemory = (NtReadVirtualMemory_t)GetProcAddress(ntdll, "NtReadVirtualMemory");
    NtTerminateProcess_t pNtTerminateProcess = (NtTerminateProcess_t)GetProcAddress(ntdll, "NtTerminateProcess");

    if (!pNtGetNextProcess || !pNtQueryInformationProcess ||
        !pNtReadVirtualMemory || !pNtTerminateProcess) {
        return false;
    }

    HANDLE current_process = NULL;
    bool found = false;

    while (true) {
        HANDLE next_process = NULL;
        NTSTATUS status = pNtGetNextProcess(current_process, PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_TERMINATE, 0, 0, &next_process);

        if (current_process) {
            CloseHandle(current_process);
            current_process = NULL;
        }

        if (!NT_SUCCESS(status)) break;

        current_process = next_process;

        BYTE buffer[sizeof(UNICODE_STRING) + MAX_PATH * 2];
        PUNICODE_STRING image_name = (PUNICODE_STRING)buffer;

        status = pNtQueryInformationProcess(current_process, ProcessImageFileName, image_name, sizeof(buffer), NULL);

        if (!NT_SUCCESS(status) || image_name->Length == 0) continue;

        wchar_t* filename = image_name->Buffer;
        for (wchar_t* p = image_name->Buffer + (image_name->Length / sizeof(wchar_t)) - 1; p > image_name->Buffer; p--) {
            if (*p == L'\\' || *p == L'/') {
                filename = p + 1;
                break;
            }
        }

        if (_wcsicmp(filename, process_name) != 0) continue;

        PROCESS_BASIC_INFORMATION pbi = { 0 };
        status = pNtQueryInformationProcess(current_process, ProcessBasicInformation, &pbi, sizeof(pbi), NULL);

        if (!NT_SUCCESS(status) || !pbi.PebBaseAddress) continue;

        PEB peb = { 0 };
        SIZE_T bytes_read = 0;
        status = pNtReadVirtualMemory(current_process, pbi.PebBaseAddress, &peb, sizeof(peb), &bytes_read);

        if (!NT_SUCCESS(status)) continue;

        RTL_USER_PROCESS_PARAMETERS params = { 0 };
        status = pNtReadVirtualMemory(current_process, peb.ProcessParameters, &params, sizeof(params), &bytes_read);

        if (!NT_SUCCESS(status)) continue;

        if (params.CommandLine.Length > 0) {
            wchar_t* cmd_line = malloc(params.CommandLine.Length + sizeof(wchar_t));
            if (cmd_line) {
                memset(cmd_line, 0, params.CommandLine.Length + sizeof(wchar_t));

                status = pNtReadVirtualMemory(
                    current_process,
                    params.CommandLine.Buffer,
                    cmd_line,
                    params.CommandLine.Length,
                    &bytes_read
                );

                if (NT_SUCCESS(status)) {
                    if (wcsstr(cmd_line, L"--utility-sub-type=network.mojom.NetworkService")) {
                        pNtTerminateProcess(current_process, 0);
                        found = true;
                    }
                }

                free(cmd_line);
            }
        }
    }

    if (current_process) {
        CloseHandle(current_process);
    }

    return true;
}