# Author: Nicolas VERDIER
# This file is part of memorpy.
#
# memorpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# memorpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with memorpy.  If not, see <http://www.gnu.org/licenses/>.

from ctypes import (
    pointer,
    sizeof,
    windll,
    create_string_buffer,
    c_ulong,
    byref,
    GetLastError,
    c_bool,
    WinError,
)
from .WinStructures import *
import copy
import struct
from . import utils
import platform
from .BaseProcess import BaseProcess, ProcessException

psapi = windll.psapi
kernel32 = windll.kernel32
advapi32 = windll.advapi32

IsWow64Process = None
if hasattr(kernel32, "IsWow64Process"):
    IsWow64Process = kernel32.IsWow64Process
    IsWow64Process.restype = c_bool
    IsWow64Process.argtypes = [c_void_p, POINTER(c_bool)]


class WinProcess(BaseProcess):
    def __init__(self, pid=None, name=None, debug=True):
        """ Create and Open a process object from its pid or from its name """
        super(WinProcess, self).__init__()
        if pid:
            self._open(int(pid), debug=debug)

        elif name:
            self._open_from_name(name, debug=debug)
        else:
            raise ValueError(
                "You need to instanciate process with at least a name or a pid"
            )

        if self.is_64bit():
            si = self.GetNativeSystemInfo()
            self.max_addr = si.lpMaximumApplicationAddress
        else:
            si = self.GetSystemInfo()
            self.max_addr = 2147418111
        self.min_addr = si.lpMinimumApplicationAddress

    def __del__(self):
        self.close()

    def is_64bit(self):
        if not "64" in platform.machine():
            return False
        iswow64 = c_bool(False)
        if IsWow64Process is None:
            return False
        if not IsWow64Process(self.h_process, byref(iswow64)):
            raise WinError()
        return not iswow64.value

    @staticmethod
    def list():
        processes = []
        arr = c_ulong * 256
        lpidProcess = arr()
        cb = sizeof(lpidProcess)
        cbNeeded = c_ulong()
        hModule = c_ulong()
        count = c_ulong()
        modname = create_string_buffer(100)
        PROCESS_QUERY_INFORMATION = 0x0400
        PROCESS_VM_READ = 0x0010

        psapi.EnumProcesses(byref(lpidProcess), cb, byref(cbNeeded))
        nReturned = cbNeeded.value / sizeof(c_ulong())

        pidProcess = [i for i in lpidProcess][:nReturned]
        for pid in pidProcess:
            proc = {"pid": int(pid)}
            hProcess = kernel32.OpenProcess(
                PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, False, pid
            )
            if hProcess:
                psapi.EnumProcessModules(
                    hProcess, byref(hModule), sizeof(hModule), byref(count)
                )
                psapi.GetModuleBaseNameA(
                    hProcess, hModule.value, modname, sizeof(modname)
                )
                proc["name"] = modname.value
                kernel32.CloseHandle(hProcess)
            processes.append(proc)
        return processes

    @staticmethod
    def processes_from_name(processName):
        processes = []
        for process in WinProcess.list():
            if processName == process.get("name", None) or (
                process.get("name", "").lower().endswith(".exe")
                and process.get("name", "")[:-4] == processName
            ):
                processes.append(process)

        if len(processes) > 0:
            return processes

    @staticmethod
    def name_from_process(dwProcessId):
        process_list = WinProcess.list()
        for process in process_list:
            if process.pid == dwProcessId:
                return process.get("name", None)

        return False

    def _open(self, dwProcessId, debug=False):
        if debug:
            ppsidOwner = DWORD()
            ppsidGroup = DWORD()
            ppDacl = DWORD()
            ppSacl = DWORD()
            ppSecurityDescriptor = SECURITY_DESCRIPTOR()

            process = kernel32.OpenProcess(262144, 0, dwProcessId)
            advapi32.GetSecurityInfo(
                kernel32.GetCurrentProcess(),
                6,
                0,
                byref(ppsidOwner),
                byref(ppsidGroup),
                byref(ppDacl),
                byref(ppSacl),
                byref(ppSecurityDescriptor),
            )
            advapi32.SetSecurityInfo(
                process,
                6,
                DACL_SECURITY_INFORMATION | UNPROTECTED_DACL_SECURITY_INFORMATION,
                None,
                None,
                ppSecurityDescriptor.dacl,
                ppSecurityDescriptor.group,
            )
            kernel32.CloseHandle(process)
        self.h_process = kernel32.OpenProcess(2035711, 0, dwProcessId)
        if self.h_process is not None:
            self.isProcessOpen = True
            self.pid = dwProcessId
            return True
        return False

    def close(self):
        if self.h_process is not None:
            ret = kernel32.CloseHandle(self.h_process) == 1
            if ret:
                self.h_process = None
                self.pid = None
                self.isProcessOpen = False
            return ret
        return False

    def _open_from_name(self, processName, debug=False):
        processes = self.processes_from_name(processName)
        if not processes:
            raise ProcessException("can't get pid from name %s" % processName)
        elif len(processes) > 1:
            raise ValueError(
                "There is multiple processes with name %s. Please select a process from its pid instead"
                % processName
            )
        if debug:
            self._open(processes[0]["pid"], debug=True)
        else:
            self._open(processes[0]["pid"], debug=False)

    def GetSystemInfo(self):
        si = SYSTEM_INFO()
        kernel32.GetSystemInfo(byref(si))
        return si

    def GetNativeSystemInfo(self):
        si = SYSTEM_INFO()
        kernel32.GetNativeSystemInfo(byref(si))
        return si

    def VirtualQueryEx(self, lpAddress):
        mbi = MEMORY_BASIC_INFORMATION()
        if not VirtualQueryEx(self.h_process, lpAddress, byref(mbi), sizeof(mbi)):
            raise ProcessException("Error VirtualQueryEx: 0x%08X" % lpAddress)
        return mbi

    def VirtualQueryEx64(self, lpAddress):
        mbi = MEMORY_BASIC_INFORMATION64()
        if not VirtualQueryEx64(self.h_process, lpAddress, byref(mbi), sizeof(mbi)):
            raise ProcessException("Error VirtualQueryEx: 0x%08X" % lpAddress)
        return mbi

    def VirtualProtectEx(self, base_address, size, protection):
        old_protect = c_ulong(0)
        if not kernel32.VirtualProtectEx(
            self.h_process, base_address, size, protection, byref(old_protect)
        ):
            raise ProcessException(
                "Error: VirtualProtectEx(%08X, %d, %08X)"
                % (base_address, size, protection)
            )
        return old_protect.value

    def iter_region(
        self, start_offset=None, end_offset=None, protec=None, optimizations=None
    ):

        offset = start_offset or self.min_addr
        end_offset = end_offset or self.max_addr

        while True:
            if offset >= end_offset:
                break
            mbi = self.VirtualQueryEx(offset)
            offset = mbi.BaseAddress
            chunk = mbi.RegionSize
            protect = mbi.Protect
            state = mbi.State
            # print "offset: %s, chunk:%s"%(offset, chunk)
            if state & MEM_FREE or state & MEM_RESERVE:
                offset += chunk
                continue
            if protec:
                if (
                    not protect & protec
                    or protect & PAGE_NOCACHE
                    or protect & PAGE_WRITECOMBINE
                    or protect & PAGE_GUARD
                ):
                    offset += chunk
                    continue
            yield offset, chunk
            offset += chunk

    def write_bytes(self, address, data):
        address = int(address)
        if not self.isProcessOpen:
            raise ProcessException(
                "Can't write_bytes(%s, %s), process %s is not open"
                % (address, data, self.pid)
            )
        buffer = create_string_buffer(data)
        sizeWriten = c_size_t(0)
        bufferSize = sizeof(buffer) - 1
        _address = address
        _length = bufferSize + 1
        try:
            old_protect = self.VirtualProtectEx(
                _address, _length, PAGE_EXECUTE_READWRITE
            )
        except:
            pass

        res = kernel32.WriteProcessMemory(
            self.h_process, address, buffer, bufferSize, byref(sizeWriten)
        )
        try:
            self.VirtualProtectEx(_address, _length, old_protect)
        except:
            pass

        return res

    def read_bytes(self, address, bytes=4, use_NtWow64ReadVirtualMemory64=False):
        # print "reading %s bytes from addr %s"%(bytes, address)
        if use_NtWow64ReadVirtualMemory64:
            if NtWow64ReadVirtualMemory64 is None:
                raise WindowsError(
                    "NtWow64ReadVirtualMemory64 is not available from a 64bit process"
                )
            RpM = NtWow64ReadVirtualMemory64
        else:
            RpM = ReadProcessMemory

        address = int(address)
        buffer = create_string_buffer(bytes)
        bytesread = c_size_t(0)
        data = b""
        length = bytes
        while length:
            if RpM(self.h_process, address, buffer, bytes, byref(bytesread)) or (
                use_NtWow64ReadVirtualMemory64 and GetLastError() == 0
            ):
                if bytesread.value:
                    data += buffer.raw[: bytesread.value]
                    length -= bytesread.value
                    address += bytesread.value
                if not len(data):
                    raise ProcessException(
                        "Error %s in ReadProcessMemory(%08x, %d, read=%d)"
                        % (GetLastError(), address, length, bytesread.value)
                    )
                return data
            else:
                if (
                    GetLastError() == 299
                ):  # only part of ReadProcessMemory has been done, let's return it
                    data += buffer.raw[: bytesread.value]
                    return data
                raise WinError()
            # data += buffer.raw[:bytesread.value]
            # length -= bytesread.value
            # address += bytesread.value
        return data

    def list_modules(self):
        module_list = []
        if self.pid is not None:
            hModuleSnap = CreateToolhelp32Snapshot(TH32CS_CLASS.SNAPMODULE, self.pid)
            if hModuleSnap is not None:
                module_entry = MODULEENTRY32()
                module_entry.dwSize = sizeof(module_entry)
                success = Module32First(hModuleSnap, byref(module_entry))
                while success:
                    if module_entry.th32ProcessID == self.pid:
                        module_list.append(copy.copy(module_entry))
                    success = Module32Next(hModuleSnap, byref(module_entry))

                kernel32.CloseHandle(hModuleSnap)
        return module_list

    def get_symbolic_name(self, address):
        for m in self.list_modules():
            if int(m.modBaseAddr) <= int(address) < int(m.modBaseAddr + m.modBaseSize):
                return "%s+0x%08X" % (m.szModule, int(address) - m.modBaseAddr)

        return "0x%08X" % int(address)

    def hasModule(self, module):
        if module[-4:] != ".dll":
            module += ".dll"
        module_list = self.list_modules()
        for m in module_list:
            if module in m.szExePath.split("\\"):
                return True
        return False

    def get_instruction(self, address):
        """
        Pydasm disassemble utility function wrapper. Returns the pydasm decoded instruction in self.instruction.
        """
        import pydasm

        try:
            data = self.read_bytes(int(address), 32)
        except:
            return "Unable to disassemble at %08x" % address

        return pydasm.get_instruction(data, pydasm.MODE_32)
