import sys
import ctypes
import pytest

import windows.generated_def as gdef
from windows.winobject.apisetmap import get_api_set_map_for_current_process

from .pfwtest import *


# Late getattr param si these can be at the end od the file
@pytest.fixture(params=[
    ("APISETMAP_V6", 6),
    ("APISETMAP_V4", 4),
    ("APISETMAP_V2", 2)
    ],ids=["Version6", "Version4","Version2"])
def dumped_apisetmap_base_and_version(request):
    variable_name, version = request.param
    data = b64decode(getattr(sys.modules[__name__], variable_name))
    ctypes_data = ctypes.c_buffer(data)
    yield ctypes.addressof(ctypes_data), version

KNOWN_APISETMAP_PREFIX = ["api-", "ext-", "MS-Win-"]

def verify_apisetmap_parsing(apisetmap_base, version=None):
    if version is not None:
        assert windows.current_process.read_dword(apisetmap_base) == version
    apisetmap = get_api_set_map_for_current_process(apisetmap_base)
    apisetmap_dict = apisetmap.resolution_dict
    # import pdb;pdb.set_trace()
    assert apisetmap_dict
    # Verify that at least one entry resolve to kernel32.dll
    # This ensure that the ApiSetMap parsing works at least a little
    assert "kernel32.dll" in apisetmap_dict.values()
    assert all(any(dll.startswith(pref) for pref in KNOWN_APISETMAP_PREFIX) for dll in apisetmap_dict)
    # This first key was found in most of the tested version by hand
    # MS-Win found on: 6.1.7600 (Win7)
    assert 'api-ms-win-core-com-l2-1-' in apisetmap_dict or "MS-Win-Core-ProcessThreads-L1-1-" in apisetmap_dict



def test_apisetmap_parsing_current_process():
    return verify_apisetmap_parsing(windows.current_process.peb.ApiSetMap)

def test_apisetmap_parsing_from_dump(dumped_apisetmap_base_and_version):
    return verify_apisetmap_parsing(*dumped_apisetmap_base_and_version)

APISETMAP_V2 = b"""
[encoded data, 5536 characters]
"""

APISETMAP_V4 = b"""
[encoded data, 82998 characters]
"""

APISETMAP_V6 = b"""
[encoded data, 121733 characters]
"""