#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import re
import glob
import random
import struct


def get_old_seed():
    with open('include/syscalls.h') as f:
        code = f.read()
    match = re.search(r'#define SW2_SEED (0x[a-fA-F0-9]{8})', code)
    assert match is not None, 'SW2_SEED not found!'
    return match.group(1)


def replace_seed(old_seed, new_seed):
    with open('include/syscalls.h') as f:
        code = f.read()
    code = code.replace(
        f'#define SW2_SEED {old_seed}',
        f'#define SW2_SEED 0x{new_seed:08X}',
        1
    )
    with open('include/syscalls.h', 'w') as f:
        f.write(code)


def get_function_hash(seed, function_name, is_syscall=True):
    function_hash = seed
    function_name = function_name.replace('_', '')
    if is_syscall and function_name[:2] == 'Nt':
        function_name = 'Zw' + function_name[2:]
    name = function_name + '\0'
    ror8 = lambda v: ((v >> 8) & (2 ** 32 - 1)) | ((v << 24) & (2 ** 32 - 1))

    for segment in [s for s in [name[i:i + 2] for i in range(len(name))] if len(s) == 2]:
        partial_name_short = struct.unpack('<H', segment.encode())[0]
        function_hash ^= partial_name_short + ror8(function_hash)

    return function_hash


def replace_syscall_hashes(seed):
    with open('source/syscalls.c') as f:
        code = f.read()
    regex = re.compile(r'__declspec\(naked\) NTSTATUS (Nt[^(]+)')
    syscall_names = re.findall(regex, code)
    syscall_names = set(syscall_names)
    syscall_definitions = code.split('#elif defined(__GNUC__)')[3]

    for syscall_name in syscall_names:
        regex = re.compile('NTSTATUS ' + syscall_name + '\\(.*?"mov ecx, (0x[A-Fa-f0-9]{8})', re.DOTALL)
        match = re.search(regex, syscall_definitions)
        assert match is not None, f'hash of syscall {syscall_name} not found!'
        old_hash = match.group(1)
        new_hash = get_function_hash(seed, syscall_name)
        print(f'{syscall_name} -> {old_hash} - 0x{new_hash:08X}')
        code = code.replace(
            old_hash,
            f'0x{new_hash:08X}'
        )

    with open('source/syscalls.c', 'w') as f:
        f.write(code)

    with open('source/syscalls-asm.asm') as f:
        code = f.read()

    for syscall_name in syscall_names:
        regex = re.compile(syscall_name + ' PROC.*?mov ecx, 0([A-Fa-f0-9]{8})h', re.DOTALL)
        match = re.search(regex, code)
        assert match is not None, f'hash of syscall {syscall_name} not found!'
        old_hash = match.group(1)
        new_hash = get_function_hash(seed, syscall_name)
        code = code.replace(
            f'0{old_hash}h',
            f'0{new_hash:08X}h',
            1
        )

    with open('source/syscalls-asm.asm', 'w') as f:
        f.write(code)


def replace_dinvoke_hashes(seed):
    for header_file in glob.glob("include/**/*.h", recursive=True):
        with open(header_file) as f:
            code = f.read()
        regex = re.compile(r'#define (\w+)_SW2_HASH (0x[a-fA-F0-9]{8})')
        matches = re.findall(regex, code)
        for function_name, old_hash in matches:
            new_hash = get_function_hash(seed, function_name, is_syscall=False)
            code = code.replace(
                f'#define {function_name}_SW2_HASH {old_hash}',
                f'#define {function_name}_SW2_HASH 0x{new_hash:08X}',
                1
            )
        if matches:
            with open(header_file, 'w') as f:
                f.write(code)


def main():
    new_seed = random.randint(2 ** 28, 2 ** 32 - 1)
    #new_seed = 0x1337c0de
    old_seed = get_old_seed()
    replace_seed(old_seed, new_seed)
    replace_syscall_hashes(new_seed)
    replace_dinvoke_hashes(new_seed)
    if os.name == 'nt':
        print('done! recompile with:\nnmake -f Makefile.msvc')
    else:
        print('done! recompile with:\nmake -f Makefile.mingw')

if __name__ == '__main__':
    main()

[((14, 12, 14, 67), 're.search', 're.search', ({(14, 22, 14, 60): '"""#define SW2_SEED (0x[a-fA-F0-9]{8})"""', (14, 62, 14, 66): 'code'}, {}), "('#define SW2_SEED (0x[a-fA-F0-9]{8})', code)", False, 'import re\n'), ((49, 12, 49, 65), 're.compile', 're.compile', ({(49, 23, 49, 64): '"""__declspec\\\\(naked\\\\) NTSTATUS (Nt[^(]+)"""'}, {}), "('__declspec\\\\(naked\\\\) NTSTATUS (Nt[^(]+)')", False, 'import re\n'), ((50, 20, 50, 43), 're.findall', 're.findall', ({(50, 31, 50, 36): 'regex', (50, 38, 50, 42): 'code'}, {}), '(regex, code)', False, 'import re\n'), ((89, 23, 89, 66), 'glob.glob', 'glob.glob', (), '', False, 'import glob\n'), ((107, 15, 107, 51), 'random.randint', 'random.randint', ({(107, 30, 107, 37): '2 ** 28', (107, 39, 107, 50): '2 ** 32 - 1'}, {}), '(2 ** 28, 2 ** 32 - 1)', False, 'import random\n'), ((55, 16, 55, 104), 're.compile', 're.compile', ({(55, 27, 55, 92): '\'NTSTATUS \' + syscall_name + \'\\\\(.*?"mov ecx, (0x[A-Fa-f0-9]{8})\'', (55, 94, 55, 103): 're.DOTALL'}, {}), '(\'NTSTATUS \' + syscall_name +\n    \'\\\\(.*?"mov ecx, (0x[A-Fa-f0-9]{8})\', re.DOTALL)', False, 'import re\n'), ((56, 16, 56, 53), 're.search', 're.search', ({(56, 26, 56, 31): 'regex', (56, 33, 56, 52): 'syscall_definitions'}, {}), '(regex, syscall_definitions)', False, 'import re\n'), ((73, 16, 73, 91), 're.compile', 're.compile', ({(73, 27, 73, 79): "syscall_name + ' PROC.*?mov ecx, 0([A-Fa-f0-9]{8})h'", (73, 81, 73, 90): 're.DOTALL'}, {}), "(syscall_name + ' PROC.*?mov ecx, 0([A-Fa-f0-9]{8})h', re.DOTALL)", False, 'import re\n'), ((74, 16, 74, 38), 're.search', 're.search', ({(74, 26, 74, 31): 'regex', (74, 33, 74, 37): 'code'}, {}), '(regex, code)', False, 'import re\n'), ((92, 16, 92, 72), 're.compile', 're.compile', ({(92, 27, 92, 71): '"""#define (\\\\w+)_SW2_HASH (0x[a-fA-F0-9]{8})"""'}, {}), "('#define (\\\\w+)_SW2_HASH (0x[a-fA-F0-9]{8})')", False, 'import re\n'), ((93, 18, 93, 41), 're.findall', 're.findall', ({(93, 29, 93, 34): 'regex', (93, 36, 93, 40): 'code'}, {}), '(regex, code)', False, 'import re\n')]