/*
 * beacon_compatibility.c -- FULL Beacon API compatibility layer
 * Step 07: Complete implementation of every Beacon API function.
 *
 * Previous steps used stubs. This file provides real, working
 * implementations so that BOFs can parse arguments, produce output,
 * build format buffers, and use basic token/process helpers.
 *
 * The output system accumulates all BeaconPrintf / BeaconOutput
 * text into a single buffer that the loader retrieves after the
 * BOF's entry point returns.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <windows.h>
#include "beacon_compatibility.h"

/* ================================================================
 * Global output buffer state
 *
 * Every BeaconPrintf / BeaconOutput call appends to this buffer.
 * After the BOF returns, the loader calls BeaconGetOutputData()
 * to harvest the accumulated text.
 * ================================================================ */
char*  beacon_compatibility_output = NULL;   /* accumulated output        */
int    beacon_compatibility_size   = 0;      /* allocated capacity        */
int    beacon_compatibility_offset = 0;      /* bytes written so far      */

/* ================================================================
 * InternalFunctions lookup table
 *
 * 30 slots, each holding { (unsigned char*)"FuncName", (unsigned char*)funcPtr }.
 * The loader populates this at startup; process_symbol() searches
 * it when resolving __imp_BeaconXxx symbols.
 * ================================================================ */
unsigned char* InternalFunctions[30][2] = { {0} };


/* ================================================================
 *                    UTILITY HELPERS
 * ================================================================ */

/*
 * swap_endianess -- byte-swap a 32-bit integer to big-endian.
 *
 * Cobalt Strike's format buffers store integers in network byte
 * order (big-endian).  This swaps the four bytes of a native
 * little-endian int.
 */
int swap_endianess(int value) {
    return ((value >> 24) & 0x000000FF) |
           ((value >>  8) & 0x0000FF00) |
           ((value <<  8) & 0x00FF0000) |
           ((value << 24) & 0xFF000000);
}

/*
 * toWideChar -- convert a narrow (ANSI/UTF-8) string to UTF-16.
 *
 * Wraps MultiByteToWideChar with CP_UTF8.
 * Returns TRUE on success, FALSE on failure.
 */
BOOL toWideChar(char* src, wchar_t* dst, int max) {
    /* Convert from UTF-8 to UTF-16LE */
    int result = MultiByteToWideChar(CP_UTF8, 0, src, -1, dst, max);
    if (result == 0) {
        return FALSE;
    }
    return TRUE;
}


/* ================================================================
 *                    DATA PARSING API
 *
 * These functions implement a cursor-based reader over a packed
 * argument buffer.  The buffer format is:
 *   [4-byte total_length] [payload bytes...]
 * BeaconDataParse skips the 4-byte length prefix so subsequent
 * reads start at the payload.
 * ================================================================ */

/*
 * BeaconDataParse -- initialize a datap parser on a raw buffer.
 *
 * Sets original to the full buffer, then advances the read cursor
 * past the 4-byte length prefix.
 */
void BeaconDataParse(datap* parser, char* buffer, int size) {
    if (parser == NULL) {
        return;
    }
    parser->original = buffer;
    parser->buffer   = buffer + 4;   /* skip 4-byte length prefix */
    parser->length   = size - 4;
    parser->size     = size - 4;
}

/*
 * BeaconDataInt -- read a 4-byte integer from the current position.
 *
 * Copies 4 bytes via memcpy (safe for unaligned access), advances
 * the cursor by 4, and decrements the remaining length.
 */
int BeaconDataInt(datap* parser) {
    int value = 0;
    if (parser->length < 4) {
        return 0;
    }
    memcpy(&value, parser->buffer, 4);
    parser->buffer += 4;
    parser->length -= 4;
    return value;
}

/*
 * BeaconDataShort -- read a 2-byte short from the current position.
 *
 * Same pattern as BeaconDataInt but for 16-bit values.
 */
short BeaconDataShort(datap* parser) {
    short value = 0;
    if (parser->length < 2) {
        return 0;
    }
    memcpy(&value, parser->buffer, 2);
    parser->buffer += 2;
    parser->length -= 2;
    return value;
}

/*
 * BeaconDataLength -- return the number of bytes remaining to read.
 */
int BeaconDataLength(datap* parser) {
    return parser->length;
}

/*
 * BeaconDataExtract -- read a length-prefixed blob from the buffer.
 *
 * Reads a 4-byte length, returns a pointer to that many bytes,
 * and advances past them.  If outsize is non-NULL, writes the
 * extracted length into it.
 */
char* BeaconDataExtract(datap* parser, int* outsize) {
    int length = 0;
    char* data = NULL;

    if (parser->length < 4) {
        if (outsize) *outsize = 0;
        return NULL;
    }

    /* Read the 4-byte length prefix */
    memcpy(&length, parser->buffer, 4);
    parser->buffer += 4;
    parser->length -= 4;

    /* Return pointer to the data and advance past it */
    data = parser->buffer;
    if (length > parser->length) {
        length = parser->length;  /* clamp to available */
    }
    parser->buffer += length;
    parser->length -= length;

    if (outsize) {
        *outsize = length;
    }
    return data;
}


/* ================================================================
 *                    OUTPUT API
 *
 * Output goes to two places simultaneously:
 *   1. The console (printf/fwrite) for immediate developer feedback.
 *   2. An internal buffer so the loader can retrieve it later,
 *      mirroring how Cobalt Strike captures BOF output.
 * ================================================================ */

/*
 * BeaconPrintf -- printf-style formatted output from a BOF.
 *
 * Prints to the console via vprintf, then measures the formatted
 * string with vsnprintf, grows the output buffer, and appends it.
 */
void BeaconPrintf(int type, char* fmt, ...) {
    va_list args;
    int needed;
    char* new_buf;

    /* Print to console for immediate feedback */
    va_start(args, fmt);
    vprintf(fmt, args);
    va_end(args);

    /* Measure how many bytes the formatted string needs */
    va_start(args, fmt);
    needed = vsnprintf(NULL, 0, fmt, args);
    va_end(args);

    if (needed < 0) {
        return;
    }

    /* Grow the output buffer to fit the new text */
    new_buf = (char*)realloc(
        beacon_compatibility_output,
        beacon_compatibility_offset + needed + 1
    );
    if (new_buf == NULL) {
        return;
    }
    beacon_compatibility_output = new_buf;
    beacon_compatibility_size   = beacon_compatibility_offset + needed + 1;

    /* Append the formatted text into the buffer */
    va_start(args, fmt);
    vsnprintf(
        beacon_compatibility_output + beacon_compatibility_offset,
        needed + 1,
        fmt, args
    );
    va_end(args);

    beacon_compatibility_offset += needed;
}

/*
 * BeaconOutput -- write raw bytes to the output buffer.
 *
 * Unlike BeaconPrintf, this copies raw data (not format strings).
 * The data is also null-terminated in the buffer for safety.
 */
void BeaconOutput(int type, char* data, int len) {
    char* new_buf;

    /* Grow the output buffer */
    new_buf = (char*)realloc(
        beacon_compatibility_output,
        beacon_compatibility_offset + len + 1
    );
    if (new_buf == NULL) {
        return;
    }
    beacon_compatibility_output = new_buf;
    beacon_compatibility_size   = beacon_compatibility_offset + len + 1;

    /* Copy the raw bytes and null-terminate */
    memcpy(
        beacon_compatibility_output + beacon_compatibility_offset,
        data, len
    );
    beacon_compatibility_offset += len;
    beacon_compatibility_output[beacon_compatibility_offset] = '\0';
}

/*
 * BeaconGetOutputData -- harvest accumulated output after BOF returns.
 *
 * Returns the output buffer pointer and its length via outsize,
 * then resets all globals so the next BOF starts clean.
 */
char* BeaconGetOutputData(int* outsize) {
    char* output = beacon_compatibility_output;

    if (outsize) {
        *outsize = beacon_compatibility_offset;
    }

    /* Reset globals -- caller owns the returned pointer */
    beacon_compatibility_output = NULL;
    beacon_compatibility_size   = 0;
    beacon_compatibility_offset = 0;

    return output;
}


/* ================================================================
 *                    FORMAT BUFFER API
 *
 * Format buffers let BOFs build structured binary data to send
 * back to the team server.  They work like a simple arena:
 * allocate once, append repeatedly, extract at the end.
 * ================================================================ */

/*
 * BeaconFormatAlloc -- allocate a format buffer of the given capacity.
 *
 * The buffer is zero-initialized.  All cursor fields start at zero.
 */
void BeaconFormatAlloc(formatp* format, int maxsz) {
    if (format == NULL) {
        return;
    }
    format->original = (char*)calloc(maxsz, 1);
    format->buffer   = format->original;
    format->length   = 0;
    format->size     = maxsz;
}

/*
 * BeaconFormatReset -- clear the buffer contents without freeing.
 *
 * Zeros the memory and resets the write cursor to the start.
 */
void BeaconFormatReset(formatp* format) {
    if (format == NULL || format->original == NULL) {
        return;
    }
    memset(format->original, 0, format->size);
    format->buffer = format->original;
    format->length = 0;
}

/*
 * BeaconFormatFree -- release the format buffer memory.
 *
 * Frees the allocation and zeros all struct fields.
 */
void BeaconFormatFree(formatp* format) {
    if (format == NULL) {
        return;
    }
    if (format->original) {
        free(format->original);
    }
    format->original = NULL;
    format->buffer   = NULL;
    format->length   = 0;
    format->size     = 0;
}

/*
 * BeaconFormatAppend -- copy raw bytes into the format buffer.
 *
 * Appends len bytes from text at the current cursor position.
 * Does NOT check for overflow -- caller must respect maxsz.
 */
void BeaconFormatAppend(formatp* format, char* text, int len) {
    if (format == NULL || format->original == NULL) {
        return;
    }
    /* Bounds check: do not write past allocated capacity */
    if (format->length + len > format->size) {
        return;
    }
    memcpy(format->original + format->length, text, len);
    format->length += len;
}

/*
 * BeaconFormatPrintf -- printf-style append into a format buffer.
 *
 * Formats the string into the remaining space after the cursor.
 * Advances the cursor by the number of characters written.
 */
void BeaconFormatPrintf(formatp* format, char* fmt, ...) {
    va_list args;
    int remaining;
    int written;

    if (format == NULL || format->original == NULL) {
        return;
    }

    remaining = format->size - format->length;
    if (remaining <= 0) {
        return;
    }

    va_start(args, fmt);
    written = vsnprintf(
        format->original + format->length,
        remaining,
        fmt, args
    );
    va_end(args);

    if (written > 0) {
        /* Clamp to available space */
        if (written > remaining - 1) {
            written = remaining - 1;
        }
        format->length += written;
    }
}

/*
 * BeaconFormatToString -- finalize and return the buffer contents.
 *
 * Returns the pointer to the buffer start and the number of bytes
 * written.  The buffer is still owned by the formatp -- call
 * BeaconFormatFree when done.
 */
char* BeaconFormatToString(formatp* format, int* size) {
    if (size) {
        *size = format->length;
    }
    return format->original;
}

/*
 * BeaconFormatInt -- append a 4-byte integer in big-endian (network) order.
 *
 * Cobalt Strike uses big-endian integers in its format protocol.
 * This swaps the byte order then appends the 4 bytes.
 */
void BeaconFormatInt(formatp* format, int value) {
    int swapped = swap_endianess(value);
    BeaconFormatAppend(format, (char*)&swapped, 4);
}


/* ================================================================
 *                    TOKEN / PROCESS API
 *
 * In a real Cobalt Strike beacon, these manage impersonation tokens
 * and process injection.  In a standalone loader, most are no-ops
 * or minimal stubs since we have no team server to coordinate with.
 * ================================================================ */

/*
 * BeaconUseToken -- impersonate using the given token handle.
 *
 * Stub: calls ImpersonateLoggedOnUser if the handle is valid,
 * but does not integrate with any broader token management.
 */
void BeaconUseToken(HANDLE token) {
    if (token != NULL && token != INVALID_HANDLE_VALUE) {
        ImpersonateLoggedOnUser(token);
    }
}

/*
 * BeaconRevertToken -- drop any impersonation and revert to self.
 *
 * Calls RevertToSelf() to restore the loader's original token.
 */
void BeaconRevertToken(void) {
    RevertToSelf();
}

/*
 * BeaconIsAdmin -- check if the current process has admin privileges.
 *
 * Stub: always returns FALSE in the standalone loader.
 * A real implementation would check the token's group membership.
 */
BOOL BeaconIsAdmin(void) {
    return FALSE;
}

/*
 * BeaconGetSpawnTo -- get the path of the process to spawn for fork&run.
 *
 * Stub: fills the buffer with "rundll32.exe", which is the default
 * spawn-to process in Cobalt Strike.
 */
void BeaconGetSpawnTo(BOOL x86, char* buffer, int length) {
    if (buffer == NULL || length <= 0) {
        return;
    }
    /* Default spawn-to binary */
    strncpy(buffer, "rundll32.exe", length - 1);
    buffer[length - 1] = '\0';
}

/*
 * BeaconSpawnTemporaryProcess -- create a temporary sacrificial process.
 *
 * Stub: returns FALSE. A real implementation would create the process
 * in a suspended state for injection.
 */
BOOL BeaconSpawnTemporaryProcess(BOOL x86, BOOL ignoreToken,
                                 STARTUPINFOA* si, PROCESS_INFORMATION* pi) {
    /* Not implemented in standalone loader */
    return FALSE;
}

/*
 * BeaconInjectProcess -- inject payload into an existing process.
 *
 * Stub: no-op. Would normally allocate memory in the target process,
 * write the payload, and create a remote thread.
 */
void BeaconInjectProcess(HANDLE hProc, int pid,
                          char* payload, int p_len, int offset,
                          char* arg, int a_len) {
    /* Not implemented in standalone loader */
}

/*
 * BeaconInjectTemporaryProcess -- inject into a newly spawned process.
 *
 * Stub: no-op. Works with BeaconSpawnTemporaryProcess to inject
 * into the sacrificial process before resuming it.
 */
void BeaconInjectTemporaryProcess(PROCESS_INFORMATION* pi,
                                   char* payload, int p_len, int offset,
                                   char* arg, int a_len) {
    /* Not implemented in standalone loader */
}

/*
 * BeaconCleanupProcess -- close handles from a spawned process.
 *
 * Closes the process and thread handles stored in the
 * PROCESS_INFORMATION struct to prevent handle leaks.
 */
void BeaconCleanupProcess(PROCESS_INFORMATION* pi) {
    if (pi == NULL) {
        return;
    }
    if (pi->hProcess) {
        CloseHandle(pi->hProcess);
        pi->hProcess = NULL;
    }
    if (pi->hThread) {
        CloseHandle(pi->hThread);
        pi->hThread = NULL;
    }
}