#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/usbdevice_fs.h>
#include <linux/kprobes.h>
#include <linux/sched.h>
#include <linux/uaccess.h> // for copy_from_user

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Leo");
MODULE_DESCRIPTION("Hook USBDEVFS_SUBMITURB ioctl on arm64");

static struct kprobe kp;

static int handler_pre(struct kprobe *p, struct pt_regs *regs)
{
    unsigned int cmd = 0;
    struct usbdevfs_urb urb_kern;
    struct usbdevfs_urb __user *urb_user;

    cmd = (unsigned int)regs->regs[1]; // ioctl 的第二个参数

    if (cmd == USBDEVFS_SUBMITURB) {
        urb_user = (struct usbdevfs_urb __user *)regs->regs[2]; // 第三个参数
        if (urb_user) {
            if (copy_from_user(&urb_kern, urb_user, sizeof(struct usbdevfs_urb)) == 0) {
                pr_info("[usbFilter] process: %s, pid: %d, ep: 0x%x, len: %d\n",
                        current->comm, current->pid,
                        urb_kern.endpoint, urb_kern.buffer_length);

                if (urb_kern.buffer && urb_kern.buffer_length > 0) {
                    unsigned char data[16] = {0};
                    unsigned int to_copy = urb_kern.buffer_length > 16 ? 16 : urb_kern.buffer_length;
                    if (copy_from_user(data, urb_kern.buffer, to_copy) == 0) {
                        char hex[3 * 16 + 1] = {0};
                        int i;
                        for (i = 0; i < to_copy; ++i)
                            snprintf(hex + i * 3, 4, "%02X ", data[i]);
                        pr_info("[usbFilter] first %u bytes (hex): %s\n", to_copy, hex);
                    } else {
                        pr_warn("[usbFilter] copy_from_user buffer failed\n");
                    }
                }
            } else {
                pr_warn("[usbFilter] copy_from_user failed\n");
            }
        }
    }
    return 0;
}

static int __init usb_hook_init(void)
{
    kp.symbol_name = "usbdev_do_ioctl";
    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\n", kp.symbol_name);
    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);