/*
 * usbhid_remap.c
 * 
 * 功能：捕获 USB HID 键盘原始数据，重映射后通过 uinput 注入系统
 * 编译：gcc -o usbhid_remap usbhid_remap.c -lusb-1.0 -Wall -O2
 * 运行：sudo ./usbhid_remap
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/input.h>
#include <linux/uinput.h>
#include <libusb-1.0/libusb.h>

#define VENDOR_ID   0x258a   // SINO WEALTH
#define PRODUCT_ID  0x002a   // Keyboard
#define INTERFACE   0        // 标准键盘接口
#define ENDPOINT    0x81     // 输入端点

#define HID_REPORT_SIZE 8
#define MAX_KEYS 6

// 运行标志
volatile int running = 1;
int uinput_fd = -1;

// 当前按键状态（用于检测变化）
uint8_t prev_keys[MAX_KEYS] = {0};
uint8_t curr_keys[MAX_KEYS] = {0};

// 修饰键状态
uint8_t prev_mods = 0;
uint8_t curr_mods = 0;

// 信号处理
void signal_handler(int sig) {
    running = 0;
}


// USB HID Usage ID 到 Linux 键码的标准映射
static const int hid_to_linux[256] = {
    [0x00]=0,       [0x01]=KEY_F1,   [0x02]=KEY_F2, [0x03]=KEY_F3, [0x04]=KEY_F4, [0x05]=KEY_F5, [0x06]=KEY_F6, [0x07]=KEY_F7, [0x08]=KEY_F8, [0x09]=KEY_F9,   [0x0A]=KEY_F10, [0x0B]=KEY_F11, [0x0C]=KEY_F12,   [0x0D]=KEY_SYSRQ,   [0x0E]=KEY_SCROLLLOCK,   [0x0F]=KEY_ESC,
    [0x10]=KEY_GRAVE,   [0x11]=KEY_1,   [0x12]=KEY_2, [0x13]=KEY_3, [0x14]=KEY_4, [0x15]=KEY_5, [0x16]=KEY_6, [0x17]=KEY_7, [0x18]=KEY_8, [0x19]=KEY_9,   [0x1A]=KEY_0, [0x1B]=KEY_MINUS, [0x1C]=KEY_EQUAL,   [0x1D]=KEY_BACKSPACE,   [0x1E]=KEY_INSERT,   [0x1F]=KEY_HOME,
    [0x20]=KEY_TAB,   [0x21]=KEY_Q,   [0x22]=KEY_W, [0x23]=KEY_E, [0x24]=KEY_R, [0x25]=KEY_T, [0x26]=KEY_Y, [0x27]=KEY_U, [0x28]=KEY_I, [0x29]=KEY_O,   [0x2A]=KEY_P, [0x2B]=KEY_LEFTBRACE, [0x2C]=KEY_RIGHTBRACE,   [0x2D]=KEY_BACKSLASH,   [0x2E]=KEY_DELETE,   [0x2F]=KEY_END,
    [0x30]=KEY_CAPSLOCK,   [0x31]=KEY_A,   [0x32]=KEY_S, [0x33]=KEY_D, [0x34]=KEY_F, [0x35]=KEY_G, [0x36]=KEY_H, [0x37]=KEY_J, [0x38]=KEY_K, [0x39]=KEY_L,   [0x3A]=KEY_SEMICOLON, [0x3B]=KEY_APOSTROPHE, [0x3C]=KEY_ENTER,   [0x3D]=KEY_4,   [0x3E]=KEY_5,   [0x3F]=KEY_6,
    [0x40]=KEY_LEFTSHIFT,   [0x41]=KEY_Z,   [0x42]=KEY_X, [0x43]=KEY_C, [0x44]=KEY_V, [0x45]=KEY_B, [0x46]=KEY_N, [0x47]=KEY_M, [0x48]=KEY_COMMA, [0x49]=KEY_DOT,   [0x4A]=KEY_SLASH, [0x4B]=KEY_RIGHTSHIFT, [0x4C]=KEY_UP,   [0x4D]=KEY_1,   [0x4E]=KEY_2,   [0x4F]=KEY_3,
    [0x50]=KEY_LEFTCTRL,   [0x51]=KEY_LEFTMETA,   [0x52]=KEY_LEFTALT, [0x53]=KEY_SPACE, [0x54]=KEY_RIGHTALT, [0x55]=KEY_RIGHTMETA, [0x56]=KEY_COMPOSE, [0x57]=KEY_RIGHTCTRL, [0x58]=KEY_LEFT, [0x59]=KEY_DOWN,   [0x5A]=KEY_RIGHT, [0x5B]=KEY_0, [0x5C]=KEY_DOT,   [0x5D]=KEY_ENTER,   [0x5E]=KEY_E,   [0x5F]=KEY_F,
    [0x60]=KEY_0,   [0x61]=KEY_1,   [0x62]=KEY_2, [0x63]=KEY_3, [0x64]=KEY_4, [0x65]=KEY_5, [0x66]=KEY_6, [0x67]=KEY_7, [0x68]=KEY_8, [0x69]=KEY_9,   [0x6A]=KEY_A, [0x6B]=KEY_B, [0x6C]=KEY_C,   [0x6D]=KEY_D,   [0x6E]=KEY_E,   [0x6F]=KEY_F,
    [0x70]=KEY_0,   [0x71]=KEY_1,   [0x72]=KEY_2, [0x73]=KEY_3, [0x74]=KEY_4, [0x75]=KEY_5, [0x76]=KEY_6, [0x77]=KEY_7, [0x78]=KEY_8, [0x79]=KEY_9,   [0x7A]=KEY_A, [0x7B]=KEY_B, [0x7C]=KEY_C,   [0x7D]=KEY_D,   [0x7E]=KEY_E,   [0x7F]=KEY_F,
};

#if 0
// USB HID Usage ID 到 Linux 键码的标准映射
static const int hid_to_linux[256] = {
    [0x00]=0,   [0x01] = KEY_ESC,   [0x02] = KEY_1, [0x03] = KEY_2, [0x04] = KEY_3, [0x05] = KEY_4,
    [0x06] = KEY_5, [0x07] = KEY_6, [0x08] = KEY_7, [0x09] = KEY_8,
    [0x0A] = KEY_9, [0x0B] = KEY_0,
    [0x0C] = KEY_MINUS, [0x0D] = KEY_EQUAL, [0x0E] = KEY_BACKSPACE,
    [0x0F] = KEY_TAB, [0x10] = KEY_Q, [0x11] = KEY_W, [0x12] = KEY_E,
    [0x13] = KEY_R, [0x14] = KEY_T, [0x15] = KEY_Y, [0x16] = KEY_U,
    [0x17] = KEY_I, [0x18] = KEY_O, [0x19] = KEY_P,
    [0x1A] = KEY_LEFTBRACE, [0x1B] = KEY_RIGHTBRACE, [0x1C] = KEY_ENTER,
    [0x1D] = KEY_LEFTCTRL, [0x1E] = KEY_A, [0x1F] = KEY_S,
    [0x20] = KEY_D, [0x21] = KEY_F, [0x22] = KEY_G, [0x23] = KEY_H,
    [0x24] = KEY_J, [0x25] = KEY_K, [0x26] = KEY_L,
    [0x27] = KEY_SEMICOLON, [0x28] = KEY_APOSTROPHE, [0x29] = KEY_GRAVE,
    [0x2A] = KEY_LEFTSHIFT, [0x2B] = KEY_BACKSLASH,
    [0x2C] = KEY_Z, [0x2D] = KEY_X, [0x2E] = KEY_C, [0x2F] = KEY_V,
    [0x30] = KEY_B, [0x31] = KEY_N, [0x32] = KEY_M,
    [0x33] = KEY_COMMA, [0x34] = KEY_DOT, [0x35] = KEY_SLASH,
    [0x36] = KEY_RIGHTSHIFT, [0x37] = KEY_KPASTERISK,
    [0x38] = KEY_LEFTALT, [0x39] = KEY_SPACE, [0x3A] = KEY_CAPSLOCK,
    [0x3B] = KEY_F1, [0x3C] = KEY_F2, [0x3D] = KEY_F3, [0x3E] = KEY_F4,
    [0x3F] = KEY_F5, [0x40] = KEY_F6, [0x41] = KEY_F7, [0x42] = KEY_F8,
    [0x43] = KEY_F9, [0x44] = KEY_F10, [0x45] = KEY_F11, [0x46] = KEY_F12,
    [0x47] = KEY_SYSRQ, [0x48] = KEY_SCROLLLOCK, [0x49] = KEY_PAUSE,
    [0x4A] = KEY_INSERT, [0x4B] = KEY_HOME, [0x4C] = KEY_PAGEUP,
    [0x4D] = KEY_DELETE, [0x4E] = KEY_END, [0x4F] = KEY_PAGEDOWN,
    [0x50] = KEY_RIGHT, [0x51] = KEY_LEFT, [0x52] = KEY_DOWN, [0x53] = KEY_UP,
    [0x54] = KEY_NUMLOCK, [0x55] = KEY_KPSLASH, [0x56] = KEY_KPASTERISK,
    [0x57] = KEY_KPMINUS, [0x58] = KEY_KPPLUS, [0x59] = KEY_KPENTER,
    [0x5A] = KEY_KP1, [0x5B] = KEY_KP2, [0x5C] = KEY_KP3, [0x5D] = KEY_KP4,
    [0x5E] = KEY_KP5, [0x5F] = KEY_KP6, [0x60] = KEY_KP7, [0x61] = KEY_KP8,
    [0x62] = KEY_KP9, [0x63] = KEY_KP0, [0x64] = KEY_KPDOT,
    [0x65] = KEY_102ND,[0x66] = KEY_102ND,[0x67] = KEY_102ND,[0x68] = KEY_102ND,
    [0x69] = KEY_102ND,[0x6a] = KEY_102ND,[0x6b] = KEY_102ND,[0x6c] = KEY_102ND,
    [0x6d] = KEY_102ND,[0x6e] = KEY_102ND,[0x6f] = KEY_102ND,
    [0x70] = KEY_102ND,[0x71] = KEY_102ND,[0x72] = KEY_102ND,[0x73] = KEY_102ND,
    [0x74] = KEY_102ND,[0x75] = KEY_102ND,[0x76] = KEY_102ND,[0x77] = KEY_102ND,
    [0x78] = KEY_102ND,[0x79] = KEY_102ND,[0x7a] = KEY_102ND,[0x7b] = KEY_102ND,
    [0x7c] = KEY_102ND,[0x7d] = KEY_102ND,[0x7e] = KEY_102ND,[0x7f] = KEY_102ND,
    [0xE0] = KEY_LEFTCTRL, [0xE1] = KEY_LEFTSHIFT, [0xE2] = KEY_LEFTALT,
    [0xE3] = KEY_LEFTMETA, [0xE4] = KEY_RIGHTCTRL, [0xE5] = KEY_RIGHTSHIFT,
    [0xE6] = KEY_RIGHTALT, [0xE7] = KEY_RIGHTMETA,
};
#endif
// 自定义重映射表（覆盖标准映射）
// 修改这里来改变按键映射
static const int custom_remap[256] = {
    // 示例：将 0x31 和 0x32 重新映射
    // [0x31] = KEY_RIGHTBRACE,  // 取消注释并修改为你需要的键
    // [0x32] = KEY_BACKSLASH,   // 取消注释并修改为你需要的键
    
    // 默认：0 表示使用标准映射
    [0 ... 255] = 0,
};

// 修饰键位定义
#define MOD_LCTRL   0x01
#define MOD_LSHIFT  0x02
#define MOD_LALT    0x04
#define MOD_LMETA   0x08
#define MOD_RCTRL   0x10
#define MOD_RSHIFT  0x20
#define MOD_RALT    0x40
#define MOD_RMETA   0x80

// 获取映射后的键码
int get_mapped_key(int hid_code) {
    if (hid_code < 0 || hid_code > 255) return 0;
    if (custom_remap[hid_code] != 0) return custom_remap[hid_code];
    return hid_to_linux[hid_code];
}

// 初始化 uinput 设备
int init_uinput(void) {
    int fd = open("/dev/uinput", O_WRONLY | O_NONBLOCK);
    if (fd < 0) {
        perror("Failed to open /dev/uinput");
        return -1;
    }
    
    struct uinput_setup setup;
    memset(&setup, 0, sizeof(setup));
    
    // 设置设备信息
    snprintf(setup.name, UINPUT_MAX_NAME_SIZE, "USB HID Remapped Keyboard");
    setup.id.bustype = BUS_USB;
    setup.id.vendor = VENDOR_ID;
    setup.id.product = PRODUCT_ID;
    setup.id.version = 1;
    
    // 启用事件类型
    if (ioctl(fd, UI_SET_EVBIT, EV_KEY) < 0) goto error;
    if (ioctl(fd, UI_SET_EVBIT, EV_SYN) < 0) goto error;
    if (ioctl(fd, UI_SET_EVBIT, EV_MSC) < 0) goto error;
    if (ioctl(fd, UI_SET_MSCBIT, MSC_SCAN) < 0) goto error;
    
    // 启用所有键
    for (int i = 0; i < 256; i++) {
        if (ioctl(fd, UI_SET_KEYBIT, i) < 0) goto error;
    }
    
    // 创建设备
    if (ioctl(fd, UI_DEV_SETUP, &setup) < 0) goto error;
    if (ioctl(fd, UI_DEV_CREATE) < 0) goto error;
    
    printf("uinput device created: %s\n", setup.name);
    return fd;
    
error:
    perror("uinput setup failed");
    close(fd);
    return -1;
}

// 销毁 uinput 设备
void destroy_uinput(int fd) {
    if (fd < 0) return;
    ioctl(fd, UI_DEV_DESTROY);
    close(fd);
    printf("uinput device destroyed\n");
}

// 发送按键事件
void send_key_event(int fd, int code, int value, int scan_code) {
    struct input_event ev;
    
    // MSC_SCAN 事件（原始扫描码）
    if (scan_code >= 0) {
        memset(&ev, 0, sizeof(ev));
        gettimeofday(&ev.time, NULL);
        ev.type = EV_MSC;
        ev.code = MSC_SCAN;
        ev.value = scan_code;
        write(fd, &ev, sizeof(ev));
    }
    
    // KEY 事件
    memset(&ev, 0, sizeof(ev));
    gettimeofday(&ev.time, NULL);
    ev.type = EV_KEY;
    ev.code = code;
    ev.value = value;
    write(fd, &ev, sizeof(ev));
    
    // SYN 事件
    memset(&ev, 0, sizeof(ev));
    gettimeofday(&ev.time, NULL);
    ev.type = EV_SYN;
    ev.code = SYN_REPORT;
    ev.value = 0;
    write(fd, &ev, sizeof(ev));
}

// 发送修饰键事件
void send_modifier_events(int fd, uint8_t prev, uint8_t curr) {
    struct {
        uint8_t mask;
        int keycode;
    } mods[] = {
        {MOD_LCTRL, KEY_LEFTCTRL}, {MOD_LSHIFT, KEY_LEFTSHIFT},
        {MOD_LALT, KEY_LEFTALT}, {MOD_LMETA, KEY_LEFTMETA},
        {MOD_RCTRL, KEY_RIGHTCTRL}, {MOD_RSHIFT, KEY_RIGHTSHIFT},
        {MOD_RALT, KEY_RIGHTALT}, {MOD_RMETA, KEY_RIGHTMETA},
    };
    
    for (int i = 0; i < 8; i++) {
        uint8_t mask = mods[i].mask;
        int prev_state = (prev & mask) ? 1 : 0;
        int curr_state = (curr & mask) ? 1 : 0;
        
        if (prev_state != curr_state) {
            printf("  Modifier %s: %s\n", 
                   mods[i].keycode == KEY_LEFTCTRL ? "LCtrl" :
                   mods[i].keycode == KEY_LEFTSHIFT ? "LShift" :
                   mods[i].keycode == KEY_LEFTALT ? "LAlt" :
                   mods[i].keycode == KEY_LEFTMETA ? "LMeta" :
                   mods[i].keycode == KEY_RIGHTCTRL ? "RCtrl" :
                   mods[i].keycode == KEY_RIGHTSHIFT ? "RShift" :
                   mods[i].keycode == KEY_RIGHTALT ? "RAlt" : "RMeta",
                   curr_state ? "DOWN" : "UP");
            send_key_event(fd, mods[i].keycode, curr_state, -1);
        }
    }
}

// 处理 HID 报告
void process_hid_report(int fd, const uint8_t *data, int len) {
    if (len < 8) return;
    
    curr_mods = data[0];
    
    // 提取按键（去重并排序）
    int key_count = 0;
    memset(curr_keys, 0, sizeof(curr_keys));
    for (int i = 2; i < 8 && key_count < MAX_KEYS; i++) {
        if (data[i] != 0) {
            // 检查是否已存在
            int exists = 0;
            for (int j = 0; j < key_count; j++) {
                if (curr_keys[j] == data[i]) { exists = 1; break; }
            }
            if (!exists) curr_keys[key_count++] = data[i];
        }
    }
    
    // 打印原始数据
    printf("Raw HID: ");
    for (int i = 0; i < len; i++) printf("%02X ", data[i]);
    printf("\n");
    
    // 处理修饰键变化
    if (prev_mods != curr_mods) {
        send_modifier_events(fd, prev_mods, curr_mods);
    }
    
    // 找出释放的键（在 prev 中但不在 curr 中）
    for (int i = 0; i < MAX_KEYS; i++) {
        if (prev_keys[i] == 0) continue;
        int found = 0;
        for (int j = 0; j < MAX_KEYS; j++) {
            if (curr_keys[j] == prev_keys[i]) { found = 1; break; }
        }
        if (!found) {
            int mapped = get_mapped_key(prev_keys[i]);
            if (mapped > 0) {
                printf("  Key UP: HID_%02X -> Linux_%d (%s)\n", 
                       prev_keys[i], mapped, 
                       mapped == KEY_A ? "A" : mapped == KEY_B ? "B" :
                       mapped == KEY_ENTER ? "Enter" : mapped == KEY_SPACE ? "Space" :
                       mapped == KEY_BACKSLASH ? "Backslash" :
                       mapped == KEY_RIGHTBRACE ? "RightBrace" : "Other");
                send_key_event(fd, mapped, 0, 0x70000 | prev_keys[i]);
            }
        }
    }
    
    // 找出按下的键（在 curr 中但不在 prev 中）
    for (int i = 0; i < MAX_KEYS; i++) {
        if (curr_keys[i] == 0) continue;
        int found = 0;
        for (int j = 0; j < MAX_KEYS; j++) {
            if (prev_keys[j] == curr_keys[i]) { found = 1; break; }
        }
        if (!found) {
            int mapped = get_mapped_key(curr_keys[i]);
            if (mapped > 0) {
                const char *remap_note = (custom_remap[curr_keys[i]] != 0) ? " [REMAPPED]" : "";
                printf("  Key DOWN: HID_%02X -> Linux_%d (%s)%s\n", 
                       curr_keys[i], mapped,
                       mapped == KEY_A ? "A" : mapped == KEY_B ? "B" :
                       mapped == KEY_ENTER ? "Enter" : mapped == KEY_SPACE ? "Space" :
                       mapped == KEY_BACKSLASH ? "Backslash" :
                       mapped == KEY_RIGHTBRACE ? "RightBrace" : "Other",
                       remap_note);
                send_key_event(fd, mapped, 1, 0x70000 | curr_keys[i]);
            } else {
                printf("  Key DOWN: HID_%02X -> (unmapped, ignored)\n", curr_keys[i]);
            }
        }
    }
    
    // 保存当前状态
    prev_mods = curr_mods;
    memcpy(prev_keys, curr_keys, sizeof(curr_keys));
}

// 打印用法
void print_usage(const char *prog) {
    printf("Usage: %s [options]\n", prog);
    printf("Options:\n");
    printf("  -v VID    Set vendor ID (hex, default: %04X)\n", VENDOR_ID);
    printf("  -p PID    Set product ID (hex, default: %04X)\n", PRODUCT_ID);
    printf("  -i IFACE  Set interface (default: %d)\n", INTERFACE);
    printf("  -d        Detach kernel driver only, don't capture\n");
    printf("  -h        Show this help\n");
    printf("\nExample:\n");
    printf("  sudo %s -v 258a -p 002a\n", prog);
}

int main(int argc, char *argv[]) {
    libusb_context *ctx = NULL;
    libusb_device_handle *handle = NULL;
    int ret;
    int opt;
    uint16_t vid = VENDOR_ID;
    uint16_t pid = PRODUCT_ID;
    int iface = INTERFACE;
    int detach_only = 0;
    
    // 解析参数
    while ((opt = getopt(argc, argv, "v:p:i:dh")) != -1) {
        switch (opt) {
            case 'v': vid = strtol(optarg, NULL, 16); break;
            case 'p': pid = strtol(optarg, NULL, 16); break;
            case 'i': iface = atoi(optarg); break;
            case 'd': detach_only = 1; break;
            case 'h': print_usage(argv[0]); return 0;
            default: print_usage(argv[0]); return 1;
        }
    }
    
    printf("=== USB HID Remapper ===\n");
    printf("Target: VID=%04X PID=%04X Interface=%d\n", vid, pid, iface);
    
    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    // 初始化 libusb
    ret = libusb_init(&ctx);
    if (ret < 0) {
        fprintf(stderr, "libusb_init failed: %s\n", libusb_error_name(ret));
        return 1;
    }
    
    // 打开设备
    handle = libusb_open_device_with_vid_pid(ctx, vid, pid);
    if (!handle) {
        fprintf(stderr, "Failed to open device %04X:%04X\n", vid, pid);
        fprintf(stderr, "Check device connection and permissions (try sudo)\n");
        libusb_exit(ctx);
        return 1;
    }
    
    // 检查并分离内核驱动
    if (libusb_kernel_driver_active(handle, iface) == 1) {
        printf("Detaching kernel driver...\n");
        ret = libusb_detach_kernel_driver(handle, iface);
        if (ret < 0) {
            fprintf(stderr, "Failed to detach driver: %s\n", libusb_error_name(ret));
            libusb_close(handle);
            libusb_exit(ctx);
            return 1;
        }
    }
    
    if (detach_only) {
        printf("Kernel driver detached. Press Enter to reattach and exit...");
        getchar();
        libusb_attach_kernel_driver(handle, iface);
        libusb_close(handle);
        libusb_exit(ctx);
        return 0;
    }
    
    // 声明接口
    ret = libusb_claim_interface(handle, iface);
    if (ret < 0) {
        fprintf(stderr, "Failed to claim interface: %s\n", libusb_error_name(ret));
        libusb_attach_kernel_driver(handle, iface);
        libusb_close(handle);
        libusb_exit(ctx);
        return 1;
    }
    
    // 初始化 uinput
    uinput_fd = init_uinput();
    if (uinput_fd < 0) {
        libusb_release_interface(handle, iface);
        libusb_attach_kernel_driver(handle, iface);
        libusb_close(handle);
        libusb_exit(ctx);
        return 1;
    }
    
    printf("Capturing... Press Ctrl+C to exit\n\n");
    
    // 主循环
    uint8_t data[HID_REPORT_SIZE];
    int transferred;
    
    while (running) {
        ret = libusb_interrupt_transfer(handle, ENDPOINT | LIBUSB_ENDPOINT_IN,
                                        data, sizeof(data), &transferred, 100);
        
        if (ret == LIBUSB_ERROR_TIMEOUT) {
            continue; // 正常超时，检查 running 标志
        } else if (ret == LIBUSB_ERROR_NO_DEVICE) {
            printf("Device disconnected\n");
            break;
        } else if (ret < 0) {
            if (running) fprintf(stderr, "Transfer error: %s\n", libusb_error_name(ret));
            usleep(10000);
            continue;
        }
        
        if (transferred > 0) {
            process_hid_report(uinput_fd, data, transferred);
            fflush(stdout);
        }
    }
    
    printf("\nShutting down...\n");
    
    // 清理
    destroy_uinput(uinput_fd);
    libusb_release_interface(handle, iface);
    
    ret = libusb_attach_kernel_driver(handle, iface);
    if (ret < 0) {
        printf("Warning: Could not reattach kernel driver: %s\n", libusb_error_name(ret));
    } else {
        printf("Kernel driver reattached\n");
    }
    
    libusb_close(handle);
    libusb_exit(ctx);
    
    return 0;
}