#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/kprobes.h>
#include <linux/sched.h>
#include <linux/uaccess.h>
#include <linux/usb.h>
#include <linux/slab.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Leo");
MODULE_DESCRIPTION("Hook usb_submit_urb() on ARM64 and replace callback");

static struct kprobe kp;

// 自定义回调上下文，保存原始回调和上下文
struct urb_context
{
    usb_complete_t original_complete;
    void *original_context;
};

// 包装回调函数
static void callback_wrapper(struct urb *urb)
{
    struct urb_context *ctx = urb->context;

    pr_info("[usbFilter] [callback_wrapper] URB 完成: endpoint=0x%x, status=%d, actual_length=%d\n",
            usb_pipeendpoint(urb->pipe),
            urb->status,
            urb->actual_length);

    if (urb->transfer_buffer && urb->actual_length > 0)
    {
        char hex[3 * 32 + 1] = {0};
        int i, len = min(32, urb->actual_length);
        unsigned char *data = (unsigned char *)urb->transfer_buffer;
        for (i = 0; i < len; ++i)
        {
            snprintf(hex + i * 3, sizeof(hex) - i * 3, "%02X ", data[i]);
        }
        pr_info("[usbFilter] [callback_wrapper] 返回数据(hex): %s\n", hex);
    }

    if (ctx && ctx->original_complete)
    {
        urb->context = ctx->original_context;
        ctx->original_complete(urb);
    }

    kfree(ctx); // 释放上下文
}

static int handler_pre(struct kprobe *p, struct pt_regs *regs)
{
    struct urb *urb_kern = (struct urb *)regs->regs[0];

    if (!urb_kern)
        return 0;

    pr_info("[usbFilter] 提交 URB 的进程: %s (pid: %d)\n", current->comm, current->pid);

    pr_info("[usbFilter] URB: %p, pipe=0x%x, flags=0x%x\n",
            urb_kern, urb_kern->pipe, urb_kern->transfer_flags);

    if (urb_kern->dev)
    {
        pr_info("[usbFilter] USB设备: VID=0x%04x, PID=0x%04x\n",
                urb_kern->dev->descriptor.idVendor,
                urb_kern->dev->descriptor.idProduct);
    }

    // 打印传输方向与类型
    pr_info("[usbFilter] pipe: 端点=%d, 方向=%s, 类型=%s\n",
            usb_pipeendpoint(urb_kern->pipe),
            usb_pipein(urb_kern->pipe) ? "IN" : "OUT",
            usb_pipetype(urb_kern->pipe) == PIPE_CONTROL ? "CONTROL" : usb_pipetype(urb_kern->pipe) == PIPE_ISOCHRONOUS ? "ISO"
                                                                   : usb_pipetype(urb_kern->pipe) == PIPE_BULK          ? "BULK"
                                                                   : usb_pipetype(urb_kern->pipe) == PIPE_INTERRUPT     ? "INTERRUPT"
                                                                                                                        : "UNKNOWN");

    if (usb_pipetype(urb_kern->pipe) == PIPE_CONTROL)
    {
        struct usb_ctrlrequest *setup = (struct usb_ctrlrequest *)urb_kern->setup_packet;
        if (setup)
        {
            pr_info("[usbFilter] 控制传输Setup包: bRequestType=0x%02x, bRequest=0x%02x, "
                    "wValue=0x%04x, wIndex=0x%04x, wLength=%u\n",
                    setup->bRequestType, setup->bRequest,
                    le16_to_cpu(setup->wValue), le16_to_cpu(setup->wIndex),
                    le16_to_cpu(setup->wLength));
        }

        // //将原有的setup包传输的数据全部设置为0
        // if (urb_kern->setup_packet && urb_kern->transfer_buffer_length > 0)
        // {
        //     memset(urb_kern->setup_packet, 0, urb_kern->transfer_buffer_length);
        //     pr_info("[usbFilter] 已将控制传输的 setup 包数据清零\n");
        // }
    }
    else
    {
        // 打印前 32 字节传输数据
        if (urb_kern->transfer_buffer && urb_kern->transfer_buffer_length > 0)
        {
            unsigned char data[32] = {0};
            unsigned int to_copy = min(32U, (unsigned int)urb_kern->transfer_buffer_length);
            memcpy(data, urb_kern->transfer_buffer, to_copy);

            char hex[3 * 32 + 1] = {0};
            int i;
            for (i = 0; i < to_copy; ++i)
            {
                snprintf(hex + i * 3, sizeof(hex) - i * 3, "%02X ", data[i]);
            }
            pr_info("[usbFilter] 数据内容(hex): %s\n", hex);
        }
    }

    // 是否为目标设备
    if (urb_kern->dev &&
        urb_kern->dev->descriptor.idVendor == 0x1a86 &&
        urb_kern->dev->descriptor.idProduct == 0x55de)
    {

        pr_info("[usbFilter] 命中目标设备，替换 URB 回调\n");

        struct urb_context *ctx = kmalloc(sizeof(*ctx), GFP_ATOMIC);
        if (!ctx)
        {
            pr_err("[usbFilter] 分配回调上下文失败\n");
            return 0;
        }

        ctx->original_complete = urb_kern->complete;
        ctx->original_context = urb_kern->context;

        urb_kern->complete = callback_wrapper;
        urb_kern->context = ctx;

        // 可以选择返回 0，让 URB 正常提交；也可以选择模拟成功阻断：
        // regs->regs[0] = 0;
        // return 1;
    }

    return 0;
}

static int __init usb_hook_init(void)
{
    kp.symbol_name = "usb_submit_urb";
    kp.pre_handler = handler_pre;

    if (register_kprobe(&kp) < 0)
    {
        pr_err("[usbFilter] 无法注册 kprobe\n");
        return -1;
    }

    pr_info("[usbFilter] 成功 hook usb_submit_urb()\n");
    return 0;
}

static void __exit usb_hook_exit(void)
{
    unregister_kprobe(&kp);
    pr_info("[usbFilter] 已卸载 usb_submit_urb hook\n");
}

module_init(usb_hook_init);
module_exit(usb_hook_exit);