/*
 * 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"

#include "even_peb.h"
#include "utlist.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}};

even_peb_t *even_fake_peb = NULL;

static even_peb_t __stdcall *even_find_peb(HANDLE pid)
{
    DbgPrint("Even Htsysm7679: Finding PEB for PID %d\r\n", pid);
    even_peb_t *s;
    DL_SEARCH_SCALAR(even_fake_peb, s, pid, pid);

    DbgPrint("Even Htsysm7679: Found at %08x\r\n", s);
    return s;
}

static NTSTATUS __stdcall even_create_peb(HANDLE pid)
{
    DbgPrint("Even Htsysm7679: Creating PEB for PID %d\r\n", pid);

    even_peb_t *s = even_find_peb(pid);
    if (!s)
    {
        DbgPrint("Even Htsysm7679: Allocating new PEB\r\n");

        s = ExAllocatePool(PagedPool, sizeof(*s));
        if (!s)
            return STATUS_INTERNAL_ERROR;

        s->pid = pid;
        s->refcnt = 0;
        DL_APPEND(even_fake_peb, s);
    }

    DbgPrint("Even Htsysm7679: fake_peb at %08x\r\n", s);

    s->refcnt++;

    return STATUS_SUCCESS;
}

static NTSTATUS __stdcall even_delete_peb(even_peb_t *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");

        DL_DELETE(even_fake_peb, peb);
        ExFreePool(peb);

        return STATUS_SUCCESS;
    }

    return STATUS_SUCCESS;
}

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

    even_peb_t *el, *tmp;
    DL_FOREACH_SAFE(even_fake_peb, el, tmp)
    {
        DL_DELETE(even_fake_peb, el);
        ExFreePool(el);
    }

    /*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;
}

static NTSTATUS __stdcall even_dispatch_open_close(PDEVICE_OBJECT DeviceObject, 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);
    }
    break;

    case IRP_MJ_CLOSE: {
        if (pid == 0)
            goto err_out;

        DbgPrint("Even Htsysm7679: Handle to symbolink link closed by PID %d\r\n", pid);

        even_peb_t *s = even_find_peb(pid);
        if (!s)
            goto err_out;

        Irp->IoStatus.Information = 0;
        Irp->IoStatus.Status = even_delete_peb(s);
    }
    break;

    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: {
        DbgPrint("Even Htsysm7679: Check 1\r\n");
        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);
    }
    break;

    case IOCTL_CHECK_2: {
        DbgPrint("Even Htsysm7679: Check 2\r\n");
        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);
    }
    break;

    case IOCTL_CHECK_3: {
        DbgPrint("Even Htsysm7679: Check 3\r\n");
        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);
    }
    break;

    case IOCTL_WRITE_PEB_1: {
        DbgPrint("Even Htsysm7679: Write PEB 1\r\n");
        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_data)) ||
            (!even_out_data))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        even_peb_t *s = NULL;

        if ((even_in_data->offset + 4) > sizeof(s->fake_peb))
        {
            DbgPrint("Even Htsysm7679: offset %d out of bounds\r\n", even_in_data->offset);
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        };

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

        DbgPrint("Even Htsysm7679: Passed Value = %08x\r\n", even_in_data->value);

        unsigned int value = ~even_in_data->value;

        DbgPrint("Even Htsysm7679: Write Offset = %08x, Value = %08x\r\n", even_in_data->offset, value);
        *(int *)&s->fake_peb[even_in_data->offset] = value;

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

    case IOCTL_WRITE_PEB_2: {
        DbgPrint("Even Htsysm7679: Write PEB 2\r\n");
        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_data)) ||
            (!even_out_data))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        even_peb_t *s = NULL;

        if ((even_in_data->offset + 4) > sizeof(s->fake_peb))
        {
            DbgPrint("Even Htsysm7679: offset %d out of bounds\r\n", even_in_data->offset);
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        };

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

        DbgPrint("Even Htsysm7679: Write Offset = %08x, Value = %08x\r\n", even_in_data->offset, even_in_data->value);
        *(int *)&s->fake_peb[even_in_data->offset] = even_in_data->value;

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

    case IOCTL_READ_PEB: {
        DbgPrint("Even Htsysm7679: Read PEB\r\n");
        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(*even_in_data)) || (!even_in_data) ||
            (io_stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(*even_out_data)) || (!even_out_data))
        {
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        }

        even_peb_t *s = NULL;

        if ((even_in_data->offset + 4) > sizeof(s->fake_peb))
        {
            DbgPrint("Even Htsysm7679: offset out of bounds\r\n");
            Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            goto err_out;
        };

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

        DbgPrint("Even Htsysm7679: Read Offset = %08x, Mask = %08x\r\n", even_in_data->offset, even_in_data->value);

        unsigned int value = even_in_data->value;
        DbgPrint("Even Htsysm7679: Value = %08x\r\n", value);

        RtlCopyBytes(&even_out_data->data, &s->fake_peb[even_in_data->offset], sizeof(even_out_data->data));

        *(int *)even_out_data->data &= value;

        DbgPrint("Even Htsysm7679: Return = %08x\r\n", *(int *)even_out_data->data);

        even_out_data->result = 1;
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = sizeof(*even_out_data);
    }
    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_dispatch_open_close;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = even_dispatch_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