#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/kprobes.h>
#include <linux/sched.h>    // for current->comm, current->pid
#include <linux/uaccess.h>  // for copy_from_user, copy_to_user (如果需要访问用户空间内存)
#include <linux/usb.h>      // 包含 USB 核心数据结构，如 struct urb

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Leo");
MODULE_DESCRIPTION("Hook usb_submit_urb ioctl on arm64"); // 这里的注释可能不完全准确，usb_submit_urb不是ioctl，是内核函数

static struct kprobe kp;

// 针对 ARM64 (AArch64) 架构，函数参数通常在 x0, x1, x2... 寄存器中。
// pt_regs 中的 regs 数组索引通常对应这些寄存器。
// 对于 AArch64，x0 是 regs[0]，x1 是 regs[1]，以此类推。
// usb_submit_urb 函数原型通常是：int usb_submit_urb(struct urb *urb, gfp_t mem_flags)
// 所以第一个参数 struct urb *urb 应该在 x0 寄存器中，对应 regs[0]。

static int handler_pre(struct kprobe *p, struct pt_regs *regs)
{
    // !!! 核心修改点：在 ARM64 上，函数第一个参数通常在 x0 (regs[0]) !!!
    struct urb *urb_kern = (struct urb *)regs->regs[0];

    // 严谨的空指针检查
    if (!urb_kern) {
        pr_err("[usbFilter] handler_pre: urb_kern is NULL. Skipping.\n");
        return 0; // KPROBE_OK, 让原始函数继续执行，避免崩溃
    }

    // 打印进程信息、PID
    pr_info("[usbFilter] process: %s, pid: %d\n",
             current->comm, current->pid);

    // 获取端点地址、传输长度和 pipe 信息
    pr_info("[usbFilter] urb_kern: %p, pipe: 0x%08x, ep: 0x%x, len: %d\n",
             urb_kern,
             urb_kern->pipe, // 新增打印 pipe 信息
             urb_kern->ep ? urb_kern->ep->desc.bEndpointAddress : 0, // 检查 urb_kern->ep
             urb_kern->transfer_buffer_length);


    if (urb_kern->transfer_buffer && urb_kern->transfer_buffer_length > 0) {
        unsigned char data[16] = {0}; // 局部缓冲区，用于存放拷贝的数据
        unsigned int to_copy = min((unsigned int)16, urb_kern->transfer_buffer_length);
        bool data_copied_successfully = false;

        // 尝试从用户空间拷贝
        if (copy_from_user(data, urb_kern->transfer_buffer, to_copy) == 0) { // 0 表示成功
            pr_info("[usbFilter] Successfully copied %u bytes using copy_from_user from user buffer at %p.\n", to_copy, urb_kern->transfer_buffer);
            data_copied_successfully = true;
        } else {
            // copy_from_user 失败
            pr_warn("[usbFilter] copy_from_user failed for buffer at %p. Attempting memcpy (assuming buffer is in kernel space).\n", urb_kern->transfer_buffer);
            
            // 警告：如果 transfer_buffer 不是有效的内核地址，memcpy 可能会导致内核崩溃。
            // 这仅作为调试时的后备尝试。
            memcpy(data, urb_kern->transfer_buffer, to_copy);
            // 如果 memcpy 没有导致崩溃，我们假设数据为了打印目的是成功拷贝的。
            pr_info("[usbFilter] memcpy attempted for %u bytes from buffer at %p (assumed kernel space).\n", to_copy, urb_kern->transfer_buffer);
            data_copied_successfully = true; // 标记为成功，以便后续打印
        }

        if (data_copied_successfully) {
            char hex[3 * 16 + 1] = {0}; // 用于存放十六进制字符串
            int i;
            for (i = 0; i < to_copy; ++i) {
                // 确保 snprintf 不会溢出 hex 缓冲区
                snprintf(hex + i * 3, sizeof(hex) - (i * 3), "%02X ", data[i]);
            }
            pr_info("[usbFilter] first %u bytes (hex): %s\n", to_copy, hex);
        }
    }

    
    // 这里一定要设置返回成功，如果不设置返回成功，将会导致一直重复发送；
    // regs->regs[0] = 0; // 设置 x0 寄存器（返回值）为 0 (成功)
    //这里return 1表示阻断，return 0 表示继续运行
    // return 1;
    return 0;
}

static int __init usb_hook_init(void)
{
    // usb_submit_urb 是一个导出的符号，Kprobe 可以 Hook
    kp.symbol_name = "usb_submit_urb";
    kp.pre_handler = handler_pre;

    if (register_kprobe(&kp) < 0) {
        pr_err("[usbFilter] register_kprobe failed\n");
        return -1;
    }
    pr_info("[usbFilter] kprobe registered for %s at %p\n", kp.symbol_name, kp.addr);
    return 0;
}

static void __exit usb_hook_exit(void)
{
    unregister_kprobe(&kp);
    pr_info("[usbFilter] kprobe unregistered\n");
}

module_init(usb_hook_init);
module_exit(usb_hook_exit);