#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/usb/storage.h>
#include <linux/usb.h>

char *buffer;

static struct usb_endpoint_descriptor *find_bulk_in_endpoint(struct usb_device *udev) {
    struct usb_host_interface *iface_desc;
    struct usb_interface *iface;
    struct usb_endpoint_descriptor *endpoint;
    int i, j;

    for (i = 0; i < udev->actconfig->desc.bNumInterfaces; i++) {
        iface = udev->actconfig->interface[i];
        iface_desc = iface->cur_altsetting;
        if (iface_desc->desc.bInterfaceClass == 0x08) {
            pr_info("Found storage interface:\n");
            pr_info("  Interface Number: %d\n", iface_desc->desc.bInterfaceNumber);
            pr_info("  Number of Endpoints: %d\n", iface_desc->desc.bNumEndpoints);
            for (j = 0; j < iface_desc->desc.bNumEndpoints; j++) {
                endpoint = &iface_desc->endpoint[j].desc;
                if ((endpoint->bmAttributes & USB_ENDPOINT_XFERTYPE_MASK) == USB_ENDPOINT_XFER_BULK &&
                    (endpoint->bEndpointAddress & USB_DIR_IN))
                    return endpoint;
            }
        }
    }
    pr_err("No bulk IN endpoint found\n");
    return NULL;
}

int read_first_sector(struct usb_device *udev) {
    struct usb_endpoint_descriptor *bulk_in_endpoint = find_bulk_in_endpoint(udev);

    if (!bulk_in_endpoint)
        return -ENODEV;

    u8 buf[512];
    int ret, actual_size;
    unsigned char scsi_cmd[10] = { 
        0x28,
        0x00,
        0x00,
        0x00,
        0x00,
        0x01,
        0x00,
        0x00,
        0x00,
        0x00
    };

    ret = usb_control_msg(udev, usb_sndctrlpipe(udev, 0), 0x28, USB_DIR_IN | USB_TYPE_CLASS | USB_RECIP_INTERFACE, 0, 0,(void *)scsi_cmd, sizeof(scsi_cmd), 1000);
    if (ret < 0) {
        printk(KERN_ERR "USB SCSI READ command failed: %d\n", ret);
        return ret;
    }
    ret = usb_bulk_msg(udev, usb_rcvbulkpipe(udev, bulk_in_endpoint->bEndpointAddress), buf, 512, &actual_size, 1000);
    if (ret < 0) {
        printk(KERN_ERR "USB bulk read failed: %d\n", ret);
        return ret;
    }
    printk(KERN_INFO "First sector data:\n");
    for (int i = 0; i < 512; i++) {
        printk(KERN_CONT "%02x ", buf[i]);
        if ((i + 1) % 16 == 0)
            printk(KERN_CONT "\n");
    }
    return 0;
}

static int usb_notify(struct notifier_block *self, unsigned long action, void *dev) {
    struct usb_device *usb_dev = dev;

    switch (action) {
        case USB_DEVICE_ADD:
            printk(KERN_INFO "USB device connected: Vendor ID=%04x, Product ID=%04x\n",
                      usb_dev->descriptor.idVendor, usb_dev->descriptor.idProduct);
            printk(KERN_INFO "Path: %s\n", dev_name(&usb_dev->dev));
            read_first_sector(usb_dev);
            break;
        case USB_DEVICE_REMOVE:
            printk(KERN_INFO "USB device disconnected: Vendor ID=%04x, Product ID=%04x\n",
                      usb_dev->descriptor.idVendor, usb_dev->descriptor.idProduct);
            break;
        default:
            break;
    }
    return NOTIFY_OK;
}

static struct notifier_block usb_notifier = {
    .notifier_call = usb_notify
};

static int __init usb_device_list_init(void) {
    printk(KERN_INFO "Listing all connected USB devices:\n");
    usb_register_notify(&usb_notifier); 
    return 0;
}

static void __exit usb_device_list_exit(void) {
    usb_unregister_notify(&usb_notifier);
    printk(KERN_INFO "USB device listing module exiting.\n");
    kfree(buffer); 
}

late_initcall(usb_device_list_init);
module_exit(usb_device_list_exit);
MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("USB Device Module");
MODULE_AUTHOR("MyName");

echo -n "USBSECRET" | sudo dd of=/dev/sdb bs=1 count=10 conv=notrunc