# -*- coding: cp936 -*-
from winappdbg import EventHandler, Debug
from winappdbg.win32 import *
from ctypes import Structure as Struct
from custom_config import strs, program_dir, os

try:
    import psyco
    psyco.full()
except ImportError:
    pass


class lvColumn(Struct):
    _fields_ = [('mask', UINT),
                ('fmt', INT),
                ('cx', INT),
                ('pszText', PVOID),
                ('cchTextMax', INT),
                ('iSubItem', INT),
                ('iOrder', INT),
                ('iImage', INT)]


class lvItem(Struct):
    _fields_ = [('mask', UINT),
                ('iItem', INT),
                ('iSubItem', INT),
                ('state', UINT),
                ('stateMask', UINT),
                ('pszText', PVOID),
                ('cchTextMax', INT),
                ('iImage', INT)]


class MsgBox(Struct):
    _fields_ = [('cbSize', UINT),
                ('hwndOwner', HWND),
                ('hInstance', HINSTANCE),
                ('lpszText', PVOID),
                ('lpszCaption', PVOID),
                ('dwStyle', DWORD),
                ('lpszIcon', PVOID),
                ('dwContextHelpId', DWORD_PTR),
                ('lpfnMsgBoxCallback', PVOID),
                ('dwLanguageId', DWORD)]


class OpenFileName(Struct):
    _fields_ = [('lStructSize', DWORD),
                ('hwndOwner', HWND),
                ('hInstance', HINSTANCE),
                ('lpstrFilter', PVOID),
                ('lpstrCustomFilter', LPSTR),
                ('nMaxCustFilter', DWORD),
                ('nFilterIndex', DWORD),
                ('lpstrFile', LPSTR),
                ('nMaxFile', DWORD),
                ('lpstrFileTitle', LPSTR),
                ('nMaxFileTitle', DWORD),
                ('lpstrInitialDir', PVOID),
                ('lpstrTitle', PVOID),
                ('Flags', DWORD),
                ('nFileOffset', WORD),
                ('nFileExtension', WORD),
                ('lpstrDefExt', PVOID),
                ('lCustData', LPARAM),
                ('lpfnHook', DWORD),
                ('lpTemplateName', PVOID)]


class MyEventHandler(EventHandler):
    apiHooks = {
        'kernel32.dll': [
            ('CreateFileA', (PVOID, DWORD, DWORD, PVOID, DWORD, DWORD, HANDLE))
        ],
        'comdlg32.dll': [
            ('GetOpenFileNameA', (PVOID,)), ('GetSaveFileNameA', (PVOID,))
        ],
        'user32.dll': [
            #  Function            Parameters
            ('CreateWindowExA', (DWORD, PVOID, PVOID, DWORD,
                                 INT, INT, INT, INT, HWND, HWND, HANDLE, LPVOID)),
            ('SetWindowTextA', (HWND, PVOID)),
            ('SendMessageA', (HWND, UINT, WPARAM, LPARAM)),
            ('AppendMenuA', (HWND, UINT, UINT, PVOID)),
            ('TabbedTextOutA', (HDC, INT, INT, PVOID, INT, INT, PVOID, INT)),
            ('MessageBoxA', (HWND, PVOID, PVOID, UINT)),
            ('MessageBoxIndirectA', (PVOID,))
        ],
    }

    def pre_CreateFileA(self, event, ra, lpFileName, dwDesiredAccess,
                        dwShareMode, lpSecurityAttributes, dwCreationDisposition,
                        dwFlagsAndAttributes, hTemplateFile):
        self.replace_linebreak(event, lpFileName)

    def replace_linebreak(self, event, pointer):
        process = event.get_process()
        string = process.peek_string(pointer)
        if '\r\n' in string:
            tobe = string.replace('\r\n', '').ljust(len(string), ' ')
            process.write(pointer, tobe)

    def pre_GetOpenFileNameA(self, event, ra, lpofn):
        handle_filter(event, lpofn)

    def pre_GetSaveFileNameA(self, event, ra, lpofn):
        handle_filter(event, lpofn)

    def pre_SendMessageA(self, event, ra, hWnd, Msg, wParam, lParam):
        if Msg == 323:
            replace_text(event, lParam)
        elif Msg == 4123:
            tmp = event.get_process().read_structure(lParam, lvColumn)
            replace_text(event, tmp.pszText)
        elif Msg == 4103 or Msg == 4102:
            tmp = event.get_process().read_structure(lParam, lvItem)
            if isinstance(tmp.pszText, int):
                replace_text(event, tmp.pszText + (Msg - 4102) * 4)

    def pre_MessageBoxIndirectA(self, event, ra, lpMsgBoxParams):
        tmp = event.get_process().read_structure(lpMsgBoxParams, MsgBox)
        replace_text(event, tmp.lpszText, space=' ')

    def pre_TabbedTextOutA(self, event, ra, hDc, X, Y, lpString, nCount, nTabPositions, TabStopPositions, nTabOrigin):
        replace_text(event, lpString)

    def pre_MessageBoxA(self, event, ra, hWnd, lpText, lpCaption, uType):
        replace_text(event, lpText, space=' ')

    def pre_AppendMenuA(self, event, ra, hMenu, uFlags, uIDNewItem, lpNewItem):
        replace_text(event, lpNewItem)

    def pre_SetWindowTextA(self, event, ra, hWnd, lpString):
        replace_text(event, lpString)

    def pre_CreateWindowExA(self, event, ra, dwExStyle, lpClassName, lpWindowName,
                            dwStyle, x, y, nWidth, nHeight, hWndParent, hMenu,
                            hInstance, lpParam):
        replace_text(event, lpWindowName)


def handle_filter(event, lpofn):
    process = event.get_process()
    tmp = process.read_structure(lpofn, OpenFileName).lpstrFilter
    while tmp is not None:
        string = process.peek_string(tmp)
        if len(string):
            if ja2cn(string) in strs['text']:
                process.write(tmp, strs['text'][ja2cn(string)])
            tmp += len(string) + 1
        else:
            tmp = None


def replace_text(event, pointer, space='\r'):
    process = event.get_process()
    if process.is_address_writeable(pointer):
        text = process.peek_string(pointer)
        matched = False
        for string in text.split("\r\n"):
            try:
                for target in strs['middle']:
                    if target in string:
                        string = string.replace(
                            target, cn2ja(strs['middle'][target]))
                if string in strs['special']:
                    tobe = strs['special'][string]
                else:
                    tobe = ja2cn(string)
                    if tobe in strs['text'] and len(tobe) >= len(strs['text'][tobe]):
                        tobe = strs['text'][tobe]
                        matched = True
                    if len(tobe) < len(string):
                        tobe = tobe.ljust(len(string), space)
                process.write(pointer, tobe)
            except (WindowsError, UnicodeDecodeError, UnicodeEncodeError):
                pass
            finally:
                if not matched:
                    for target in strs['special']:
                        if target == string:
                            process.write(pointer, strs['special'][target])
                    for item in sorted(strs['prefix'].keys(), key=len, reverse=True):
                        if string.startswith(cn2ja(item)):
                            process.write(pointer, strs['prefix'][item].ljust(len(item), space))
                            break
                    for item in strs['suffix']:
                        if string.endswith(cn2ja(item)):
                            process.write(pointer + len(string) -
                                          len(item), strs['suffix'][item])
                            break
                    try:
                        print (ja2cn(string))
                    except (UnicodeDecodeError, UnicodeEncodeError):
                        print (string)
            pointer += len(string) + 2


def ja2cn(word):
    return word.decode('shift-jis').encode('gbk')


def cn2ja(word):
    return word.decode('gbk').encode('shift-jis')


def main():
    with Debug(MyEventHandler(), bKillOnExit=True) as debug:
        debug.execl(os.path.join(program_dir, os.path.pardir, strs['config']['binary']))
        debug.loop()


if __name__ == '__main__':
    main()
