S社NB30機種向SHDIAG

日系机圈子里的，看到这个标题就应该明白这是干什么用的了吧？
说实话写的很烂，能凑合用的水平，也不知道那么多依赖的更新之后，还能不能用。
稍稍做了一些小修补，由于在国外也没有机器测试，还是希望有人可以来测一下。
这个小东西是很久以前的商用品了，现在也觉得没啥价值，故直接公开
想法及思路来自tewilove
依赖PyUSB

import sys
import usb.core
import usb.util
import random

NB30_USB_VID = 0x04dd
NB30_USB_PID = 0x933a


def shdiag_xchar(c):
    xft = [0, 2, 4, 6, 1, 3, 5, 7]
    r = 0
    for i in range(8):
        r |= (((c >> i) & 1) << xft[i])
    return r


def shdiag_encode_data(in_data):
    ip = 0
    op = 1
    out_data = [0] * (len(in_data) * 2)
    v = (random.randint(0, 0x7f) + 0xf)
    x = shdiag_xchar(v)
    out_data[0] = v

    for char in in_data:
        if char in '\r\n':
            out_data[op] = ord(char)
            op += 1
            if char == '\n':
                x = shdiag_xchar(v)
                out_data[op] = v
                op += 1
        else:
            out_data[op] = (ord(char) ^ (x & 0x7f)) + 0xf
            x = (x >> 1) | ((x & 1) << 7)
            op += 1

    return out_data[:op]


def shdiag_decode_data(in_data):
    ip = 1
    op = 0
    out_data = [0] * (len(in_data) - 1)
    x = shdiag_xchar(in_data[0])

    for char in in_data[1:]:
        if char in '\r\n':
            out_data[op] = char
            op += 1
            if char == '\n':
                x = shdiag_xchar(in_data[ip + 1])
        else:
            out_data[op] = (char - 0xf) ^ (x & 0x7f)
            x = (x >> 1) | ((x & 1) << 7)
            op += 1

    return out_data[:op]


def nb30port_open():
    dev = usb.core.find(idVendor=NB30_USB_VID, idProduct=NB30_USB_PID)
    if dev is None:
        raise Exception("Device not found!")
    dev.set_configuration()
    cfg = dev.get_active_configuration()
    intf = cfg[(1, 1)]
    ep_in = usb.util.find_descriptor(intf, custom_match=lambda e: usb.util.endpoint_direction(
        e.bEndpointAddress) == usb.util.ENDPOINT_IN)
    ep_out = usb.util.find_descriptor(intf, custom_match=lambda e: usb.util.endpoint_direction(
        e.bEndpointAddress) == usb.util.ENDPOINT_OUT)
    return dev, ep_in, ep_out


def nb30port_close(dev):
    usb.util.dispose_resources(dev)


def nb30port_read(ep_in):
    return ep_in.read(ep_in.wMaxPacketSize, 1000)


def nb30port_write(ep_out, data):
    return ep_out.write(data, 1000)


def main():
    if len(sys.argv) != 2:
        print("YuuSHDIAG v0.2")
        print("Usage: python shdiag.py <SHDIAG_COMMAND>")
        return 1

    dev, ep_in, ep_out = nb30port_open()

    res = nb30port_read(ep_in)
    if res:
        res_decode = shdiag_decode_data(res)
        print(''.join(map(chr, res_decode)))

    req = '{}\r\n'.format(sys.argv[1])
    req_encode = shdiag_encode_data(req)
    if req_encode and len(req_encode) == nb30port_write(ep_out, req_encode):
        res = nb30port_read(ep_in)
        if res:
            res_decode = shdiag_decode_data(res)
            print(''.join(map(chr, res_decode)))

    nb30port_close(dev)


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(str(e))