/*
 * even.sys - a stub for Htsysm7679
 * Copyright (C) 2024 Zephyr Lykos <self@mochaa.ws>
 *
 * This program 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 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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 this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 *
 */

#include <ddk/ntddk.h>
#include "even_ioctl.h"
#include "even_names.h"

#define uthash_nonfatal_oom(obj) do {return STATUS_INTERNAL_ERROR;} while (0)

#include "even_peb.h"
#include "uthash.h"

/*
 * Htsysm7679 constatns
 */

USE_EVEN_IOCTL_OUT_DATA(1);
USE_EVEN_IOCTL_OUT_DATA(2);

const struct even_ioctl_out_data1 htsysm_check_1 = {.result = 1, .data = {0x64, 0xBC, 0xD4, 0x56}};
const struct even_ioctl_out_data2 htsysm_check_2 = {
    .result = 1,
    .data = {0x00, 0x00, 0x59, 0x26, 0x04, 0xF8, 0xFF, 0xFF},
};
const struct even_ioctl_out_data1 htsysm_check_3 = {.result = 1, .data = {0x48, 0xB3, 0xF7, 0x97}};

static VOID __stdcall even_unload(IN PDRIVER_OBJECT DriverObject)
{
    ANSI_STRING SymbolicLinkNameA;
    UNICODE_STRING SymbolicLinkNameW;

    DbgPrint("Even Htsysm7679: DriverUnload called\r\n"); // catch this using DBGVIEW from www.sysinternals.com

    /*RtlInitUnicodeString( &SymbolicLinkName, EVEN_DOSDEVICE_NAME );*/
    /* If I want to use normal string combining logic in my_names.h, I need to mess with ANSI vs. Unicode */
    RtlInitString(&SymbolicLinkNameA, EVEN_DOSDEVICE_NAME);
    RtlAnsiStringToUnicodeString(&SymbolicLinkNameW, &SymbolicLinkNameA,
                                 TRUE); /* Init from an ANSI string, and do allocate target buffer */

    IoDeleteSymbolicLink(&SymbolicLinkNameW);
    IoDeleteDevice(DriverObject->DeviceObject);
}

static NTSTATUS __stdcall even_dispatch_create(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return STATUS_SUCCESS;
}

struct even_peb *fake_peb = NULL;

static NTSTATUS __stdcall even_create_peb(HANDLE pid)
{
        struct even_peb *s;

        HASH_FIND_INT(fake_peb, &pid, s);
        if (s == NULL) {
            s = ExAllocatePool(PagedPool, sizeof(*s));
            s->pid = pid;

            PROCESS_BASIC_INFORMATION ProcessInformation;
            NtQueryInformationProcess(
                    pid,
                    ProcessBasicInformation,
                    &ProcessInformation,
                    sizeof(ProcessInformation),
                    NULL
            );

            RtlCopyMemory(&s->fake_peb, ProcessInformation.PebBaseAddress, sizeof(*ProcessInformation.PebBaseAddress));

            HASH_ADD_INT(fake_peb, pid, s);
        }

        s->refcnt += 1;

        return STATUS_SUCCESS;
}

static NTSTATUS __stdcall even_delete_peb(struct even_peb* peb)
{
    DbgPrint("Even Htsysm7679: PEB Refcnt for PID %d: %d\r\n", peb->pid, peb->refcnt);
    if (peb->refcnt <= 0) {
        DbgPrint("Even Htsysm7679: PEB Refcnt <= 0, deleting\r\n");
        HASH_DEL(fake_peb, peb);
        return STATUS_SUCCESS;
    }

    peb->refcnt -= 1;
    return STATUS_SUCCESS;
}

static NTSTATUS __stdcall *even_find_peb(IN HANDLE pid, OUT struct even_peb *s)
{
    HASH_FIND_INT(fake_peb, &pid, s);
    return STATUS_SUCCESS;
}

static NTSTATUS __stdcall even_open_close(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    UNREFERENCED_PARAMETER(DeviceObject);

    PIO_STACK_LOCATION io_stack;
    io_stack = IoGetCurrentIrpStackLocation(Irp);
    if (!(io_stack)) /* do we have a valid io_stack? */
    {
        Irp->IoStatus.Status = STATUS_INTERNAL_ERROR;
        goto err_out;
    }

    HANDLE pid = PsGetCurrentProcessId();
    switch (io_stack->MajorFunction)
    {
    case IRP_MJ_CREATE:
        if (pid == 0) goto err_out;
        DbgPrint("Even Htsysm7679: Handle to symbolink link opened by PID %d\r\n", pid);

        Irp->IoStatus.Information = 0;
        Irp->IoStatus.Status = even_create_peb(pid);
        IoCompleteRequest(Irp, IO_NO_INCREMENT);

        return STATUS_SUCCESS;
    case IRP_MJ_CLOSE:
        if (pid == 0) goto err_out;
        DbgPrint("Even Htsysm7679: Handle to symbolink link closed by PID %d\r\n", pid);

        struct even_peb s;
        even_find_peb(pid, &s);

        Irp->IoStatus.Information = 0;
        Irp->IoStatus.Status = even_delete_peb(&s);
        IoCompleteRequest(Irp, IO_NO_INCREMENT);

        return STATUS_SUCCESS;
    default:
        break;
    }

err_out:

    /* complete IRP request */
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return Irp->IoStatus.Status;
}

static NTSTATUS __stdcall even_dispatch_device_control(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    UNREFERENCED_PARAMETER(DeviceObject);

    PIO_STACK_LOCATION io_stack;

    /* default return values */
    Irp->IoStatus.Status = STATUS_NOT_IMPLEMENTED;
    Irp->IoStatus.Information = 0;

    /* get current I/O stack location */
    io_stack = IoGetCurrentIrpStackLocation(Irp);
    if (!(io_stack)) /* do we have a valid io_stack? */
    {
        Irp->IoStatus.Status = STATUS_INTERNAL_ERROR;
        goto err_out;
    }

    /* if this fires, we've hooked up the ioctl() handler to the wrong callback slot */
    if (io_stack->MajorFunction != IRP_MJ_DEVICE_CONTROL)
    {
        Irp->IoStatus.Status = STATUS_INTERNAL_ERROR;
        goto err_out;
    }

    switch (io_stack->Parameters.DeviceIoControl.IoControlCode)
    {
    case IOCTL_CHECK_1: {
        if (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(htsysm_check_1))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        RtlCopyBytes(Irp->AssociatedIrp.SystemBuffer, &htsysm_check_1, sizeof(htsysm_check_1));
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(htsysm_check_1);
        DbgPrint("Even Htsysm7679: Check 1\r\n");
    }
    break;

    case IOCTL_CHECK_2: {
        if (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(htsysm_check_2))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        RtlCopyBytes(Irp->AssociatedIrp.SystemBuffer, &htsysm_check_2, sizeof(htsysm_check_2));
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(htsysm_check_2);
        DbgPrint("Even Htsysm7679: Check 2\r\n");
    }
    break;

    case IOCTL_CHECK_3: {
        if (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(htsysm_check_3))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        RtlCopyBytes(Irp->AssociatedIrp.SystemBuffer, &htsysm_check_3, sizeof(htsysm_check_3));
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(htsysm_check_3);
        DbgPrint("Even Htsysm7679: Check 3\r\n");
    }
    break;

    case IOCTL_WRITE_PEB_1: {
        struct even_ioctl_in_data *even_in_data = (struct even_ioctl_in_data *)Irp->AssociatedIrp.SystemBuffer;
        struct even_ioctl_out_data *even_out_data = (struct even_ioctl_out_data *)Irp->AssociatedIrp.SystemBuffer;

        if ((io_stack->Parameters.DeviceIoControl.InputBufferLength < sizeof(struct even_ioctl_in_data)) ||
            (!even_in_data) ||
            (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(struct even_ioctl_out_data1)) ||
            (!even_out_data))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        struct even_peb s;

        if ((even_in_data->offset + 4) > sizeof(s.fake_peb))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        };

        HANDLE pid = PsGetCurrentProcessId();
        if (even_find_peb(pid, &s) != STATUS_SUCCESS) {
            Irp->IoStatus.Status = STATUS_INTERNAL_ERROR;
            goto err_out;
        }

        void* fake_peb;
        fake_peb = &s.fake_peb;

        int value;
        RtlCopyBytes(&value, even_in_data->value, sizeof(value));
        value = ~value;

        DbgPrint("Even Htsysm7679: Write PEB 1\r\n");
        DbgPrint("Even Htsysm7679: Offset = %0x, Value = %0x\r\n", even_in_data->offset, value);
        for (size_t i = 0; i < 4; i++)
        {
            ((char *)fake_peb)[even_in_data->offset + i] = ((char *)&value)[i];
        }

        even_out_data->result = 1;
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(*even_out_data);
    }
    break;

    case IOCTL_WRITE_PEB_2: {
        struct even_ioctl_in_data *even_in_data = (struct even_ioctl_in_data *)Irp->AssociatedIrp.SystemBuffer;
        struct even_ioctl_out_data *even_out_data = (struct even_ioctl_out_data *)Irp->AssociatedIrp.SystemBuffer;

        if ((io_stack->Parameters.DeviceIoControl.InputBufferLength < sizeof(struct even_ioctl_in_data)) ||
            (!even_in_data) ||
            (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(struct even_ioctl_out_data1)) ||
            (!even_out_data))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        struct even_peb s;

        if ((even_in_data->offset + 4) > sizeof(s.fake_peb))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        };

        HANDLE pid = PsGetCurrentProcessId();
        if (even_find_peb(pid, &s) != STATUS_SUCCESS) {
            Irp->IoStatus.Status = STATUS_INTERNAL_ERROR;
            goto err_out;
        }

        void* fake_peb;
        fake_peb = &s.fake_peb;

        int value;
        RtlCopyBytes(&value, even_in_data->value, sizeof(value));

        for (size_t i = 0; i < 4; i++)
        {
            ((char *)fake_peb)[even_in_data->offset + i] = ((char *)&value)[i];
        }

        even_out_data->result = 1;
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(*even_out_data);
        DbgPrint("Even Htsysm7679: Write PEB 2\r\n");
        DbgPrint("Even Htsysm7679: Offset = %0x, Value = %0x\r\n", even_in_data->offset, value);
    }
    break;

    case IOCTL_READ_PEB: {
        struct even_ioctl_in_data *even_in_data = (struct even_ioctl_in_data *)Irp->AssociatedIrp.SystemBuffer;
        struct even_ioctl_out_data1 *even_out_data = (struct even_ioctl_out_data1 *)Irp->AssociatedIrp.SystemBuffer;

        if ((io_stack->Parameters.DeviceIoControl.InputBufferLength < sizeof(struct even_ioctl_in_data)) ||
            (!even_in_data) ||
            (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(struct even_ioctl_out_data1)) ||
            (!even_out_data))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        struct even_peb s;

        if ((even_in_data->offset + 4) > sizeof(s.fake_peb))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        };

        HANDLE pid = PsGetCurrentProcessId();
        if (even_find_peb(pid, &s) != STATUS_SUCCESS) {
            Irp->IoStatus.Status = STATUS_INTERNAL_ERROR;
            goto err_out;
        }


        void* fake_peb;
        fake_peb = &s.fake_peb;

        char value[4];
        RtlCopyBytes(&value, fake_peb + (even_in_data->offset), sizeof(value));

        for (size_t i = 0; i < 4; i++)
        {
            value[i] &= value[i];
        }

        even_out_data->result = 1;
        RtlCopyMemory(even_out_data->data, value, sizeof(value));

        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(*even_out_data);
        DbgPrint("Even Htsysm7679: Read PEB\r\n");
        DbgPrint("Even Htsysm7679: Offset = %0x, Value = %0x\r\n", even_in_data->offset, *(int *)even_in_data->value);
        DbgPrint("Even Htsysm7679: Return = %0x\r\n", *(int *)value);
    }
    break;

    default: /* STATUS_NOT_IMPLEMENTED takes action */
        break;
    }

err_out:

    /* complete IRP request */
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return Irp->IoStatus.Status;
}

NTSTATUS __stdcall
DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath)
{
    NTSTATUS status;
    /* a particular Device Instance, theoretically one of many,
       managed by this driver : */
    PDEVICE_OBJECT DeviceObject;
    UNICODE_STRING DeviceNameW;
    UNICODE_STRING SymbolicLinkNameW;
    ANSI_STRING DeviceNameA;
    ANSI_STRING SymbolicLinkNameA;

    DbgPrint("Even Htsysm7679: DriverEntry Called\r\n"); // catch this using DBGVIEW from www.sysinternals.com

    /* support for service stopping */
    DriverObject->DriverUnload = even_unload;
    /* create support */
    DriverObject->MajorFunction[IRP_MJ_CREATE] = even_dispatch_create;
    /* IOCTL support */
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = even_dispatch_device_control;
    /* Open/close support */
    DriverObject->MajorFunction[IRP_MJ_CREATE] = even_open_close;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = even_open_close;

    /* initialize counted unicode strings */
    RtlInitString(&DeviceNameA, EVEN_DEVICE_NAME);
    RtlAnsiStringToUnicodeString(&DeviceNameW, &DeviceNameA,
                                 TRUE); /* Init from an ANSI string, and do allocate target buffer */
    RtlInitString(&SymbolicLinkNameA, EVEN_DOSDEVICE_NAME);
    RtlAnsiStringToUnicodeString(&SymbolicLinkNameW, &SymbolicLinkNameA,
                                 TRUE); /* Init from an ANSI string, and do allocate target buffer */

    /* create device object */
    status = IoCreateDevice(DriverObject, 0, &DeviceNameW, 0xaa01, 0, FALSE, &DeviceObject);
    if (!NT_SUCCESS(status))
        return status;

    DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;

    /* create user-visible name for the device */
    status = IoCreateSymbolicLink(&SymbolicLinkNameW, &DeviceNameW);
    if (!NT_SUCCESS(status))
        return status;

    return STATUS_SUCCESS;
}

// vim: sw=4 et