/***
    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"

namespace Aurora::Process
{
    static AuThreadPrimitives::SpinLock gSpinLock;
    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;
    }

    static const AuString& ConstructAuDllSuffix()
    {
        static AuString dllSuffixString {};
        if (dllSuffixString.empty()) dllSuffixString = ConstructAuDllSuffix();
        dllSuffixString = ConstructAuDllSuffixUncached();
        return dllSuffixString;
    }

    #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;

        // WinVerifyTrust verifies signatures as specified by the GUID 
        // and Wintrust_Data.
        lStatus = WinVerifyTrust(
            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;

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

        return status;
    }
    #endif

    static bool LoadModule(const AuString &name, const AuString &path)
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        auto handle = LoadLibraryW(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

        gModuleHandles.insert(AuMakePair(name, (void *)handle));
        return true;
    }

    static bool 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 = CreateFileW(widePath.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

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

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

        if (request.verify || (kIsMainSigned || request.enableMitigations))
        {
        #if defined(AURORA_PLATFORM_WIN32)
            if (!VerifyEmbeddedSignature(widePath.c_str(), mitigateTimeOfUse))
            {
                SysPushErrorNested("Couldn't verify file {}", path);
                fail = true;
                AuWin32CloseHandle(mitigateTimeOfUse);
                return false;
            }
        #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 false;
            }

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

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        auto returnValue = LoadModule(absPath);
    #else
        auto returnValue = LoadModule(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 bool TryLoadModule(const AuString &path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request, bool &fail)
    {
        AuString a = path + "/" + auDll;
        AuString b = path + "/" + genericDll;

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        AuString aAbs = path + "/" + auDll + ".";
        AuString bAbs = path + "/" + 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 bool TryLoadModule(EModulePath path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request)
    {
        AuString pathA, pathB;

        switch (path)
        {
        case EModulePath::eModulePathCWD:
            if (!Process::GetWorkingDirectory(pathA)) return false;

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

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

            if (pathA.empty()) return false;

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

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

        if (pathA.size())
        {
            if (TryLoadModule(pathA, auDll, genericDll, request, fail))
            {
                return true;
            }
        
            if (fail)
            {
                return false;
            }
        }

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

            if (fail)
            {
                return false;
            }
        }

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

        return {};
    }

    AUKN_SYM bool LoadModule(const ModuleLoadRequest &request)
    {
        AU_LOCK_GUARD(gSpinLock);
        
        auto h = gModuleHandles.find(request.mod);
        if (h != gModuleHandles.end())
        {
            return true;
        }

        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 (TryLoadModule(path, au, base, request))
            {
                return true;
            }
        }

        return false;
    }

    AUKN_SYM AuMach GetProcAddress(AuString mod, AuString symbol)
    {
        AU_LOCK_GUARD(gSpinLock);

        auto h = gModuleHandles.find(mod);
        if (h == gModuleHandles.end())
        {
            SysPushErrorGen("Module {} is not loaded", mod);
            return {};
        }
        auto handle = h->second;
        
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        auto ret = reinterpret_cast<AuMach>(::GetProcAddress(reinterpret_cast<HMODULE>(handle), symbol.c_str()));
    #else
        auto ret = reinterpret_cast<AuMach>(dlsym(handle, symbol.c_str()));
    #endif

        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
    }


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

        InitProcessMap();

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

    void DeinitProcess()
    {
        DeinitProcessMap();
    }
}