helpers/api_knowledge.py

"""API knowledge base for v4 behavior analysis (defensive triage, capa-style).

Maps Windows API names (Win32, Ntdll, User32, networking) to neutral behavior tags.
Two uses:
  * teach the model what each API does (API-level questions), and
  * derive per-function yes/no capability answers from the set of APIs a function calls.

Tags are behavioral and neutral; this is descriptive documentation, not an implementation.
"""

# tag -> yes/no question to ask about a function whose calls carry that tag
TAG_QUESTIONS = {
    "other_process": "Does `function` open or access another process (a process other than itself)?",
    "process_mem_write": "Does `function` write to another process's memory?",
    "process_mem_read": "Does `function` read another process's memory?",
    "mem_alloc": "Does `function` allocate memory at runtime?",
    "mem_protect": "Does `function` change memory protection flags (for example make memory executable)?",
    "create_thread": "Does `function` create a thread?",
    "remote_thread": "Does `function` create a thread inside another process?",
    "create_process": "Does `function` create or launch a new process?",
    "file_io": "Does `function` read from or write to files?",
    "registry": "Does `function` read from or write to the Windows registry?",
    "network": "Does `function` perform network communication?",
    "dynamic_import": "Does `function` resolve an API or module address at runtime?",
    "enum_process": "Does `function` enumerate running processes, threads, or modules?",
    "keyboard_input": "Does `function` read keyboard or other user input state?",
    "window_ui": "Does `function` interact with windows or the graphical user interface?",
    "crypto": "Does `function` perform cryptographic operations via the OS crypto APIs?",
    "service": "Does `function` create or control a Windows service?",
    "token_priv": "Does `function` query or modify security tokens or privileges?",
    "sync": "Does `function` use synchronization primitives (mutex, event, semaphore, waits)?",
    "inspect_state": "Does `function` query debugger, timing, or system-inspection state?",
    "shutdown_power": "Does `function` reboot, shut down, or change system power state?",
    "clipboard": "Does `function` read from or write to the system clipboard?",
    "environment": "Does `function` read process environment variables or command line?",
    "driver_io": "Does `function` create or dispatch driver I/O (device objects, IRPs, IOCTLs)?",
    "driver_object": "Does `function` create, open, or manipulate kernel objects by name (registry, symbolic links, sections, devices)?",
    "driver_callback": "Does `function` register a kernel-mode callback or notification (image load, process create, registry, filter)?",
    "kernel_sync": "Does `function` use kernel synchronization primitives (KEVENT, KMUTEX, ERESOURCE, spinlocks)?",
    "dpc_apc": "Does `function` schedule a DPC, APC, or timer routine?",
}
TAGS = list(TAG_QUESTIONS)

# API -> list of tags. Names are stored as the base (no A/W suffix). The lookup helpers accept
# A/W and __imp_/__imp__ decorations. Multiple tags per API are common (e.g. WriteProcessMemory
# implies both "other_process" and "process_mem_write").
API_TAGS = {
    # process / thread
    "OpenProcess": ["other_process"],
    "OpenThread": ["other_process"],
    "NtOpenProcess": ["other_process"], "ZwOpenProcess": ["other_process"],
    "NtOpenThread": ["other_process"], "ZwOpenThread": ["other_process"],
    "TerminateProcess": ["other_process"], "NtTerminateProcess": ["other_process"],
    "WriteProcessMemory": ["other_process", "process_mem_write"],
    "NtWriteVirtualMemory": ["other_process", "process_mem_write"],
    "ZwWriteVirtualMemory": ["other_process", "process_mem_write"],
    "ReadProcessMemory": ["other_process", "process_mem_read"],
    "NtReadVirtualMemory": ["other_process", "process_mem_read"],
    "ZwReadVirtualMemory": ["other_process", "process_mem_read"],
    "CreateThread": ["create_thread"], "NtCreateThread": ["create_thread"],
    "CreateRemoteThread": ["create_thread", "remote_thread", "other_process"],
    "CreateRemoteThreadEx": ["create_thread", "remote_thread", "other_process"],
    "NtCreateThreadEx": ["create_thread", "remote_thread", "other_process"],
    "ZwCreateThreadEx": ["create_thread", "remote_thread", "other_process"],
    "RtlCreateUserThread": ["create_thread", "remote_thread", "other_process"],
    "QueueUserAPC": ["other_process"], "NtQueueApcThread": ["other_process"],
    "ZwQueueApcThread": ["other_process"], "QueueUserAPC2": ["other_process"],
    "CreateProcess": ["create_process"], "CreateProcessAsUser": ["create_process"],
    "CreateProcessWithLogon": ["create_process"], "CreateProcessWithToken": ["create_process"],
    "CreateProcessInternal": ["create_process"],
    "NtCreateProcess": ["create_process"], "NtCreateProcessEx": ["create_process"],
    "NtCreateUserProcess": ["create_process"], "ZwCreateUserProcess": ["create_process"],
    "ShellExecute": ["create_process"], "ShellExecuteEx": ["create_process"],
    "WinExec": ["create_process"], "system": ["create_process"],
    "ResumeThread": [], "SuspendThread": ["other_process"],
    "SetThreadContext": ["other_process"], "GetThreadContext": ["other_process"],
    "NtSetContextThread": ["other_process"], "NtGetContextThread": ["other_process"],
    "OpenProcessToken": ["token_priv", "other_process"],
    "OpenThreadToken": ["token_priv"],
    "GetCurrentProcess": [], "GetCurrentThread": [], "GetCurrentProcessId": [],
    "GetCurrentThreadId": [], "ExitProcess": [], "ExitThread": [],

    # enumeration
    "CreateToolhelp32Snapshot": ["enum_process"],
    "Process32First": ["enum_process"], "Process32Next": ["enum_process"],
    "Thread32First": ["enum_process"], "Thread32Next": ["enum_process"],
    "Module32First": ["enum_process"], "Module32Next": ["enum_process"],
    "EnumProcesses": ["enum_process"], "EnumProcessModules": ["enum_process"],
    "EnumProcessModulesEx": ["enum_process"],
    "NtQuerySystemInformation": ["enum_process", "inspect_state"],
    "ZwQuerySystemInformation": ["enum_process", "inspect_state"],
    "NtQueryInformationProcess": ["other_process", "inspect_state"],

    # memory
    "VirtualAlloc": ["mem_alloc"], "VirtualAllocEx": ["mem_alloc", "other_process"],
    "VirtualAllocExNuma": ["mem_alloc", "other_process"],
    "NtAllocateVirtualMemory": ["mem_alloc"],
    "ZwAllocateVirtualMemory": ["mem_alloc"],
    "HeapAlloc": ["mem_alloc"], "HeapReAlloc": ["mem_alloc"],
    "LocalAlloc": ["mem_alloc"], "GlobalAlloc": ["mem_alloc"],
    "RtlAllocateHeap": ["mem_alloc"], "malloc": ["mem_alloc"], "calloc": ["mem_alloc"],
    "realloc": ["mem_alloc"], "operator new": ["mem_alloc"],
    "VirtualFree": [], "VirtualFreeEx": ["other_process"], "HeapFree": [],
    "LocalFree": [], "GlobalFree": [], "RtlFreeHeap": [], "free": [],
    "VirtualProtect": ["mem_protect"], "VirtualProtectEx": ["mem_protect", "other_process"],
    "NtProtectVirtualMemory": ["mem_protect"],
    "ZwProtectVirtualMemory": ["mem_protect"],
    "MapViewOfFile": ["mem_alloc"], "MapViewOfFileEx": ["mem_alloc"],
    "CreateFileMapping": ["mem_alloc"], "OpenFileMapping": ["mem_alloc"],
    "NtMapViewOfSection": ["mem_alloc"],

    # dynamic import
    "LoadLibrary": ["dynamic_import"], "LoadLibraryEx": ["dynamic_import"],
    "LdrLoadDll": ["dynamic_import"],
    "GetProcAddress": ["dynamic_import"], "LdrGetProcedureAddress": ["dynamic_import"],
    "GetModuleHandle": ["dynamic_import"], "GetModuleHandleEx": ["dynamic_import"],
    "LdrGetDllHandle": ["dynamic_import"],

    # file I/O
    "CreateFile": ["file_io"], "OpenFile": ["file_io"],
    "ReadFile": ["file_io"], "WriteFile": ["file_io"],
    "ReadFileEx": ["file_io"], "WriteFileEx": ["file_io"],
    "SetFilePointer": ["file_io"], "SetFilePointerEx": ["file_io"],
    "DeleteFile": ["file_io"], "MoveFile": ["file_io"], "MoveFileEx": ["file_io"],
    "CopyFile": ["file_io"], "CopyFileEx": ["file_io"],
    "GetFileAttributes": ["file_io"], "GetFileSize": ["file_io"],
    "FindFirstFile": ["file_io"], "FindNextFile": ["file_io"], "FindClose": [],
    "CreateDirectory": ["file_io"], "RemoveDirectory": ["file_io"],
    "NtCreateFile": ["file_io"], "NtOpenFile": ["file_io"],
    "NtReadFile": ["file_io"], "NtWriteFile": ["file_io"],
    "ZwCreateFile": ["file_io"], "ZwOpenFile": ["file_io"],
    "ZwReadFile": ["file_io"], "ZwWriteFile": ["file_io"],
    "fopen": ["file_io"], "fread": ["file_io"], "fwrite": ["file_io"],
    "fclose": [], "_open": ["file_io"], "_read": ["file_io"], "_write": ["file_io"],

    # registry
    "RegOpenKey": ["registry"], "RegOpenKeyEx": ["registry"],
    "RegCreateKey": ["registry"], "RegCreateKeyEx": ["registry"],
    "RegSetValue": ["registry"], "RegSetValueEx": ["registry"],
    "RegQueryValue": ["registry"], "RegQueryValueEx": ["registry"],
    "RegDeleteValue": ["registry"], "RegDeleteKey": ["registry"], "RegDeleteKeyEx": ["registry"],
    "RegEnumKey": ["registry"], "RegEnumKeyEx": ["registry"],
    "RegEnumValue": ["registry"], "RegCloseKey": [],
    "NtOpenKey": ["registry"], "NtCreateKey": ["registry"],
    "NtSetValueKey": ["registry"], "NtQueryValueKey": ["registry"],
    "NtDeleteKey": ["registry"], "NtDeleteValueKey": ["registry"],
    "ZwOpenKey": ["registry"], "ZwCreateKey": ["registry"],

    # network - sockets
    "WSAStartup": ["network"], "WSACleanup": [],
    "socket": ["network"], "WSASocket": ["network"],
    "connect": ["network"], "WSAConnect": ["network"],
    "bind": ["network"], "listen": ["network"],
    "accept": ["network"], "WSAAccept": ["network"],
    "send": ["network"], "recv": ["network"],
    "WSASend": ["network"], "WSARecv": ["network"],
    "sendto": ["network"], "recvfrom": ["network"],
    "closesocket": [], "shutdown": [],
    "getaddrinfo": ["network"], "GetAddrInfoW": ["network"],
    "gethostbyname": ["network"], "inet_addr": ["network"],
    "inet_pton": ["network"], "inet_ntop": ["network"],
    "select": ["network"], "WSAPoll": ["network"],
    # network - high level
    "InternetOpen": ["network"], "InternetConnect": ["network"],
    "InternetOpenUrl": ["network"], "InternetReadFile": ["network"],
    "InternetWriteFile": ["network"], "InternetSetOption": ["network"],
    "HttpOpenRequest": ["network"], "HttpSendRequest": ["network"],
    "HttpAddRequestHeaders": ["network"], "HttpQueryInfo": ["network"],
    "WinHttpOpen": ["network"], "WinHttpConnect": ["network"],
    "WinHttpOpenRequest": ["network"], "WinHttpSendRequest": ["network"],
    "WinHttpReceiveResponse": ["network"], "WinHttpReadData": ["network"],
    "WinHttpWriteData": ["network"], "WinHttpQueryHeaders": ["network"],
    "URLDownloadToFile": ["network", "file_io"],
    "DnsQuery": ["network"], "DnsQuery_A": ["network"], "DnsQuery_W": ["network"],

    # crypto
    "CryptAcquireContext": ["crypto"], "CryptReleaseContext": ["crypto"],
    "CryptGenKey": ["crypto"], "CryptImportKey": ["crypto"], "CryptExportKey": ["crypto"],
    "CryptDestroyKey": ["crypto"],
    "CryptEncrypt": ["crypto"], "CryptDecrypt": ["crypto"],
    "CryptHashData": ["crypto"], "CryptCreateHash": ["crypto"],
    "CryptGetHashParam": ["crypto"], "CryptDeriveKey": ["crypto"],
    "CryptGenRandom": ["crypto"], "BCryptGenRandom": ["crypto"],
    "BCryptOpenAlgorithmProvider": ["crypto"], "BCryptCloseAlgorithmProvider": ["crypto"],
    "BCryptCreateHash": ["crypto"], "BCryptHashData": ["crypto"], "BCryptFinishHash": ["crypto"],
    "BCryptEncrypt": ["crypto"], "BCryptDecrypt": ["crypto"],
    "BCryptGenerateSymmetricKey": ["crypto"], "BCryptDestroyKey": ["crypto"],
    "NCryptOpenStorageProvider": ["crypto"], "NCryptEncrypt": ["crypto"], "NCryptDecrypt": ["crypto"],

    # services
    "OpenSCManager": ["service"], "OpenService": ["service"],
    "CreateService": ["service"], "StartService": ["service"],
    "ControlService": ["service"], "DeleteService": ["service"],
    "QueryServiceStatus": ["service"], "QueryServiceStatusEx": ["service"],
    "EnumServicesStatus": ["service"], "EnumServicesStatusEx": ["service"],
    "ChangeServiceConfig": ["service"], "ChangeServiceConfig2": ["service"],
    "CloseServiceHandle": [],

    # tokens / privileges
    "LookupPrivilegeValue": ["token_priv"], "AdjustTokenPrivileges": ["token_priv"],
    "GetTokenInformation": ["token_priv"], "SetTokenInformation": ["token_priv"],
    "DuplicateToken": ["token_priv"], "DuplicateTokenEx": ["token_priv"],
    "ImpersonateLoggedOnUser": ["token_priv"], "RevertToSelf": ["token_priv"],
    "LogonUser": ["token_priv"], "LogonUserEx": ["token_priv"],
    "CreateProcessAsUser": ["create_process", "token_priv"],
    "NtOpenProcessToken": ["token_priv", "other_process"],
    "NtAdjustPrivilegesToken": ["token_priv"],

    # synchronization
    "CreateMutex": ["sync"], "OpenMutex": ["sync"], "ReleaseMutex": ["sync"],
    "CreateEvent": ["sync"], "OpenEvent": ["sync"], "SetEvent": ["sync"], "ResetEvent": ["sync"],
    "CreateSemaphore": ["sync"], "OpenSemaphore": ["sync"], "ReleaseSemaphore": ["sync"],
    "WaitForSingleObject": ["sync"], "WaitForSingleObjectEx": ["sync"],
    "WaitForMultipleObjects": ["sync"], "WaitForMultipleObjectsEx": ["sync"],
    "MsgWaitForMultipleObjects": ["sync"], "SignalObjectAndWait": ["sync"],
    "Sleep": ["sync"], "SleepEx": ["sync"], "NtDelayExecution": ["sync"],
    "EnterCriticalSection": ["sync"], "LeaveCriticalSection": ["sync"],
    "InitializeCriticalSection": [], "DeleteCriticalSection": [],
    "AcquireSRWLockExclusive": ["sync"], "AcquireSRWLockShared": ["sync"],

    # system state / inspection (self-info + timing)
    "IsDebuggerPresent": ["inspect_state"], "CheckRemoteDebuggerPresent": ["inspect_state"],
    "NtQueryInformationThread": ["inspect_state"],
    "OutputDebugString": ["inspect_state"],
    "GetTickCount": ["inspect_state"], "GetTickCount64": ["inspect_state"],
    "QueryPerformanceCounter": ["inspect_state"], "QueryPerformanceFrequency": ["inspect_state"],
    "GetSystemTime": ["inspect_state"], "GetSystemTimeAsFileTime": ["inspect_state"],
    "GetLocalTime": ["inspect_state"], "GetLastError": [], "SetLastError": [],
    "GetVersion": ["inspect_state"], "GetVersionEx": ["inspect_state"],
    "GetComputerName": ["inspect_state"], "GetUserName": ["inspect_state"],
    "GetSystemInfo": ["inspect_state"], "GetNativeSystemInfo": ["inspect_state"],

    # user input (kernel/user surface)
    "GetAsyncKeyState": ["keyboard_input"], "GetKeyState": ["keyboard_input"],
    "GetKeyboardState": ["keyboard_input"], "SetWindowsHookEx": ["keyboard_input", "window_ui"],
    "UnhookWindowsHookEx": ["window_ui"], "CallNextHookEx": ["window_ui"],
    "RegisterHotKey": ["keyboard_input"], "UnregisterHotKey": [],
    "GetMessage": ["window_ui"], "PeekMessage": ["window_ui"], "DispatchMessage": ["window_ui"],
    "TranslateMessage": ["window_ui"], "SendInput": ["keyboard_input"],
    "keybd_event": ["keyboard_input"], "mouse_event": ["keyboard_input"],

    # window / UI
    "CreateWindow": ["window_ui"], "CreateWindowEx": ["window_ui"],
    "ShowWindow": ["window_ui"], "UpdateWindow": ["window_ui"], "DestroyWindow": ["window_ui"],
    "MessageBox": ["window_ui"], "MessageBoxEx": ["window_ui"],
    "FindWindow": ["window_ui"], "FindWindowEx": ["window_ui"],
    "GetForegroundWindow": ["window_ui"], "SetForegroundWindow": ["window_ui"],
    "GetWindowText": ["window_ui"], "SetWindowText": ["window_ui"],
    "EnumWindows": ["window_ui"], "EnumChildWindows": ["window_ui"],
    "GetWindowThreadProcessId": ["window_ui", "other_process"],
    "PostMessage": ["window_ui"], "SendMessage": ["window_ui"],

    # clipboard
    "OpenClipboard": ["clipboard"], "CloseClipboard": [],
    "GetClipboardData": ["clipboard"], "SetClipboardData": ["clipboard"],
    "EmptyClipboard": ["clipboard"],

    # environment / command line
    "GetCommandLine": ["environment"], "GetEnvironmentVariable": ["environment"],
    "SetEnvironmentVariable": ["environment"], "GetEnvironmentStrings": ["environment"],
    "ExpandEnvironmentStrings": ["environment"],
    "getenv": ["environment"], "_wgetenv": ["environment"],

    # shutdown / power
    "ExitWindowsEx": ["shutdown_power"], "InitiateSystemShutdown": ["shutdown_power"],
    "InitiateSystemShutdownEx": ["shutdown_power"], "InitiateShutdown": ["shutdown_power"],
    "SetSystemPowerState": ["shutdown_power"],

    # --- kernel / driver APIs ---

    # more Nt/Zw system calls used from kernel/user land
    "NtClose": [], "ZwClose": [],
    "NtLoadDriver": ["service"], "ZwLoadDriver": ["service"],
    "NtUnloadDriver": ["service"], "ZwUnloadDriver": ["service"],
    "NtQueryDirectoryObject": ["enum_process", "driver_object"],
    "ZwQueryDirectoryObject": ["enum_process", "driver_object"],
    "NtCreateSection": ["mem_alloc", "driver_object"],
    "ZwCreateSection": ["mem_alloc", "driver_object"],
    "NtOpenSection": ["driver_object"], "ZwOpenSection": ["driver_object"],
    "NtMapViewOfSection": ["mem_alloc"], "ZwMapViewOfSection": ["mem_alloc"],
    "NtUnmapViewOfSection": [], "ZwUnmapViewOfSection": [],
    "NtCreateSymbolicLinkObject": ["driver_object"], "ZwCreateSymbolicLinkObject": ["driver_object"],
    "NtOpenSymbolicLinkObject": ["driver_object"], "ZwOpenSymbolicLinkObject": ["driver_object"],
    "NtCreateDirectoryObject": ["driver_object"], "ZwCreateDirectoryObject": ["driver_object"],
    "NtOpenDirectoryObject": ["driver_object"], "ZwOpenDirectoryObject": ["driver_object"],
    "NtDeviceIoControlFile": ["driver_io"], "ZwDeviceIoControlFile": ["driver_io"],
    "NtFsControlFile": ["driver_io"], "ZwFsControlFile": ["driver_io"],
    "NtDuplicateObject": ["driver_object"], "ZwDuplicateObject": ["driver_object"],
    "NtEnumerateKey": ["registry"], "ZwEnumerateKey": ["registry"],
    "NtEnumerateValueKey": ["registry"], "ZwEnumerateValueKey": ["registry"],
    "NtQueryKey": ["registry"], "ZwQueryKey": ["registry"],
    "NtNotifyChangeKey": ["registry", "driver_callback"], "ZwNotifyChangeKey": ["registry", "driver_callback"],
    "NtQueryDirectoryFile": ["file_io"], "ZwQueryDirectoryFile": ["file_io"],
    "NtOpenSection": ["driver_object"], "NtQueryInformationFile": ["file_io"],
    "ZwQueryInformationFile": ["file_io"], "NtSetInformationFile": ["file_io"],
    "ZwSetInformationFile": ["file_io"],

    # IoManager (WDM / kernel)
    "IoCreateDevice": ["driver_io", "driver_object"],
    "IoCreateDeviceSecure": ["driver_io", "driver_object"],
    "IoDeleteDevice": ["driver_io"],
    "IoCreateSymbolicLink": ["driver_object"], "IoDeleteSymbolicLink": ["driver_object"],
    "IoAttachDeviceToDeviceStack": ["driver_io"], "IoDetachDevice": ["driver_io"],
    "IoAttachDevice": ["driver_io"], "IoGetDeviceObjectPointer": ["driver_object"],
    "IoAllocateIrp": ["driver_io"], "IoBuildDeviceIoControlRequest": ["driver_io"],
    "IoBuildSynchronousFsdRequest": ["driver_io"], "IoBuildAsynchronousFsdRequest": ["driver_io"],
    "IoCallDriver": ["driver_io"], "IoFreeIrp": [],
    "IoCompleteRequest": ["driver_io"], "IoCancelIrp": ["driver_io"],
    "IoSetCompletionRoutine": ["driver_io", "driver_callback"],
    "IoSetCompletionRoutineEx": ["driver_io", "driver_callback"],
    "IoRegisterDeviceInterface": ["driver_io"], "IoRegisterFileSystem": ["driver_io"],
    "IoRegisterShutdownNotification": ["driver_callback", "shutdown_power"],
    "IoRegisterFsRegistrationChange": ["driver_callback"],
    "IoQueueWorkItem": ["dpc_apc"], "IoAllocateWorkItem": ["dpc_apc"], "IoFreeWorkItem": [],
    "IoCreateFile": ["file_io"], "IoOpenDeviceRegistryKey": ["registry"],
    "IoWMIRegistrationControl": ["driver_callback"],
    "IofCallDriver": ["driver_io"], "IofCompleteRequest": ["driver_io"],

    # Ex - executive
    "ExAllocatePool": ["mem_alloc"], "ExAllocatePoolWithTag": ["mem_alloc"],
    "ExAllocatePool2": ["mem_alloc"], "ExAllocatePool3": ["mem_alloc"],
    "ExFreePool": [], "ExFreePoolWithTag": [],
    "ExInitializeFastMutex": ["kernel_sync"], "ExAcquireFastMutex": ["kernel_sync"],
    "ExReleaseFastMutex": ["kernel_sync"],
    "ExInitializeResourceLite": ["kernel_sync"], "ExAcquireResourceExclusiveLite": ["kernel_sync"],
    "ExAcquireResourceSharedLite": ["kernel_sync"], "ExReleaseResourceLite": ["kernel_sync"],
    "ExDeleteResourceLite": [],
    "ExRegisterCallback": ["driver_callback"], "ExUnregisterCallback": [],
    "ExCreateCallback": ["driver_callback"], "ExNotifyCallback": ["driver_callback"],
    "ExQueueWorkItem": ["dpc_apc"], "ExInitializeWorkItem": [],
    "ExSetTimerResolution": ["inspect_state"],
    "ExGetPreviousMode": ["inspect_state"],

    # Ke - kernel
    "KeInitializeEvent": ["kernel_sync"], "KeSetEvent": ["kernel_sync"],
    "KeResetEvent": ["kernel_sync"], "KeClearEvent": ["kernel_sync"],
    "KeWaitForSingleObject": ["kernel_sync"], "KeWaitForMultipleObjects": ["kernel_sync"],
    "KeInitializeMutex": ["kernel_sync"], "KeReleaseMutex": ["kernel_sync"],
    "KeInitializeSemaphore": ["kernel_sync"], "KeReleaseSemaphore": ["kernel_sync"],
    "KeInitializeSpinLock": ["kernel_sync"],
    "KeAcquireSpinLock": ["kernel_sync"], "KeReleaseSpinLock": ["kernel_sync"],
    "KeAcquireSpinLockAtDpcLevel": ["kernel_sync"], "KeReleaseSpinLockFromDpcLevel": ["kernel_sync"],
    "KeAcquireInStackQueuedSpinLock": ["kernel_sync"], "KeReleaseInStackQueuedSpinLock": ["kernel_sync"],
    "KeInitializeDpc": ["dpc_apc"], "KeInsertQueueDpc": ["dpc_apc"], "KeRemoveQueueDpc": [],
    "KeInitializeTimer": ["dpc_apc"], "KeInitializeTimerEx": ["dpc_apc"],
    "KeSetTimer": ["dpc_apc"], "KeSetTimerEx": ["dpc_apc"], "KeCancelTimer": [],
    "KeInitializeApc": ["dpc_apc"], "KeInsertQueueApc": ["dpc_apc"],
    "KeStackAttachProcess": ["other_process"], "KeUnstackDetachProcess": [],
    "KeGetCurrentThread": [], "KeGetCurrentIrql": ["inspect_state"],
    "KeQueryPerformanceCounter": ["inspect_state"], "KeQuerySystemTime": ["inspect_state"],
    "KeBugCheck": ["shutdown_power"], "KeBugCheckEx": ["shutdown_power"],
    "KeDelayExecutionThread": ["sync"], "KeRaiseIrqlToDpcLevel": ["kernel_sync"],
    "KeLowerIrql": ["kernel_sync"],

    # Mm - memory manager
    "MmAllocateNonCachedMemory": ["mem_alloc"], "MmFreeNonCachedMemory": [],
    "MmAllocateContiguousMemory": ["mem_alloc"], "MmFreeContiguousMemory": [],
    "MmMapIoSpace": ["mem_alloc"], "MmUnmapIoSpace": [],
    "MmMapLockedPages": ["mem_alloc"], "MmMapLockedPagesSpecifyCache": ["mem_alloc"],
    "MmUnmapLockedPages": [], "MmProbeAndLockPages": ["mem_protect"],
    "MmUnlockPages": [], "MmBuildMdlForNonPagedPool": [],
    "MmGetSystemAddressForMdlSafe": [],
    "MmProtectMdlSystemAddress": ["mem_protect"],
    "MmIsAddressValid": ["inspect_state"],
    "MmSecureVirtualMemory": ["mem_protect"], "MmUnsecureVirtualMemory": [],

    # Ob - object manager
    "ObReferenceObjectByHandle": ["driver_object"],
    "ObReferenceObjectByPointer": ["driver_object"],
    "ObReferenceObjectByName": ["driver_object"],
    "ObDereferenceObject": [], "ObfDereferenceObject": [], "ObfReferenceObject": ["driver_object"],
    "ObOpenObjectByName": ["driver_object"], "ObOpenObjectByPointer": ["driver_object"],
    "ObRegisterCallbacks": ["driver_callback"],
    "ObUnRegisterCallbacks": [],
    "ObGetObjectSecurity": ["token_priv"], "ObSetHandleAttributes": ["driver_object"],

    # Ps - process/thread
    "PsCreateSystemThread": ["create_thread"], "PsTerminateSystemThread": [],
    "PsLookupProcessByProcessId": ["other_process"],
    "PsLookupThreadByThreadId": ["other_process"],
    "PsGetCurrentProcess": [], "PsGetCurrentProcessId": [],
    "PsGetCurrentThread": [], "PsGetCurrentThreadId": [],
    "PsGetProcessId": [], "PsGetProcessImageFileName": ["enum_process"],
    "PsSetCreateProcessNotifyRoutine": ["driver_callback", "enum_process"],
    "PsSetCreateProcessNotifyRoutineEx": ["driver_callback", "enum_process"],
    "PsSetCreateProcessNotifyRoutineEx2": ["driver_callback", "enum_process"],
    "PsSetCreateThreadNotifyRoutine": ["driver_callback", "enum_process"],
    "PsSetCreateThreadNotifyRoutineEx": ["driver_callback", "enum_process"],
    "PsSetLoadImageNotifyRoutine": ["driver_callback", "enum_process"],
    "PsSetLoadImageNotifyRoutineEx": ["driver_callback", "enum_process"],
    "PsRemoveCreateThreadNotifyRoutine": [],
    "PsRemoveLoadImageNotifyRoutine": [],
    "PsIsProtectedProcess": ["inspect_state"], "PsIsSystemThread": ["inspect_state"],

    # Cm - configuration manager (registry callbacks)
    "CmRegisterCallback": ["registry", "driver_callback"],
    "CmRegisterCallbackEx": ["registry", "driver_callback"],
    "CmUnRegisterCallback": [],
    "CmCallbackGetKeyObjectID": ["registry"], "CmGetCallbackVersion": [],

    # Rtl - runtime library (kernel + user)
    "RtlInitUnicodeString": [], "RtlInitAnsiString": [],
    "RtlAnsiStringToUnicodeString": [], "RtlUnicodeStringToAnsiString": [],
    "RtlCopyMemory": [], "RtlMoveMemory": [], "RtlZeroMemory": [], "RtlSecureZeroMemory": [],
    "RtlCompareMemory": [], "RtlEqualMemory": [],
    "RtlAllocateHeap": ["mem_alloc"], "RtlFreeHeap": [],
    "RtlCreateHeap": ["mem_alloc"], "RtlDestroyHeap": [],
    "RtlGetVersion": ["inspect_state"], "RtlGetNtProductType": ["inspect_state"],
    "RtlQueryRegistryValues": ["registry"], "RtlWriteRegistryValue": ["registry"],
    "RtlDeleteRegistryValue": ["registry"], "RtlCheckRegistryKey": ["registry"],
    "RtlCreateRegistryKey": ["registry"],
    "RtlRandomEx": [], "RtlRandom": [], "RtlComputeCrc32": [],
    "RtlImageDirectoryEntryToData": ["dynamic_import"],
    "RtlImageNtHeader": ["dynamic_import"],
    "RtlLookupFunctionEntry": ["inspect_state"], "RtlVirtualUnwind": ["inspect_state"],
    "RtlCaptureContext": ["inspect_state"], "RtlCaptureStackBackTrace": ["inspect_state"],

    # Se - security
    "SeSinglePrivilegeCheck": ["token_priv"],
    "SeAccessCheck": ["token_priv"],
    "SePrivilegeCheck": ["token_priv"],
    "SeCaptureSubjectContext": ["token_priv"],
    "SeReleaseSubjectContext": [],

    # Fs / Flt - filter manager
    "FltRegisterFilter": ["driver_callback", "file_io"],
    "FltUnregisterFilter": [],
    "FltStartFiltering": ["driver_callback"],
    "FltBuildDefaultSecurityDescriptor": ["token_priv"],
    "FltCreateCommunicationPort": ["driver_io"],
    "FltCloseCommunicationPort": [],
    "FltSendMessage": ["driver_io"],
    "FltGetRoutineAddress": ["dynamic_import"],
    "FltGetFileNameInformation": ["file_io"],
    "FltReleaseFileNameInformation": [],
    "FsRtlIsNameInExpression": ["file_io"],
    "FsRtlCheckLockForReadAccess": ["file_io"],
    "FsRtlCheckLockForWriteAccess": ["file_io"],

    # Etw / logging
    "EtwRegister": ["driver_callback"], "EtwUnregister": [],
    "EtwWrite": [], "EtwEventEnabled": [],
    "DbgPrint": ["inspect_state"], "DbgPrintEx": ["inspect_state"],
    "KdPrintEx": ["inspect_state"],

    # driver entry helpers
    "IoGetCurrentProcess": [], "IoGetRequestorProcess": ["other_process"],
    "IoGetRequestorProcessId": ["other_process"],
    "IoGetTopLevelIrp": [], "IoSetTopLevelIrp": [],

    # WDF (Windows Driver Framework)
    "WdfDriverCreate": ["driver_io"], "WdfDeviceCreate": ["driver_io"],
    "WdfDeviceInitAssignName": ["driver_object"],
    "WdfDeviceInitSetIoType": ["driver_io"],
    "WdfIoQueueCreate": ["driver_io"], "WdfRequestComplete": ["driver_io"],
    "WdfRequestCompleteWithInformation": ["driver_io"],
    "WdfRequestGetParameters": ["driver_io"],
    "WdfMemoryCreate": ["mem_alloc"],
    "WdfSpinLockCreate": ["kernel_sync"], "WdfSpinLockAcquire": ["kernel_sync"],
    "WdfSpinLockRelease": ["kernel_sync"],
    "WdfWaitLockCreate": ["kernel_sync"], "WdfWaitLockAcquire": ["kernel_sync"],
    "WdfWaitLockRelease": ["kernel_sync"],
    "WdfDpcCreate": ["dpc_apc"], "WdfDpcEnqueue": ["dpc_apc"],
    "WdfTimerCreate": ["dpc_apc"], "WdfTimerStart": ["dpc_apc"],
    "WdfWorkItemCreate": ["dpc_apc"], "WdfWorkItemEnqueue": ["dpc_apc"],
    "WdfRegistryOpenKey": ["registry"], "WdfRegistryCreateKey": ["registry"],
    "WdfRegistryQueryValue": ["registry"], "WdfRegistryAssignValue": ["registry"],
    "WdfIoTargetCreate": ["driver_io"], "WdfIoTargetSendIoctlSynchronously": ["driver_io"],
    "WdfFileObjectGetFileName": ["file_io"],
}


def _strip_decoration(name):
    """Normalize an import name: drop __imp_/__imp__ decoration, MSVC/MinGW A/W suffix, drop ordinal
    forms like `foo@8`, drop stdcall `_` prefix on 32-bit imports. Leave Nt/Zw as-is."""
    n = name.strip()
    if not n:
        return n
    for prefix in ("__imp_", "__imp__", "_imp_"):
        if n.startswith(prefix):
            n = n[len(prefix):]
    # cdecl @size suffix
    if "@" in n:
        n = n.split("@", 1)[0]
    # api-ms-win- forwarders and ordinal names left alone
    return n


def lookup(name):
    """Return the list of tags for an API name, tolerating A/W and decorated variants."""
    n = _strip_decoration(name)
    if n in API_TAGS:
        return API_TAGS[n]
    # try dropping trailing A/W (Win32 ansi/wide)
    if len(n) > 1 and n[-1] in "AW" and n[:-1] in API_TAGS:
        return API_TAGS[n[:-1]]
    return []


def tags_for_calls(names):
    """Union of tags across a list of API names."""
    out = set()
    for n in names:
        out.update(lookup(n))
    return out


def sanity():
    """Quick self-check: report table stats and a few example lookups."""
    from collections import Counter
    per_tag = Counter()
    for tags in API_TAGS.values():
        for t in tags:
            per_tag[t] += 1
    missing = [t for t in TAGS if t not in per_tag]
    print(f"APIs: {len(API_TAGS)}    tags: {len(TAGS)}    tags with no APIs: {missing}")
    for t, n in sorted(per_tag.items()):
        print(f"  {t:18} {n:4}")
    for probe in ["CreateFileW", "__imp_CreateRemoteThread", "NtWriteVirtualMemory",
                  "WSAStartup", "GetProcAddress", "OpenClipboard", "unknown_func"]:
        print(f"  lookup({probe!r}) -> {lookup(probe)}")


if __name__ == "__main__":
    sanity()