using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using WhiteMagic.Native;

namespace WhiteMagic.Injection;

/// <summary>
/// Injects DLLs into an open target process by creating a remote thread or by
/// hijacking an existing thread.
/// </summary>
/// <remarks>
/// <para>
/// The DLL path is sent to <c>LoadLibraryW</c>, so it is encoded as a null-terminated
/// UTF-16 string in the target process.
/// </para>
/// <para>
/// Injection requires the target process to have the same bitness as the current
/// process, because the emitted x86/x64 stubs and the captured thread context must
/// match the target architecture.
/// </para>
/// </remarks>
public sealed class DllInjector
{
    private readonly MemoryBase _memory;
    private readonly bool _currentIs64Bit;

    /// <summary>
    /// Initializes a new <see cref="DllInjector"/> for the target represented by
    /// <paramref name="memory"/>.
    /// </summary>
    /// <param name="memory">A reader/writer for the target process.</param>
    public DllInjector(MemoryBase memory)
    {
        ArgumentNullException.ThrowIfNull(memory);
        _memory = memory;
        _currentIs64Bit = Environment.Is64BitProcess;
    }

    /// <summary>
    /// Gets the <see cref="MemoryBase"/> the injector is operating on.
    /// </summary>
    public MemoryBase Memory => _memory;

    /// <summary>
    /// Injects a DLL into the target process by creating a remote thread that loads it.
    /// </summary>
    /// <param name="dllPath">The path to the DLL. The file must exist.</param>
    /// <returns>The base address of the loaded module in the target process.</returns>
    /// <exception cref="ArgumentException"><paramref name="dllPath"/> is null or empty.</exception>
    /// <exception cref="FileNotFoundException"><paramref name="dllPath"/> does not exist.</exception>
    /// <exception cref="InvalidOperationException">The target bitness does not match the caller.</exception>
    /// <exception cref="InvalidOperationException">The remote load failed or timed out.</exception>
    public IntPtr InjectWithRemoteThread(string dllPath)
    {
        ValidateAndCheckBitness(dllPath);

        IntPtr loadLibrary = ResolveLoadLibraryW();
        byte[] pathBytes = Encoding.Unicode.GetBytes(dllPath + '\0');

        int pointerSize = _currentIs64Bit ? 8 : 4;
        int stubSize = _currentIs64Bit ? 39 : 18;
        int pathOffset = Align(stubSize, pointerSize);
        int resultOffset = Align(pathOffset + pathBytes.Length, pointerSize);
        int totalSize = resultOffset + pointerSize + 4096;

        IntPtr remoteBase = NativeMethods.VirtualAllocEx(
            _memory.Handle,
            IntPtr.Zero,
            totalSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);

        if (remoteBase == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"VirtualAllocEx failed (error {error}).");
        }

        try
        {
            IntPtr pathAddress = remoteBase + pathOffset;
            IntPtr resultAddress = remoteBase + resultOffset;

            if (_memory.WriteBytes(pathAddress, pathBytes) != pathBytes.Length)
                throw new InvalidOperationException("Failed to write the DLL path into the target process.");

            byte[] stub = _currentIs64Bit
                ? BuildRemoteThreadStubX64(pathAddress, resultAddress, loadLibrary)
                : BuildRemoteThreadStubX86(pathAddress, resultAddress, loadLibrary);

            if (_memory.WriteBytes(remoteBase, stub) != stub.Length)
                throw new InvalidOperationException("Failed to write the remote thread stub.");

            using SafeMemoryHandle thread = NativeMethods.CreateRemoteThread(
                _memory.Handle,
                IntPtr.Zero,
                0,
                remoteBase,
                IntPtr.Zero,
                ThreadCreationFlags.RunImmediately,
                out _);

            if (thread.IsInvalid)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException($"CreateRemoteThread failed (error {error}).");
            }

            const uint timeoutMs = 30000;
            uint wait = NativeMethods.WaitForSingleObject(thread, timeoutMs);
            if (wait == 0xFFFFFFFF)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException($"WaitForSingleObject failed (error {error}).");
            }

            if (wait == 0x00000102)
                throw new InvalidOperationException("Remote thread timed out while loading the DLL.");

            IntPtr result = _memory.Read<IntPtr>(resultAddress);
            if (result == IntPtr.Zero)
                throw new InvalidOperationException("LoadLibrary returned zero; the DLL could not be loaded.");

            return result;
        }
        finally
        {
            // The DLL is already loaded; the temporary stub, path and result slot can be released.
            NativeMethods.VirtualFreeEx(_memory.Handle, remoteBase, 0, MemoryFreeType.Release);
        }
    }

    /// <summary>
    /// Injects a DLL by hijacking an existing thread in the target process.
    /// </summary>
    /// <param name="threadId">The operating-system identifier of the thread to hijack.</param>
    /// <param name="dllPath">The path to the DLL. The file must exist.</param>
    /// <returns>The base address of the loaded module in the target process.</returns>
    /// <exception cref="ArgumentException"><paramref name="threadId"/> is not a positive value or
    /// <paramref name="dllPath"/> is null or empty.</exception>
    /// <exception cref="FileNotFoundException"><paramref name="dllPath"/> does not exist.</exception>
    /// <exception cref="InvalidOperationException">The target bitness does not match the caller.</exception>
    /// <exception cref="InvalidOperationException">The hijack, load, or context restore failed.</exception>
    public IntPtr InjectWithThreadHijack(int threadId, string dllPath)
    {
        if (threadId <= 0)
            throw new ArgumentException("Thread ID must be a positive value.", nameof(threadId));

        ValidateAndCheckBitness(dllPath);

        IntPtr loadLibrary = ResolveLoadLibraryW();
        byte[] pathBytes = Encoding.Unicode.GetBytes(dllPath + '\0');

        int pointerSize = _currentIs64Bit ? 8 : 4;
        int stubSize = _currentIs64Bit ? 40 : 19;
        int pathOffset = Align(stubSize, pointerSize);
        int resultOffset = Align(pathOffset + pathBytes.Length, pointerSize);
        int totalSize = resultOffset + pointerSize + 4096;

        IntPtr remoteBase = NativeMethods.VirtualAllocEx(
            _memory.Handle,
            IntPtr.Zero,
            totalSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);

        if (remoteBase == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"VirtualAllocEx failed (error {error}).");
        }

        try
        {
            IntPtr pathAddress = remoteBase + pathOffset;
            IntPtr resultAddress = remoteBase + resultOffset;

            if (_memory.WriteBytes(pathAddress, pathBytes) != pathBytes.Length)
                throw new InvalidOperationException("Failed to write the DLL path into the target process.");

            byte[] stub = _currentIs64Bit
                ? BuildHijackStubX64(pathAddress, resultAddress, loadLibrary)
                : BuildHijackStubX86(pathAddress, resultAddress, loadLibrary);

            if (_memory.WriteBytes(remoteBase, stub) != stub.Length)
                throw new InvalidOperationException("Failed to write the hijack stub.");

            nint stackTop = (nint)(remoteBase + totalSize);
            stackTop = AlignDown(stackTop, pointerSize);
            if (_currentIs64Bit)
                stackTop = AlignDown(stackTop, 16);

            using SafeMemoryHandle thread = NativeMethods.OpenThread(
                ThreadAccess.SuspendResume | ThreadAccess.GetContext | ThreadAccess.SetContext | ThreadAccess.QueryInformation,
                false,
                threadId);

            if (thread.IsInvalid)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException($"OpenThread failed (error {error}).");
            }

            if (NativeMethods.SuspendThread(thread) == 0xFFFFFFFF)
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException($"SuspendThread failed (error {error}).");
            }

            bool restored = false;

            try
            {
                IntPtr result;

                if (_currentIs64Bit)
                {
                    var originalContext = new Context64 { ContextFlags = ContextFlags.Amd64Full };
                    if (!NativeMethods.GetThreadContext(thread, ref originalContext))
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"GetThreadContext failed (error {error}).");
                    }

                    var redirectContext = originalContext;
                    redirectContext.Rip = (ulong)(nint)remoteBase;
                    redirectContext.Rsp = (ulong)stackTop;

                    if (!NativeMethods.SetThreadContext(thread, ref redirectContext))
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"SetThreadContext failed (error {error}).");
                    }

                    if (NativeMethods.ResumeThread(thread) == 0xFFFFFFFF)
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"ResumeThread failed (error {error}).");
                    }

                    result = WaitForResult(resultAddress, TimeSpan.FromSeconds(5));

                    if (NativeMethods.SuspendThread(thread) == 0xFFFFFFFF)
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"SuspendThread failed while capturing result (error {error}).");
                    }

                    if (!NativeMethods.SetThreadContext(thread, ref originalContext))
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"SetThreadContext restore failed (error {error}).");
                    }

                    if (NativeMethods.ResumeThread(thread) == 0xFFFFFFFF)
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"ResumeThread restore failed (error {error}).");
                    }

                    restored = true;
                }
                else
                {
                    // 32-bit process targeting a 32-bit process. The target context is
                    // a native x86 CONTEXT; the WOW64 APIs are for 64-bit callers only.
                    var originalContext = new Context32 { ContextFlags = ContextFlags.X86Full };
                    if (!NativeMethods.GetThreadContext(thread, ref originalContext))
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"GetThreadContext failed (error {error}).");
                    }

                    var redirectContext = originalContext;
                    redirectContext.Eip = (uint)(nint)remoteBase;
                    redirectContext.Esp = (uint)(nint)stackTop;

                    if (!NativeMethods.SetThreadContext(thread, ref redirectContext))
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"SetThreadContext failed (error {error}).");
                    }

                    if (NativeMethods.ResumeThread(thread) == 0xFFFFFFFF)
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"ResumeThread failed (error {error}).");
                    }

                    result = WaitForResult(resultAddress, TimeSpan.FromSeconds(5));

                    if (NativeMethods.SuspendThread(thread) == 0xFFFFFFFF)
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"SuspendThread failed while capturing result (error {error}).");
                    }

                    if (!NativeMethods.SetThreadContext(thread, ref originalContext))
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"SetThreadContext restore failed (error {error}).");
                    }

                    if (NativeMethods.ResumeThread(thread) == 0xFFFFFFFF)
                    {
                        int error = Marshal.GetLastPInvokeError();
                        throw new InvalidOperationException($"ResumeThread restore failed (error {error}).");
                    }

                    restored = true;
                }

                if (result == IntPtr.Zero)
                    throw new InvalidOperationException("LoadLibrary returned zero; the DLL could not be loaded.");

                return result;
            }
            catch
            {
                // If we never successfully restored the thread's original context, the
                // thread may still be executing (or about to execute) code inside the
                // injected allocation. Freeing that memory now would crash the target
                // process, so leak the block and leave the thread suspended.
                if (!restored)
                {
                    remoteBase = IntPtr.Zero;
                }

                throw;
            }
        }
        finally
        {
            NativeMethods.VirtualFreeEx(_memory.Handle, remoteBase, 0, MemoryFreeType.Release);
        }
    }

    private void ValidateAndCheckBitness(string dllPath)
    {
        if (string.IsNullOrWhiteSpace(dllPath))
            throw new ArgumentException("DLL path cannot be null or empty.", nameof(dllPath));

        if (!File.Exists(dllPath))
            throw new FileNotFoundException("The specified DLL was not found.", dllPath);

        if (_memory.Is64Bit != _currentIs64Bit)
        {
            throw new InvalidOperationException(
                "The target process bitness does not match the current process bitness.");
        }
    }

    private static IntPtr ResolveLoadLibraryW()
    {
        // kernel32.dll is loaded at the same base address in every process at a given
        // bitness, so resolving the export in the current process gives the correct
        // target address for the remote process.
        IntPtr kernel32 = NativeMethods.LoadLibrary("kernel32.dll");
        if (kernel32 == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"Unable to obtain a handle to kernel32.dll (error {error}).");
        }

        IntPtr loadLibrary = NativeMethods.GetProcAddress(kernel32, "LoadLibraryW");
        if (loadLibrary == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"Unable to resolve LoadLibraryW (error {error}).");
        }

        return loadLibrary;
    }

    private IntPtr WaitForResult(IntPtr resultAddress, TimeSpan timeout)
    {
        Stopwatch watch = Stopwatch.StartNew();
        while (watch.Elapsed < timeout)
        {
            IntPtr value = _memory.Read<IntPtr>(resultAddress);
            if (value != IntPtr.Zero)
                return value;

            System.Threading.Thread.Sleep(5);
        }

        return IntPtr.Zero;
    }

    private static byte[] BuildRemoteThreadStubX86(IntPtr pathAddress, IntPtr resultAddress, IntPtr loadLibrary)
    {
        var buffer = new List<byte>(18);

        // push pathAddress
        buffer.Add(0x68);
        EmitU32(buffer, (uint)(nint)pathAddress);

        // mov ecx, LoadLibraryW
        buffer.Add(0xB9);
        EmitU32(buffer, (uint)(nint)loadLibrary);

        // call ecx
        buffer.Add(0xFF);
        buffer.Add(0xD1);

        // mov [resultAddress], eax
        buffer.Add(0xA3);
        EmitU32(buffer, (uint)(nint)resultAddress);

        // ret
        buffer.Add(0xC3);

        return buffer.ToArray();
    }

    private static byte[] BuildRemoteThreadStubX64(IntPtr pathAddress, IntPtr resultAddress, IntPtr loadLibrary)
    {
        var buffer = new List<byte>(39);

        // mov rcx, pathAddress
        buffer.Add(0x48);
        buffer.Add(0xB9);
        EmitU64(buffer, (ulong)(nint)pathAddress);

        // mov rax, LoadLibraryW
        buffer.Add(0x48);
        buffer.Add(0xB8);
        EmitU64(buffer, (ulong)(nint)loadLibrary);

        // call rax
        buffer.Add(0xFF);
        buffer.Add(0xD0);

        // mov rdx, rax
        buffer.Add(0x48);
        buffer.Add(0x89);
        buffer.Add(0xC2);

        // mov rax, resultAddress
        buffer.Add(0x48);
        buffer.Add(0xB8);
        EmitU64(buffer, (ulong)(nint)resultAddress);

        // mov [rax], rdx
        buffer.Add(0x48);
        buffer.Add(0x89);
        buffer.Add(0x10);

        // ret
        buffer.Add(0xC3);

        return buffer.ToArray();
    }

    private static byte[] BuildHijackStubX86(IntPtr pathAddress, IntPtr resultAddress, IntPtr loadLibrary)
    {
        var buffer = new List<byte>(19);

        // push pathAddress
        buffer.Add(0x68);
        EmitU32(buffer, (uint)(nint)pathAddress);

        // mov ecx, LoadLibraryW
        buffer.Add(0xB9);
        EmitU32(buffer, (uint)(nint)loadLibrary);

        // call ecx
        buffer.Add(0xFF);
        buffer.Add(0xD1);

        // mov [resultAddress], eax
        buffer.Add(0xA3);
        EmitU32(buffer, (uint)(nint)resultAddress);

        // jmp $ (infinite loop so the main injector can suspend and restore context)
        buffer.Add(0xEB);
        buffer.Add(0xFE);

        return buffer.ToArray();
    }

    private static byte[] BuildHijackStubX64(IntPtr pathAddress, IntPtr resultAddress, IntPtr loadLibrary)
    {
        var buffer = new List<byte>(40);

        // mov rcx, pathAddress
        buffer.Add(0x48);
        buffer.Add(0xB9);
        EmitU64(buffer, (ulong)(nint)pathAddress);

        // mov rax, LoadLibraryW
        buffer.Add(0x48);
        buffer.Add(0xB8);
        EmitU64(buffer, (ulong)(nint)loadLibrary);

        // call rax
        buffer.Add(0xFF);
        buffer.Add(0xD0);

        // mov rdx, rax
        buffer.Add(0x48);
        buffer.Add(0x89);
        buffer.Add(0xC2);

        // mov rax, resultAddress
        buffer.Add(0x48);
        buffer.Add(0xB8);
        EmitU64(buffer, (ulong)(nint)resultAddress);

        // mov [rax], rdx
        buffer.Add(0x48);
        buffer.Add(0x89);
        buffer.Add(0x10);

        // jmp $ (infinite loop)
        buffer.Add(0xEB);
        buffer.Add(0xFE);

        return buffer.ToArray();
    }

    private static void EmitU32(List<byte> buffer, uint value)
    {
        buffer.Add((byte)value);
        buffer.Add((byte)(value >> 8));
        buffer.Add((byte)(value >> 16));
        buffer.Add((byte)(value >> 24));
    }

    private static void EmitU64(List<byte> buffer, ulong value)
    {
        EmitU32(buffer, (uint)value);
        EmitU32(buffer, (uint)(value >> 32));
    }

    private static int Align(int value, int alignment)
    {
        return (value + alignment - 1) / alignment * alignment;
    }

    private static nint AlignDown(nint value, int alignment)
    {
        return (nint)((nuint)value & ~((nuint)alignment - 1));
    }
}