# -*- coding: utf-8 -*-

import os, stat, shutil, string


if __name__ == '__main__':

    ko = '/usr/lib/modules/i915.ko'
    subid = b'\x86\x80\x00\x00\x92\x3E\x00\x00'
    subend = b'\x30\x82\x01\xF7\x06\x09\x2A\x86'
    isChmod = False

    id = ''
    while True:
        id = input("[I] 请输入4位十六进制显卡ID(eg:7270):")
        if len(id) == 4 and all(c in string.hexdigits for c in id):
            break
        print("[E] 输入错误, 请重新输入.")
    
    if not os.path.exists(ko):
        print("[E] 未找到 {} 文件.".format(ko))
        exit(1)

    if not os.path.exists(ko+'.bak'):
        shutil.copyfile(ko, ko+'.bak')
        print('[I] 备份 {} 到 {} '.format(ko, ko+'.bak'))

    sBak = os.stat(ko).st_mode
    if sBak & stat.S_IWUSR != stat.S_IWUSR:
        os.chmod(ko, sBak + stat.S_IWUSR)
        isChmod = True

    data = None
    with open(ko+'.bak', 'rb') as f:
        data = f.read()

    idx = data.find(subid)
    if idx != -1:
        print('[I] 替换 {} 到 {} 的 {} 位置'.format(id, ko, idx))
        data = data.replace(subid, b''.fromhex('86800000'+id[2:4]+id[0:2]+'0000'))
    else:
        print("[E] {} 中未找到 {}, 替换失败.".format(ko, '86800000923E0000'))
        exit(2)

    idx = data.find(subend)
    if idx != -1:
        print('[I] 删除 {} 的 {} 位置的校验'.format(ko, idx))
        data = data[:idx]
    else:
        print('[W] {} 中未找到校验位, 跳过.'.format(idx))

    with open(ko, 'wb') as f:
        f.write(data)
    
    if isChmod:
        os.chmod(ko, sBak)

    print('[I] 修改完成, 请重启确认, (源文件备份到{}, 如有问题请还原).'.format(ko+'.bak'))