Sigma-C2/agent/transport_schannel.c

#define SECURITY_WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#include <schannel.h>
#include <sspi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "transport.h"
#include "log.h"

#define BUFFER_SIZE 4096

// Structure to hold SChannel connection state
typedef struct {
    SOCKET sock;
    CredHandle cred;
    SecHandle ctx;
    SecPkgContext_StreamSizes sizes;
    BOOL initialized;
} SchannelHandle;

// Initialize SChannel and perform TLS handshake
BOOL init_schannel(SchannelHandle* handle, char* domain, unsigned short port, char* sni_hostname) {
    // Initialize Winsock
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        LOG_ERROR("Error: WSAStartup failed: %d\n", WSAGetLastError());
        return FALSE;
    }

    // Create socket
    handle->sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (handle->sock == INVALID_SOCKET) {
        LOG_ERROR("Error: Failed to create socket: %d\n", WSAGetLastError());
        WSACleanup();
        return FALSE;
    }

    // Set up server address
    struct sockaddr_in server = { .sin_family = AF_INET, .sin_port = htons(port) };
    if (InetPtonA(AF_INET, domain, &server.sin_addr) <= 0) {
        LOG_ERROR("Error: Failed to resolve domain %s\n", domain);
        closesocket(handle->sock);
        WSACleanup();
        return FALSE;
    }

    // Connect
    if (connect(handle->sock, (struct sockaddr*)&server, sizeof(server)) == SOCKET_ERROR) {
        LOG_ERROR("Error: Failed to connect to %s:%d: %d\n", domain, port, WSAGetLastError());
        closesocket(handle->sock);
        WSACleanup();
        return FALSE;
    }

    // Initialize Schannel credentials
    SCHANNEL_CRED schannelCred = { .dwVersion = SCHANNEL_CRED_VERSION, .grbitEnabledProtocols = SP_PROT_TLS1_2_CLIENT,
        .dwFlags = SCH_CRED_NO_SERVERNAME_CHECK | SCH_CRED_NO_DEFAULT_CREDS | SCH_CRED_MANUAL_CRED_VALIDATION };
    SECURITY_STATUS status = AcquireCredentialsHandle(NULL, UNISP_NAME, SECPKG_CRED_OUTBOUND,
        NULL, &schannelCred, NULL, NULL, &handle->cred, NULL);
    if (status != SEC_E_OK) {
        LOG_ERROR("Error: AcquireCredentialsHandle failed: 0x%x\n", status);
        closesocket(handle->sock);
        WSACleanup();
        return FALSE;
    }

    // Perform TLS handshake
    SecBufferDesc outBuffer, inBuffer;
    SecBuffer outBuffers[2] = { {0, SECBUFFER_TOKEN, NULL}, {0, SECBUFFER_EMPTY, NULL} };
    SecBuffer inBuffers[2] = { {BUFFER_SIZE, SECBUFFER_TOKEN, malloc(BUFFER_SIZE)}, {0, SECBUFFER_EMPTY, NULL} };
    outBuffer.cBuffers = 2; outBuffer.pBuffers = outBuffers; outBuffer.ulVersion = SECBUFFER_VERSION;
    inBuffer.cBuffers = 2; inBuffer.pBuffers = inBuffers; inBuffer.ulVersion = SECBUFFER_VERSION;
    DWORD flags = ISC_REQ_SEQUENCE_DETECT | ISC_REQ_REPLAY_DETECT | ISC_REQ_CONFIDENTIALITY |
                  ISC_REQ_EXTENDED_ERROR | ISC_REQ_ALLOCATE_MEMORY | ISC_REQ_STREAM;
    BOOL first = TRUE;
    int bufferLen;

    while (TRUE) {
        status = InitializeSecurityContext(&handle->cred, first ? NULL : &handle->ctx, sni_hostname, flags, 0, 0,
            first ? NULL : &inBuffer, 0, first ? &handle->ctx : NULL, &outBuffer, &flags, NULL);
        if (status == SEC_E_OK || status == SEC_I_CONTINUE_NEEDED) {
            if (outBuffers[0].cbBuffer && outBuffers[0].pvBuffer) {
                send(handle->sock, outBuffers[0].pvBuffer, outBuffers[0].cbBuffer, 0);
                FreeContextBuffer(outBuffers[0].pvBuffer);
                outBuffers[0].pvBuffer = NULL;
            }
        }
        if (status == SEC_E_OK) {
            break;
        }
        if (status != SEC_I_CONTINUE_NEEDED) {
            LOG_ERROR("Error: InitializeSecurityContext failed: 0x%x\n", status);
            free(inBuffers[0].pvBuffer);
            closesocket(handle->sock);
            DeleteSecurityContext(&handle->ctx);
            FreeCredentialsHandle(&handle->cred);
            WSACleanup();
            return FALSE;
        }
        bufferLen = recv(handle->sock, inBuffers[0].pvBuffer, BUFFER_SIZE, 0);
        if (bufferLen <= 0) {
            LOG_ERROR("Error: recv failed: %d\n", WSAGetLastError());
            free(inBuffers[0].pvBuffer);
            closesocket(handle->sock);
            DeleteSecurityContext(&handle->ctx);
            FreeCredentialsHandle(&handle->cred);
            WSACleanup();
            return FALSE;
        }
        inBuffers[0].cbBuffer = bufferLen;
        first = FALSE;
    }
    free(inBuffers[0].pvBuffer);

    // Get stream sizes for encryption/decryption
    status = QueryContextAttributes(&handle->ctx, SECPKG_ATTR_STREAM_SIZES, &handle->sizes);
    if (status != SEC_E_OK) {
        LOG_ERROR("Error: QueryContextAttributes failed: 0x%x\n", status);
        closesocket(handle->sock);
        DeleteSecurityContext(&handle->ctx);
        FreeCredentialsHandle(&handle->cred);
        WSACleanup();
        return FALSE;
    }

    handle->initialized = TRUE;
    return TRUE;
}

// Send encrypted data
int schannel_send(void* handle, char* data, size_t len) {
    SchannelHandle* schannel = (SchannelHandle*)handle;
    if (!schannel->initialized) return -1;

    // Allocate buffer for encrypted data
    char* sendBuffer = (char*)malloc(schannel->sizes.cbHeader + len + schannel->sizes.cbTrailer);
    if (!sendBuffer) return -1;
    memcpy(sendBuffer + schannel->sizes.cbHeader, data, len);

    // Set up encryption buffers
    SecBuffer encryptBuffers[4] = {
        {schannel->sizes.cbHeader, SECBUFFER_STREAM_HEADER, sendBuffer},
        {len, SECBUFFER_DATA, sendBuffer + schannel->sizes.cbHeader},
        {schannel->sizes.cbTrailer, SECBUFFER_STREAM_TRAILER, sendBuffer + schannel->sizes.cbHeader + len},
        {0, SECBUFFER_EMPTY, NULL}
    };
    SecBufferDesc encryptDesc = { SECBUFFER_VERSION, 4, encryptBuffers };

    // Encrypt and send
    SECURITY_STATUS status = EncryptMessage(&schannel->ctx, 0, &encryptDesc, 0);
    if (status != SEC_E_OK) {
        free(sendBuffer);
        return -1;
    }
    int totalSize = encryptBuffers[0].cbBuffer + encryptBuffers[1].cbBuffer + encryptBuffers[2].cbBuffer;
    int sent = send(schannel->sock, sendBuffer, totalSize, 0);
    free(sendBuffer);
    return sent == totalSize ? len : -1;
}

// Receive and decrypt data
int schannel_recv(void* handle, char* buffer, size_t len) {
    SchannelHandle* schannel = (SchannelHandle*)handle;
    if (!schannel->initialized) return -1;

    char recvBuffer[BUFFER_SIZE];
    int recvLen = recv(schannel->sock, recvBuffer, BUFFER_SIZE, 0);
    if (recvLen <= 0) return recvLen;

    // Set up decryption buffers
    SecBuffer decryptBuffers[4] = {
        {recvLen, SECBUFFER_DATA, recvBuffer},
        {0, SECBUFFER_EMPTY, NULL},
        {0, SECBUFFER_EMPTY, NULL},
        {0, SECBUFFER_EMPTY, NULL}
    };
    SecBufferDesc decryptDesc = { SECBUFFER_VERSION, 4, decryptBuffers };

    SECURITY_STATUS status = DecryptMessage(&schannel->ctx, &decryptDesc, 0, NULL);
    if (status != SEC_E_OK) return -1;

    // Find and copy decrypted data
    for (int i = 0; i < 4; i++) {
        if (decryptBuffers[i].BufferType == SECBUFFER_DATA && decryptBuffers[i].cbBuffer > 0) {
            size_t copyLen = min(decryptBuffers[i].cbBuffer, len);
            memcpy(buffer, decryptBuffers[i].pvBuffer, copyLen);
            return copyLen;
        }
    }
    return 0;
}

// Cleanup SChannel resources
void schannel_cleanup(void* handle) {
    SchannelHandle* schannel = (SchannelHandle*)handle;
    if (schannel->initialized) {
        DeleteSecurityContext(&schannel->ctx);
        FreeCredentialsHandle(&schannel->cred);
        if (schannel->sock != INVALID_SOCKET) {
            shutdown(schannel->sock, SD_BOTH);
            closesocket(schannel->sock);
        }
        WSACleanup();
    }
    free(schannel);
}

// Initialize SChannel transport
Transport* InitSchannelTransport(char* domain, unsigned short port, char* sni_hostname) {
    SchannelHandle* handle = (SchannelHandle*)calloc(1, sizeof(SchannelHandle));
    if (!handle) {
        LOG_ERROR("Error: Failed to allocate memory for Schannel handle\n");
        return NULL;
    }

    if (!init_schannel(handle, domain, port, sni_hostname)) {
        free(handle);
        return NULL;
    }

    Transport* transport = (Transport*)malloc(sizeof(Transport));
    if (!transport) {
        schannel_cleanup(handle);
        return NULL;
    }

    transport->handle = handle;
    transport->send = schannel_send;
    transport->recv = schannel_recv;
    transport->cleanup = schannel_cleanup;
    return transport;
}