/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: AuProcess.cpp
    Date: 2021-8-20
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Process.hpp"

#if defined(AURORA_IS_POSIX_DERIVED)
    #include <stdlib.h>
    #include <dlfcn.h>
    #include <sys/types.h>
    #include <sys/stat.h>
    #include <signal.h>
#endif

#if defined(AURORA_PLATFORM_WIN32)
    #include <Softpub.h>
    #include <wincrypt.h>
    #include <wintrust.h>
#endif

#if defined(AURORA_IS_MODERNNT_DERIVED)
    #include "AuProcessSectionView.NT.hpp"
#endif

#include <Source/IO/FS/FS.hpp>
#include <Source/IO/FS/Resources.hpp>

#include "AuProcessMap.hpp"

#if defined(AURORA_COMPILER_CLANG)
    // warning: enumeration values 'kEnumCount' and 'kEnumInvalid' not handled in switch [-Wswitch
    #pragma clang diagnostic ignored "-Wswitch"
    // Yea, I don't give a shit.
#endif

namespace Aurora::Process
{
    static AuFutexMutex gSpinLock;
    static AuList<AuString> gClassPath;
    static AuHashMap<AuString, void *> gModuleHandles;
    static const bool kIsMainSigned = false;

    void LoadProcessSectionViewSymbol();

    static constexpr const char *GetPlatformString(Build::EPlatform platform)
    {
        switch (platform)
        {
        case Build::EPlatform::ePlatformWin32:
            return ".Win32";
        case Build::EPlatform::ePlatformLinux:
            return ".Linux";
        case Build::EPlatform::ePlatformAndroid:
            return ".Android";
        case Build::EPlatform::ePlatformAppleMacOS:
            return ".Mac";
        case Build::EPlatform::ePlatformIos:
            return ".iOS";
        case Build::EPlatform::eKernelBsd:
            return ".BSD";
        default:
            return nullptr;
        }
    }

    static constexpr const char *GetPlatformExt(Build::EPlatform platform)
    {
        switch (platform)
        {
        case Build::EPlatform::ePlatformWin32:
            return ".dll";
        case Build::EPlatform::ePlatformLinux:
        case Build::EPlatform::ePlatformAndroid:
        case Build::EPlatform::eKernelBsd:
            return ".so";
        case Build::EPlatform::ePlatformAppleMacOS:
        case Build::EPlatform::ePlatformIos:
            return ".dylib";
        default:
            return nullptr;
        }
    }
    
    static constexpr const char *GetArchString(Build::EArchitecture architecture)
    {
        switch (architecture)
        {
        case Build::EArchitecture::eX86_32:
            return ".x86_32";
        case Build::EArchitecture::eX86_64:
            return ".x86_64";
        case Build::EArchitecture::eAArch64:
            return ".arm";
        default:
            return nullptr;
        }
    }

    static AuString ConstructAuDllSuffixUncached()
    {
        auto platform = GetPlatformString(Build::kCurrentPlatform);
        auto architecture = GetArchString(Build::kCurrentArchitecture);
        auto ext = GetPlatformExt(Build::kCurrentPlatform);

        AuString ret;

    #if defined(DEBUG)
        ret += ".Debug";
    #elif defined(STAGING)
        ret += ".Stage";
    #elif defined(SHIP)
        ret += ".Ship";
    #endif
        
        if (platform)
        {
            ret += platform;
        }
        
        if (architecture)
        {
            ret += architecture;
        }

        if (ext)
        {
            ret += ext;
        }

        return ret;
    }

    AUKN_SYM const AuString& ConstructAuDllSuffix()
    {
        static AuString dllSuffixString {};
        static AuInitOnce gInitOnce;
        gInitOnce.Call([]
        {
            dllSuffixString = ConstructAuDllSuffixUncached();
        });
        return dllSuffixString;
    }

    static AuString GetModuleNameFromFileName(const AuString &filename)
    {
        static const auto &kStringSuffixA = ConstructAuDllSuffix();
        static const auto kStringSuffixB = GetPlatformExt(Build::kCurrentPlatform);

        if (filename.ends_with(kStringSuffixA))
        {
            return filename.substr(0, filename.size() - kStringSuffixA.size());
        }

        if (filename.ends_with(kStringSuffixB))
        {
            return filename.substr(0, filename.size() - strlen(kStringSuffixB));
        }

        return filename;
    }

    #if defined(AURORA_PLATFORM_WIN32)
    bool VerifyEmbeddedSignature(const wchar_t *path, HANDLE handle)
    {
        LONG lStatus;
        DWORD dwLastError;

        WINTRUST_FILE_INFO FileData;
        memset(&FileData, 0, sizeof(FileData));
        FileData.cbStruct = sizeof(WINTRUST_FILE_INFO);
        FileData.pcwszFilePath = path;
        FileData.hFile = handle;
        FileData.pgKnownSubject = NULL;

        GUID WVTPolicyGUID = WINTRUST_ACTION_GENERIC_VERIFY_V2;
        WINTRUST_DATA WinTrustData;

        // Initialize the WinVerifyTrust input data structure.

        // Default all fields to 0.
        memset(&WinTrustData, 0, sizeof(WinTrustData));

        WinTrustData.cbStruct = sizeof(WinTrustData);
    
        // Use default code signing EKU.
        WinTrustData.pPolicyCallbackData = NULL;

        // No data to pass to SIP.
        WinTrustData.pSIPClientData = NULL;

        // Disable WVT UI.
        WinTrustData.dwUIChoice = WTD_UI_NONE;

        // Revocation checking.constexpr
        WinTrustData.fdwRevocationChecks = WTD_REVOKE_WHOLECHAIN; 

        // Verify an embedded signature on a file.
        WinTrustData.dwUnionChoice = WTD_CHOICE_FILE;

        // Verify action.
        WinTrustData.dwStateAction = WTD_STATEACTION_VERIFY;

        // Verification sets this value.
        WinTrustData.hWVTStateData = NULL;

        // Not used.
        WinTrustData.pwszURLReference = NULL;

        // This is not applicable if there is no UI because it changes 
        // the UI to accommodate running applications instead of 
        // installing applications.
        WinTrustData.dwUIContext = 0;

        // Set pFile.
        WinTrustData.pFile = &FileData;

        if (!pWinVerifyTrust)
        {
            return gRuntimeConfig.debug.bWin32VerifyTrustFailMissingAPI;
        }

        // WinVerifyTrust verifies signatures as specified by the GUID 
        // and Wintrust_Data.
        lStatus = pWinVerifyTrust(
            NULL,
            &WVTPolicyGUID,
            &WinTrustData);

        bool status {};

        switch (lStatus) 
        {
            case ERROR_SUCCESS:
                status = true;
                break;
            case TRUST_E_NOSIGNATURE:
                // The file was not signed or had a signature 
                // that was not valid.

                // Get the reason for no signature.
                dwLastError = GetLastError();
                if (TRUST_E_NOSIGNATURE == dwLastError ||
                        TRUST_E_SUBJECT_FORM_UNKNOWN == dwLastError ||
                        TRUST_E_PROVIDER_UNKNOWN == dwLastError) 
                {
                    SysPushErrorCrypt("The file is not signed.");
                } 
                else 
                {
                    SysPushErrorCrypt("An unknown error occurred trying to verify the signature of the file.");
                }

                break;
            case TRUST_E_EXPLICIT_DISTRUST:
                SysPushErrorCrypt("The signature is present, but specifically disallowed.");

                break;
            case TRUST_E_SUBJECT_NOT_TRUSTED:
                SysPushErrorCrypt("The signature is present, but not trusted.");

                break;
            case CRYPT_E_SECURITY_SETTINGS:
                SysPushErrorCrypt("CRYPT_E_SECURITY_SETTINGS - The hash "
                    "representing the subject or the publisher wasn't "
                    "explicitly trusted by the admin and admin policy "
                    "has disabled user trust. No signature, publisher "
                    "or timestamp errors.");

                break;
            default:
                SysPushErrorCrypt("Dependency Injection Error");
                break;
        }

        WinTrustData.dwStateAction = WTD_STATEACTION_CLOSE;

        if (!pWinVerifyTrust)
        {
            return gRuntimeConfig.debug.bWin32VerifyTrustFailMissingAPI;
        }

        lStatus = pWinVerifyTrust(
            NULL,
            &WVTPolicyGUID,
            &WinTrustData);

        return status;
    }
    #endif

    static void *LoadModule(const AuString &name, const AuString &path)
    {
        auto itr = gModuleHandles.find(path);
        if (itr != gModuleHandles.end())
        {
            return itr->second;
        }

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        if (!pLoadLibraryW)
        {
            return {};
        }

        auto handle = pLoadLibraryW(Locale::ConvertFromUTF8(path).c_str());
        if (handle == INVALID_HANDLE_VALUE)
        {
            SysPushErrorNested("Could't link dynamic library {}", path);
            return {};
        }
    #else
        auto handle = dlopen(path.c_str(), RTLD_DEEPBIND);
        if (handle == nullptr)
        {
            SysPushErrorNested("Could't link dynamic library {}", path);
            return {};
        }
    #endif

        auto pRet = (void *)handle;
        gModuleHandles.insert(AuMakePair(name, pRet));
        gModuleHandles.insert(AuMakePair(path, pRet));
        return pRet;
    }

    static void *TryLoadModule(const AuString &path,
    #if defined(AURORA_IS_MODERNNT_DERIVED)
                              const AuString &abs,
    #endif
                              const ModuleLoadRequest &request,
                              bool &fail
    )
    {
        fail = false;
        
        auto pathNrml = AuIOFS::NormalizePathRet(path);

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        auto widePath = Locale::ConvertFromUTF8(pathNrml);
        auto mitigateTimeOfUse = Win32Open(widePath.c_str(), GENERIC_READ, FILE_SHARE_READ, false, OPEN_EXISTING, 0, 0);

        if (mitigateTimeOfUse == INVALID_HANDLE_VALUE)
        {
            SysPushErrorNested("Couldn't open existing file. Race exploit?");
            fail = true;
            return {};
        }

        auto absPath = AuIOFS::NormalizePathRet(abs);
    #endif

        if (request.verify || ((kIsMainSigned || request.forceMitigations) && request.enableMitigations))
        {
        #if defined(AURORA_PLATFORM_WIN32)
            if (!VerifyEmbeddedSignature(widePath.c_str(), mitigateTimeOfUse))
            {
                SysPushErrorNested("Couldn't verify file {}", path);
                fail = true;
                AuWin32CloseHandle(mitigateTimeOfUse);
                return {};
            }
        #else
            AuLogWarn("Can't verify {} on this platform", path);
        #endif
        }

    #if defined(AURORA_IS_POSIX_DERIVED)
        if (request.enableMitigations && request.unixCheckPlusX)
        {
            struct stat sb;

            if (stat(pathNrml.data(), &sb) != 0)
            {
                SysPushErrorNested("Couldn't open existing file. Race exploit?");
                fail = true;
                return {};
            }

            if ((sb.st_mode & S_IXUSR) == 0)
            {
                fail = true;
                return {};
            }
        }
    #endif

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        auto returnValue = LoadModule(request.mod, absPath);
    #else
        auto returnValue = LoadModule(request.mod, pathNrml);
    #endif
        
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        // Still doesnt stop hotswaps, and idc
        // Glitching would require software to time the attack
        // ...requiring a certain degree of access (code execution) + effort   
        AuWin32CloseHandle(mitigateTimeOfUse);
    #endif

        return returnValue;
    }

    static void *TryLoadModule(const AuString &path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request, bool &fail)
    {
        AuString a = path + AuString({ AuFS::kPathSplitter }) + auDll;
        AuString b = path + AuString({ AuFS::kPathSplitter }) + genericDll;

    #if defined(AURORA_IS_MODERNNT_DERIVED)
    #if defined(AURORA_PLATFORM_WIN32)
        if (gRuntimeConfig.win32Config.bProcessCheckWinLdrForModNameFirst)
        {
            auto aW = AuLocale::ConvertFromUTF8(genericDll);
            auto bW = AuLocale::ConvertFromUTF8(auDll);

            auto pDLL = ::GetModuleHandleW(aW.c_str());
            if (!pDLL)
            {
                pDLL = ::GetModuleHandleW(bW.c_str());
            }

            if (pDLL)
            {
                gModuleHandles.insert(AuMakePair(request.mod, (void *)pDLL));
                return pDLL;
            }
        }
    #endif

        AuString aAbs = path + AuString({ AuFS::kPathSplitter }) + auDll + ".";
        AuString bAbs = path + AuString({ AuFS::kPathSplitter }) + genericDll + ".";
    #endif
        
        if (AuIOFS::FileExists(a))
        {
            return TryLoadModule(a,
        #if defined(AURORA_IS_MODERNNT_DERIVED)
                                 aAbs,
        #endif
                                 request,
                                 fail
            );
        }

        if (AuIOFS::FileExists(b))
        {
            return TryLoadModule(b,
        #if defined(AURORA_IS_MODERNNT_DERIVED)
                                 bAbs,
        #endif
                                 request,
                                 fail
            );
        }

        return {};
    }

    static void *TryLoadModule(EModulePath path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request)
    {
        AuString pathA, pathB;
        AuList<AuString> arrayPaths;

        switch (path)
        {
        case EModulePath::eClassPath:
        {
            AU_LOCK_GUARD(gSpinLock);
            arrayPaths = gClassPath;
            break;
        }
        case EModulePath::eModulePathCWD:
        {
            if (!Process::GetWorkingDirectory(pathA))
            {
                return {};
            }

            break;
        }
        case EModulePath::eProcessDirectory:
        {
            if (!Process::GetProcDirectory(pathA))
            {
                return {};
            }

            break;
        }
        case EModulePath::eModulePathSystemDir:
        {
            pathA = AuIOFS::GetSystemLibPath().value_or(AuString {});
            pathB = AuIOFS::GetSystemLibPath2().value_or(AuString {});

            if (pathA.empty())
            {
                return {};
            }

            break;
        }
        case EModulePath::eModulePathUserDir:
        {
            pathA = AuIOFS::GetUserLibPath().value_or(AuString {});
            pathB = AuIOFS::GetUserLibPath2().value_or(AuString {});

            if (pathA.empty())
            {
                return {};
            }

            break;
        }
        //case EModulePath::eOSSpecified:
        //
        //    break;
        case EModulePath::eSpecified:
            
            break;
        }
        
        bool fail {};

        if (pathA.size())
        {
            if (auto pRet = TryLoadModule(pathA, auDll, genericDll, request, fail))
            {
                return pRet;
            }
        
            if (fail)
            {
                return {};
            }
        }

        if (pathB.size())
        {
            if (auto pRet = TryLoadModule(pathB, auDll, genericDll, request, fail))
            {
                return pRet;
            }

            if (fail)
            {
                return {};
            }
        }

        for (const auto &dir : arrayPaths)
        {
            if (auto pRet = TryLoadModule(dir, auDll, genericDll, request, fail))
            {
                return pRet;
            }

            if (fail)
            {
                return {};
            }
        }

        if (path == EModulePath::eSpecified && request.specifiedSearchPaths)
        {
            for (const auto &val : *request.specifiedSearchPaths)
            {
                if (auto pRet = TryLoadModule(val, auDll, genericDll, request, fail))
                {
                    return pRet;
                }
                
                if (fail)
                {
                    return {};
                }
            }
        }

        return {};
    }

    AUKN_SYM bool LoadModule(const ModuleLoadRequest &request)
    {
        return bool(LoadModuleEx(request));
    }

    AUKN_SYM void *LoadModuleEx(const ModuleLoadRequest &request)
    {
        AU_LOCK_GUARD(gSpinLock);
        
        auto h = gModuleHandles.find(request.mod);
        if (h != gModuleHandles.end())
        {
            return h->second;
        }

        auto au   = request.mod + ConstructAuDllSuffix();
        auto base = request.mod;
        auto ext  = GetPlatformExt(Build::kCurrentPlatform);

        if (ext && !Build::kIsNtDerived)
        {
            base += ext;
        }

        if (request.version.size())
        {
            au += "." + request.version;
            base += "." + request.version;
        }

        if (ext && Build::kIsNtDerived)
        {
            base += ext;
        }
        
        auto searchPath = request.searchPath ? request.searchPath : &kUserOverloadableSearchPath;
        for (EModulePath path : *searchPath)
        {
            if (auto pRet = TryLoadModule(path, au, base, request))
            {
                return pRet;
            }
        }

        return {};
    }

    AUKN_SYM AuMach GetProcAddressEx(void *pHandle, const AuString &symbol)
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        if (!pGetProcAddress)
        {
            return {};
        }

        auto ret = reinterpret_cast<AuMach>(pGetProcAddress(reinterpret_cast<HMODULE>(pHandle), symbol.c_str()));
    #else
        auto ret = reinterpret_cast<AuMach>(dlsym(hapHandlendle, symbol.c_str()));
    #endif
        return ret;
    }

    AUKN_SYM void *GetProcHandle(const AuString &name)
    {
        AU_LOCK_GUARD(gSpinLock);

        auto h = gModuleHandles.find(name);
        if (h == gModuleHandles.end())
        {
            SysPushErrorGen("Module {} is not loaded", name);
            return {};
        }

        return h->second;
    }

    AUKN_SYM AuMach GetProcAddress(const AuString &mod, const AuString &symbol)
    {
        auto ret = GetProcAddressEx(GetProcHandle(mod), symbol);

        if (!ret)
        {
            SysPushErrorGen("Couldn't resolve symbol in module {} of mangled name '{}'", mod, symbol);
        }

        return ret;
    }

    AUKN_SYM void Exit(AuUInt32 exitcode)
    {
        Aurora::RuntimeShutdown();

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        TerminateProcess(GetCurrentProcess(), exitcode);
        Win32Terminate();
    #elif defined(AURORA_IS_POSIX_DERIVED)
        // TODO: if main thread, long jump back, and return exit code
        ::kill(getpgrp(), SIGKILL);
        while (true)
        {
            ::sched_yield();
        }
    #else
        // ???
        *(AuUInt32 *)0 = exitcode;
        *(AuUInt32 *)0xFFFF = exitcode;
        *(AuUInt32 *)0xFFFFFFFF = exitcode;
    #endif
    }

    static void PreloadDLLsDoOnce(const AuList<AuString> &dirs)
    {
        AU_DEBUG_MEMCRUNCH;

        AuList<void *> cookieList;

    #if defined(AURORA_PLATFORM_WIN32)
        for (const auto dir : dirs)
        {
            void *pCookie {};

            if (pRemoveDllDirectory)
            {
                if (pAddDllDirectory)
                {
                    pCookie = pAddDllDirectory(AuLocale::ConvertFromUTF8(dir).c_str());
                }
            }

            cookieList.push_back(pCookie);
        }
    #endif

        for (const auto dir : dirs)
        {
            AuList<AuString> files;
            AuList<AuString> sharedObjects;

            if (!AuFS::FilesInDirectory(dir, files))
            {
                SysPushErrorIO("Couldn't readdir: {}", dir);
                continue;
            }

            auto endingPattern = ConstructAuDllSuffix();
            auto endingAlt = GetPlatformExt(Build::kCurrentPlatform);

            for (const auto &file : files)
            {
                if (!(file.ends_with(endingPattern.c_str()) ||
                      file.ends_with(endingAlt)))
                {
                    continue;
                }

                sharedObjects.push_back(file);
            }

            {
                AuUInt uSuccess {};
                do
                {
                    uSuccess = 0;

                    for (auto itr = sharedObjects.begin();
                         itr != sharedObjects.end();  )
                    {
                        auto &sharedDLL = *itr;

                        AuString a = dir + AuString({ AuFS::kPathSplitter }) + sharedDLL;

                    #if defined(AURORA_IS_MODERNNT_DERIVED)
                        AuString aAbs = dir + AuString({ AuFS::kPathSplitter }) + sharedDLL + ".";
                    #endif

                    #if defined(AURORA_IS_MODERNNT_DERIVED)
                        if (GetProcHandle(aAbs))
                    #else
                        if (GetProcHandle(a))
                    #endif
                        {
                            itr++;
                            continue;
                        }

                        auto modName = GetModuleNameFromFileName(sharedDLL);
                        AuLogDbg("[attempt] Loading shared object: {} (path: {})", modName, sharedDLL);

                        bool fail {};
                        ModuleLoadRequest request { AuList<AuString>{ dir }, modName };

                        if (TryLoadModule(a,
                        #if defined(AURORA_IS_MODERNNT_DERIVED)
                                          aAbs,
                        #endif
                                          request,
                                          fail
                            ))
                        {
                            if (!fail)
                            {
                                AuLogDbg("[attempt] Loaded shared object: {} (path: {})", modName, sharedDLL);
                                uSuccess++;
                                itr = sharedObjects.erase(itr);
                            }
                            else
                            {
                                itr++;
                            }
                        }
                        else
                        {
                            itr++;
                        }
                    }
                }
                while (uSuccess);
            }
        }

    #if defined(AURORA_PLATFORM_WIN32)
        for (const auto pCookie : AuExchange(cookieList, {}))
        {
            if (pRemoveDllDirectory)
            {
                if (pCookie)
                {
                    pRemoveDllDirectory(pCookie);
                }
            }
        }
    #endif
    }

    static void PreloadDLLs(const AuString &dir)
    {
        AU_DEBUG_MEMCRUNCH;
        if (gRuntimeConfig.processConfig.bAlwaysPreloadEntireClassPath)
        {
            PreloadDLLsDoOnce(gClassPath);
        }
        else
        {
            PreloadDLLsDoOnce({ dir });
        }
    }

    AUKN_SYM bool SetBinaryClassPath(const AuList<AuString> &list, bool preloadAll)
    {
        AU_DEBUG_MEMCRUNCH;
        AU_LOCK_GUARD(gSpinLock);
        AuExchange(gClassPath, list);

        if (((preloadAll && gRuntimeConfig.processConfig.bEnablePreload)) ||
             (gRuntimeConfig.processConfig.bForcePreload))
        {
            if (gRuntimeConfig.processConfig.bAlwaysPreloadEntireClassPath)
            {
                PreloadDLLsDoOnce(gClassPath);
            }
            else
            {
                PreloadDLLsDoOnce(list);
            }
        }

        return true;
    }

    AUKN_SYM bool AddBinaryClassPath(const AuString &dir, bool preloadAll)
    {
        AU_LOCK_GUARD(gSpinLock);

        if (!AuTryInsert(gClassPath, dir))
        {
            return false;
        }

        if (((preloadAll && gRuntimeConfig.processConfig.bEnablePreload)) ||
             (gRuntimeConfig.processConfig.bForcePreload))
        {
            PreloadDLLs(dir);
        }

        return true;
    }

    AUKN_SYM AuList<AuString> GetBinaryClassPath()
    {
        AU_LOCK_GUARD(gSpinLock);
        return gClassPath;
    }

    void InitProcess()
    {
        // TODO (reece): test if self is signed -> xref kIsMainSigned

        InitProcessMap();

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        LoadProcessSectionViewSymbol();
    #endif
    }

    void DeinitProcess()
    {
        DeinitProcessMap();
    }
}