#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>
#include <linux/version.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Leo");
MODULE_DESCRIPTION("Monitor usb_submit_urb() data submission only");

static struct kprobe kp;

// 兼容 x86_64 和 arm64 获取第一个参数
static struct urb *get_urb_from_regs(struct pt_regs *regs)
{
#if defined(CONFIG_ARM64)
    return (struct urb *)regs->regs[0];
#elif defined(CONFIG_X86_64)
    return (struct urb *)regs->di;
#else
#error "Unsupported architecture"
#endif
}

static int handler_pre(struct kprobe *p, struct pt_regs *regs)
{
    struct urb *urb_kern = get_urb_from_regs(regs);

    if (!urb_kern)
        return 0;

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

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

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

    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");

    // 打印控制传输的setup包内容
    if (usb_pipetype(urb_kern->pipe) == PIPE_CONTROL && urb_kern->setup_packet) {
        char setup_hex[3 * 8 + 1] = {0};
        int i;
        unsigned char *setup = (unsigned char *)urb_kern->setup_packet;
        for (i = 0; i < 8; ++i) {
            snprintf(setup_hex + i * 3, sizeof(setup_hex) - i * 3, "%02X ", setup[i]);
        }
        pr_info("[usbFilter] 控制传输setup包(8字节hex): %s\n", setup_hex);
    }

    // 打印控制传输的数据内容
    if (urb_kern->transfer_buffer && urb_kern->transfer_buffer_length > 0) {
        unsigned int to_copy = min(32U, (unsigned int)urb_kern->transfer_buffer_length);
        unsigned char data[32] = {0};
        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] 提交数据(前32字节hex): %s\n", hex);
    }

    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);