#!/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()

[((208, 262), 're.search', 're.search', (['"""#define SW2_SEED (0x[a-fA-F0-9]{8})"""', 'code'], {}), "('#define SW2_SEED (0x[a-fA-F0-9]{8})', code)\n", (217, 262), False, 'import re\n'), ((1357, 1411), 're.compile', 're.compile', (['"""__declspec\\\\(naked\\\\) NTSTATUS (Nt[^(]+)"""'], {}), "('__declspec\\\\(naked\\\\) NTSTATUS (Nt[^(]+)')\n", (1367, 1411), False, 'import re\n'), ((1431, 1454), 're.findall', 're.findall', (['regex', 'code'], {}), '(regex, code)\n', (1441, 1454), False, 'import re\n'), ((2820, 2863), 'glob.glob', 'glob.glob', (['"""include/**/*.h"""'], {'recursive': '(True)'}), "('include/**/*.h', recursive=True)\n", (2829, 2863), False, 'import glob\n'), ((3497, 3533), 'random.randint', 'random.randint', (['(2 ** 28)', '(2 ** 32 - 1)'], {}), '(2 ** 28, 2 ** 32 - 1)\n', (3511, 3533), False, 'import random\n'), ((1617, 1709), 're.compile', 're.compile', (['(\'NTSTATUS \' + syscall_name + \'\\\\(.*?"mov ecx, (0x[A-Fa-f0-9]{8})\')', 're.DOTALL'], {}), '(\'NTSTATUS \' + syscall_name +\n    \'\\\\(.*?"mov ecx, (0x[A-Fa-f0-9]{8})\', re.DOTALL)\n', (1627, 1709), False, 'import re\n'), ((1722, 1759), 're.search', 're.search', (['regex', 'syscall_definitions'], {}), '(regex, syscall_definitions)\n', (1731, 1759), False, 'import re\n'), ((2286, 2361), 're.compile', 're.compile', (["(syscall_name + ' PROC.*?mov ecx, 0([A-Fa-f0-9]{8})h')", 're.DOTALL'], {}), "(syscall_name + ' PROC.*?mov ecx, 0([A-Fa-f0-9]{8})h', re.DOTALL)\n", (2296, 2361), False, 'import re\n'), ((2378, 2400), 're.search', 're.search', (['regex', 'code'], {}), '(regex, code)\n', (2387, 2400), False, 'import re\n'), ((2946, 3002), 're.compile', 're.compile', (['"""#define (\\\\w+)_SW2_HASH (0x[a-fA-F0-9]{8})"""'], {}), "('#define (\\\\w+)_SW2_HASH (0x[a-fA-F0-9]{8})')\n", (2956, 3002), False, 'import re\n'), ((3021, 3044), 're.findall', 're.findall', (['regex', 'code'], {}), '(regex, code)\n', (3031, 3044), False, 'import re\n')]