module hooking.jmphook;

import std.c.windows.windows;
import hooking.processmemory;
public import hooking.define;

class JmpHook(TResult, TArguments...)
{

    alias extern(Windows) TResult function(TArguments) StdCallFunc;
    alias extern(C) TResult function(TArguments) CdeclFunc;

    // Fields
    private static __gshared JmpHook[void*] _hooks;

    private void* _address;
    private void* _callback;
    private CallingConvention _callingConvention;
    private byte _length;
    private byte[] _newBytes;
    private byte[] _originalBytes;

    // Constructors
    public this(void* address, void* callback, CallingConvention callingConvention, byte length = 5)
    {

        if (address == null)
            throw new Exception("Address cannot be null");

        if (length < 5)
            throw new Exception("Length cannot be less than 5");

        _address = address;
        _callback = callback;
        _callingConvention = callingConvention;
        _length = length;
        _newBytes = new byte[_length + 5];

        _originalBytes = ProcessMemory.readBytes(address, length);

    }

    public this(string moduleName, string procedureName, void* callback, CallingConvention callingConvention, byte length = 5)
    {
        this(getModuleProcedureAddress(moduleName, procedureName), callback, callingConvention, length);
    }

    // Public
    @property
    public bool isInstalled()
    {
        JmpHook hook = _hooks.get(_address, null);
        return hook !is null;
    }

    @property
    public void* newFunctionAddress()
    {
        return &_newBytes[0];
    }

    public TResult invokeOriginal(void* thisPtr, TArguments arguments)
    {

        switch (_callingConvention)
        {

            case CallingConvention.StdCall:
                asm { mov ECX, thisPtr; }
                return (cast(StdCallFunc)&_newBytes[0])(arguments);

            case CallingConvention.Cdecl:
                asm { mov ECX, thisPtr; }
                return (cast(CdeclFunc)&_newBytes[0])(arguments);

            default:
                throw new Exception("Calling convention not supported");

        }

    }

    public void install()
    {

        if (!isInstalled)
        {

            byte* newBytesPtr = &_newBytes[0];

            if (_originalBytes[0] == cast(byte)0xE9) // Check if far jump instruction
            {

                uint jumpOffset = *cast(uint*)(_address + 1);
                void* jumpDestination = cast(void*)(_address + 5) + (jumpOffset);

                ProcessMemory.writeBytes(newBytesPtr, _originalBytes);

                JmpPatch jmpOldHookPatch = JmpPatch(newBytesPtr, jumpDestination);
                ProcessMemory.write!(JmpPatch)(newBytesPtr, jmpOldHookPatch);

            }
            else if (_originalBytes[0] == cast(byte)0xEB) // Check if short jump instruction
            {

                byte jumpOffset = *cast(byte*)(_address + 1);
                void* jumpDestination = cast(void*)(_address + 2) + (jumpOffset);

                ProcessMemory.writeBytes(newBytesPtr, _originalBytes);

                JmpPatch jmpOldHookPatch = JmpPatch(newBytesPtr, jumpDestination);
                ProcessMemory.write!(JmpPatch)(newBytesPtr, jmpOldHookPatch);

            }
            else
                ProcessMemory.writeBytes(newBytesPtr, _originalBytes);

            // Write the (jmp hookaddress + hooklength) to the new bytes
            JmpPatch jmpBackPatch = JmpPatch(&_newBytes[_length], _address + _length);
            ProcessMemory.write!(JmpPatch)(&_newBytes[_length], jmpBackPatch);

            // Write the jmp patch
            JmpPatch jmpPatch = JmpPatch(_address, _callback);
            ProcessMemory.write!(JmpPatch)(_address, jmpPatch);

            if (_length > 5)
            {

                byte[] nops = new byte[_length - 5];

                for (int i = 0; i < _length - 5; i++)
                    nops[i] = cast(byte)0x90;

                ProcessMemory.writeBytes(_address + 5, nops);

            }

            _hooks[_address] = this;

        }

    }

    public void uninstall()
    {
        if (isInstalled)
        {
            ProcessMemory.writeBytes(_address, _originalBytes);
            _hooks[_address] = null;
        }
    }

}