/*
 * beacon_compatibility.c -- Beacon API compatibility layer implementation
 * Step 08: Full implementation of all Beacon API functions for standalone
 *          BOF execution outside of Cobalt Strike.
 *
 * This file is compiled as part of the COFFLoader (a normal C program
 * with full CRT access). It provides the function implementations that
 * BOFs call through resolved function pointers. The BOF itself has no
 * CRT, but these implementations freely use vprintf, realloc, calloc, etc.
 */

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

/* ================================================================
 * Global output buffer
 *
 * In Cobalt Strike, BeaconPrintf/BeaconOutput send data over the
 * C2 channel. Here we accumulate output in a dynamically growing
 * buffer that the caller retrieves after the BOF finishes.
 * ================================================================ */
char*  beacon_compatibility_output = NULL;
int    beacon_compatibility_size   = 0;
int    beacon_compatibility_offset = 0;

/* ================================================================
 * InternalFunctions table
 *
 * 30-slot array populated by RunCOFF() before relocation processing.
 * process_symbol() searches this table when resolving Beacon API
 * function names. Initialized to all NULLs here.
 * ================================================================ */
unsigned char* InternalFunctions[30][2] = { {NULL} };

/* ================================================================
 * Data Parsing Functions
 *
 * Arguments from Cobalt Strike arrive as a packed buffer with a
 * 4-byte size prefix. These functions provide cursor-based reading.
 * ================================================================ */

void BeaconDataParse(datap* parser, char* buffer, int size) {
    if (parser == NULL) return;

    parser->original = buffer;
    parser->buffer   = buffer + 4;    /* skip 4-byte size prefix */
    parser->length   = size - 4;      /* remaining after prefix  */
    parser->size     = size - 4;
}

int BeaconDataInt(datap* parser) {
    if (parser == NULL || parser->length < 4) return 0;

    int32_t value;
    memcpy(&value, parser->buffer, sizeof(int32_t));

    parser->buffer += 4;
    parser->length -= 4;

    return value;
}

short BeaconDataShort(datap* parser) {
    if (parser == NULL || parser->length < 2) return 0;

    short value;
    memcpy(&value, parser->buffer, sizeof(short));

    parser->buffer += 2;
    parser->length -= 2;

    return value;
}

int BeaconDataLength(datap* parser) {
    if (parser == NULL) return 0;
    return parser->length;
}

char* BeaconDataExtract(datap* parser, int* size) {
    if (parser == NULL || parser->length < 4) {
        if (size) *size = 0;
        return NULL;
    }

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

    /* Return pointer to the data */
    char* data = parser->buffer;
    if (size) *size = length;

    /* Advance past the data */
    parser->buffer += length;
    parser->length -= length;

    return data;
}

/* ================================================================
 * Output Functions
 *
 * BeaconPrintf: formatted output (like printf) + buffer accumulation
 * BeaconOutput: raw byte output + buffer accumulation
 * ================================================================ */

void BeaconPrintf(int type, char* fmt, ...) {
    va_list args;
    va_start(args, fmt);

    /* 1. Print to console (COFFLoader runs as a CLI tool) */
    vprintf(fmt, args);

    va_end(args);
    va_start(args, fmt);

    /* 2. Calculate required buffer size */
    int len = vsnprintf(NULL, 0, fmt, args);
    va_end(args);

    if (len <= 0) return;

    /* 3. Grow the output buffer */
    char* newbuf = (char*)realloc(
        beacon_compatibility_output,
        beacon_compatibility_offset + len + 1
    );
    if (newbuf == NULL) return;
    beacon_compatibility_output = newbuf;

    /* 4. Format the string into the buffer */
    va_start(args, fmt);
    vsnprintf(
        beacon_compatibility_output + beacon_compatibility_offset,
        len + 1,
        fmt,
        args
    );
    va_end(args);

    beacon_compatibility_offset += len;
}

void BeaconOutput(int type, char* data, int len) {
    /* Grow the output buffer */
    char* newbuf = (char*)realloc(
        beacon_compatibility_output,
        beacon_compatibility_offset + len + 1
    );
    if (newbuf == NULL) return;
    beacon_compatibility_output = newbuf;

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

/* ================================================================
 * Format Buffer Functions
 *
 * String-builder pattern: allocate a buffer, append formatted data,
 * retrieve the result, then free.
 * ================================================================ */

void BeaconFormatAlloc(formatp* format, int maxsz) {
    if (format == NULL) return;
    format->original = (char*)calloc(1, maxsz);
    format->buffer   = format->original;
    format->length   = 0;
    format->size     = maxsz;
}

void BeaconFormatReset(formatp* format) {
    if (format == NULL) return;
    memset(format->original, 0, format->size);
    format->buffer = format->original;
    format->length = 0;
}

void BeaconFormatFree(formatp* format) {
    if (format == NULL) return;
    free(format->original);
    format->original = NULL;
    format->buffer   = NULL;
    format->length   = 0;
    format->size     = 0;
}

void BeaconFormatAppend(formatp* format, char* text, int len) {
    if (format == NULL || format->length + len > format->size) return;
    memcpy(format->buffer, text, len);
    format->buffer += len;
    format->length += len;
}

void BeaconFormatPrintf(formatp* format, char* fmt, ...) {
    if (format == NULL) return;

    va_list args;
    va_start(args, fmt);
    int remaining = format->size - format->length;
    int len = vsnprintf(format->buffer, remaining, fmt, args);
    va_end(args);

    if (len > 0 && len < remaining) {
        format->buffer += len;
        format->length += len;
    }
}

char* BeaconFormatToString(formatp* format, int* size) {
    if (size) *size = format->length;
    return format->original;
}

/* ================================================================
 * BeaconFormatInt + endian swap helper
 *
 * Cobalt Strike's structured data uses big-endian (network byte
 * order). swap_endianess converts native little-endian to big-endian.
 * ================================================================ */

int swap_endianess(int value) {
    return ((value >> 24) & 0x000000FF) |
           ((value >>  8) & 0x0000FF00) |
           ((value <<  8) & 0x00FF0000) |
           ((value << 24) & 0xFF000000);
}

void BeaconFormatInt(formatp* format, int value) {
    int swapped = swap_endianess(value);
    BeaconFormatAppend(format, (char*)&swapped, sizeof(int));
}

/* ================================================================
 * Token / Process Stubs
 *
 * COFFLoader runs as a standalone CLI tool, not an implant. These
 * functions provide minimal implementations or no-ops.
 * ================================================================ */

void BeaconUseToken(HANDLE token) {
    /* In a real implant: ImpersonateLoggedOnUser(token) */
    (void)token;
}

void BeaconRevertToken(void) {
    /* In a real implant: RevertToSelf() */
}

BOOL BeaconIsAdmin(void) {
    /* Simplified stub -- always returns FALSE */
    return FALSE;
}

void BeaconGetSpawnTo(BOOL x86, char* buffer, int length) {
    /* In a real implant: return the spawn-to process path */
    if (buffer && length > 0) {
        memset(buffer, 0, length);
    }
}

BOOL BeaconSpawnTemporaryProcess(BOOL x86, BOOL ignoreToken,
                                  STARTUPINFOA* si, PROCESS_INFORMATION* pi) {
    /* Stub -- not implemented in standalone loader */
    (void)x86; (void)ignoreToken; (void)si; (void)pi;
    return FALSE;
}

void BeaconInjectProcess(HANDLE hProc, int pid,
                          char* payload, int p_len, int offset,
                          char* arg, int a_len) {
    /* Stub -- not implemented in standalone loader */
    (void)hProc; (void)pid; (void)payload; (void)p_len;
    (void)offset; (void)arg; (void)a_len;
}

void BeaconInjectTemporaryProcess(PROCESS_INFORMATION* pi,
                                   char* payload, int p_len, int offset,
                                   char* arg, int a_len) {
    /* Stub -- not implemented in standalone loader */
    (void)pi; (void)payload; (void)p_len;
    (void)offset; (void)arg; (void)a_len;
}

void BeaconCleanupProcess(PROCESS_INFORMATION* pi) {
    /* Stub -- not implemented in standalone loader */
    (void)pi;
}

/* ================================================================
 * Utility: ANSI to wide string conversion
 * ================================================================ */

BOOL toWideChar(char* src, wchar_t* dst, int max) {
    return MultiByteToWideChar(CP_ACP, 0, src, -1, dst, max);
}

/* ================================================================
 * Output Retrieval
 *
 * Called after RunCOFF() completes to get the accumulated output
 * from BeaconPrintf/BeaconOutput calls. Resets the global buffer
 * so the next BOF invocation starts fresh.
 * ================================================================ */

char* BeaconGetOutputData(int* outsize) {
    char* output = beacon_compatibility_output;
    *outsize = beacon_compatibility_offset;

    /* Reset for next BOF execution */
    beacon_compatibility_output = NULL;
    beacon_compatibility_size   = 0;
    beacon_compatibility_offset = 0;

    return output;
}