#include <stdint.h>
#include <stdarg.h>
#include <windef.h>
#include <winbase.h>
#include <winnt.h>
#include <ntstatus.h>
#include <winioctl.h>
#define MAX_KEY_LENGTH 255
#define MAX_VALUE_NAME 16383
#define NT_PFX L"\\GLOBAL?"

// Basic NT Types
typedef long NTSTATUS, *PNTSTATUS;

typedef struct _UNICODE_STRING
{
    USHORT Length;
    USHORT MaximumLength;
    PWSTR Buffer;
} UNICODE_STRING, *PUNICODE_STRING;
typedef const UNICODE_STRING* PCUNICODE_STRING;

PVOID NTAPI
RtlAllocateHeap(VOID *HeapHandle, ULONG Flags, SIZE_T Size);

BOOLEAN NTAPI
RtlFreeHeap(HANDLE HeapHandle, ULONG Flags, VOID *P);

VOID NTAPI
RtlInitUnicodeString(PUNICODE_STRING DestinationString, PCWSTR SourceString);

BOOLEAN NTAPI
RtlEqualUnicodeString(PCUNICODE_STRING String1, PCUNICODE_STRING String2,
                      BOOLEAN CaseInsensitive);

typedef struct _CLIENT_ID {
    HANDLE UniqueProcess;
    HANDLE UniqueThread;
} CLIENT_ID, *PCLIENT_ID;

typedef struct _PEB {
    BYTE Reserved1[4];
    VOID *Reserved2[5];
    VOID *ProcessHeap;
    BYTE Reserved4[88];
    VOID *Reserved5[53];
    BYTE Reserved6[128];
    VOID *Reserved7[1];
    ULONG SessionId;
} PEB, *PPEB;

typedef struct __TEB {
  NT_TIB Tib;
  VOID  *EnvironmentPointer;
  CLIENT_ID Cid;
  VOID *ActiveRpcInfo;
  VOID *ThreadLocalStoragePointer;
  PEB *ProcessEnvironmentBlock;
  unsigned int LastStatusValue;
  BYTE Reserved1[1844];
  PVOID Reserved2[412];
  PVOID TlsSlots[64];
  BYTE Reserved3[8];
  PVOID Reserved4[26];
  PVOID ReservedForOle;
  PVOID Reserved5[4];
  PVOID *TlsExpansionSlots;
} _TEB;
#define NtCurrentPeb() (((_TEB*)NtCurrentTeb())->ProcessEnvironmentBlock)
#define RtlGetProcessHeap() (NtCurrentPeb()->ProcessHeap)

typedef struct _devlist_t {
    void *next;
    ULONG idx;
    UNICODE_STRING devpath;
} devlist_t;

static inline void devlist_append(devlist_t **list, unsigned long dev_idx, unsigned long subdev_idx, wchar_t *devpathname) {
    devlist_t *new_list = RtlAllocateHeap(RtlGetProcessHeap(), 0, sizeof(devlist_t));
    new_list->next = *list;
    new_list->idx = ((dev_idx << 16)|subdev_idx);
    RtlInitUnicodeString(&(new_list->devpath), devpathname);
    *list = new_list;
}

//
// Registry Specific Access Rights.
//
#define NT_KEY_QUERY_VALUE         (0x0001)
#define NT_KEY_SET_VALUE           (0x0002)
#define NT_KEY_CREATE_SUB_KEY      (0x0004)
#define NT_KEY_ENUMERATE_SUB_KEYS  (0x0008)
#define NT_KEY_NOTIFY              (0x0010)
#define NT_KEY_CREATE_LINK         (0x0020)
#define NT_KEY_READ                ((STANDARD_RIGHTS_READ|NT_KEY_QUERY_VALUE|NT_KEY_ENUMERATE_SUB_KEYS|NT_KEY_NOTIFY) & (~SYNCHRONIZE))
#define NT_KEY_WRITE               ((STANDARD_RIGHTS_WRITE|NT_KEY_SET_VALUE|NT_KEY_CREATE_SUB_KEY) & (~SYNCHRONIZE))
#define NT_KEY_EXECUTE             ((NT_KEY_READ) & (~SYNCHRONIZE))
#define NT_KEY_ALL_ACCESS          ((STANDARD_RIGHTS_ALL|NT_KEY_QUERY_VALUE|NT_KEY_SET_VALUE|NT_KEY_CREATE_SUB_KEY|NT_KEY_ENUMERATE_SUB_KEYS|NT_KEY_NOTIFY|NT_KEY_CREATE_LINK) & (~SYNCHRONIZE))

//
// Key query structures
//
typedef struct _KEY_BASIC_INFORMATION {
    LARGE_INTEGER LastWriteTime;
    ULONG   TitleIndex;
    ULONG   NameLength;
    WCHAR   Name[1];            // Variable length string
} KEY_BASIC_INFORMATION, *PKEY_BASIC_INFORMATION;

typedef struct _KEY_FULL_INFORMATION {
    LARGE_INTEGER LastWriteTime;
    ULONG   TitleIndex;
    ULONG   ClassOffset;
    ULONG   ClassLength;
    ULONG   SubKeys;
    ULONG   MaxNameLen;
    ULONG   MaxClassLen;
    ULONG   Values;
    ULONG   MaxValueNameLen;
    ULONG   MaxValueDataLen;
    WCHAR   Class[1];           // Variable length
} KEY_FULL_INFORMATION, *PKEY_FULL_INFORMATION;
#define KeyInfo(key) ((KEY_BASIC_INFORMATION*)key)

typedef enum _KEY_INFORMATION_CLASS {
    KeyBasicInformation,
    KeyNodeInformation,
    KeyFullInformation,
    KeyNameInformation,
    KeyCachedInformation,
    KeyFlagsInformation,
    MaxKeyInfoClass  // MaxKeyInfoClass should always be the last enum
} KEY_INFORMATION_CLASS;

typedef struct _OBJECT_ATTRIBUTES
{
    ULONG Length;
    HANDLE RootDirectory;
    PUNICODE_STRING ObjectName; ULONG Attributes; PVOID SecurityDescriptor; PVOID SecurityQualityOfService;
} OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;
#define OBJ_CASE_INSENSITIVE 0x00000040L

unsigned int NTAPI
NtOpenKey(HANDLE *KeyHandle, ACCESS_MASK DesiredAccess, OBJECT_ATTRIBUTES *oa);
unsigned int NTAPI
NtClose(HANDLE KeyHandle);
unsigned int NTAPI
NtEnumerateKey(HANDLE KeyHandle, ULONG Index, KEY_INFORMATION_CLASS KeyInformationClass,
               VOID *KeyInformation, ULONG Length, ULONG *ResultLength);
unsigned int NTAPI
NtQueryKey(HANDLE KeyHandle, KEY_INFORMATION_CLASS KeyInformationClass,
           VOID *KeyInformation, ULONG Length, ULONG *ResultLength);

// WaveRT drivers can be found by grabbing the result from the below query under KSCATEGORY_AUDIO ({6994AD04-93EF-11D0-A3CC-00A0C9223196})
// and checking to see if it is also present under KSCATEGORY_REALTIME ({EB115FFC-10C8-4964-831D-6DCB02E6F23F}). If yes, it's a WaveRT driver.
// If no, it's a WavePCI (or god forbid, a WaveCyclic driver).
devlist_t *NtWdmKSEnumerateAudioDevices(UNICODE_STRING *KeyName)
{ 
    KEY_FULL_INFORMATION keyInfo, AudioKeyInfo; 
    NTSTATUS r;
    DWORD i, j;
    ULONG wtf;
    HANDLE hAudioKey;
    OBJECT_ATTRIBUTES oa;
    devlist_t *list = NULL;

    memset(&oa, 0, sizeof(oa));
    oa.Length = sizeof(oa);
    oa.RootDirectory = NULL;
    oa.Attributes = OBJ_CASE_INSENSITIVE;
    oa.ObjectName = KeyName;

    // Open the key in question.
    r = NtOpenKey(&hAudioKey, NT_KEY_READ, &oa);
    if (r != STATUS_SUCCESS) {
      return list;
    }
 
    // Get the class name and the value count. 
    r = NtQueryKey(hAudioKey, KeyFullInformation, &keyInfo, sizeof(KEY_FULL_INFORMATION), &wtf);
    if (r != STATUS_SUCCESS) {
        return list;
    }

    // Enumerate the subkeys, until RegEnumKeyEx fails.
    if (keyInfo.SubKeys) {
        for (i=0; i<keyInfo.SubKeys; i++) {
            HANDLE hSubKey, hFilterKey;
            UNICODE_STRING AudioDevKey;
            ULONG KeyInfoSizeNeeded;
            WCHAR cur_key[1024];
            SIZE_T cur_key_len = 0;
            SIZE_T sz = 0x1000;
            unsigned char *buf = NULL;
            buf = RtlAllocateHeap(RtlGetProcessHeap(), HEAP_ZERO_MEMORY, sz);
            r = NtEnumerateKey(hAudioKey, i, KeyBasicInformation, buf, 0x1000, &KeyInfoSizeNeeded);
            if ((r != STATUS_SUCCESS) && (KeyInfoSizeNeeded > 0x1000)) {
                return list;
            }
            memcpy(cur_key+6, KeyInfo(buf)->Name, KeyInfo(buf)->NameLength);
            memcpy(cur_key, NT_PFX, 8 * sizeof(WCHAR));
            cur_key[9] = L'\\';
            cur_key[KeyInfo(buf)->NameLength + 6] = L'\0';
            cur_key_len = KeyInfo(buf)->NameLength+6;

            memset(&oa, 0, sizeof(oa));
            oa.Length = sizeof(oa);
            oa.RootDirectory = hAudioKey;
            oa.Attributes = OBJ_CASE_INSENSITIVE;
            RtlInitUnicodeString(&AudioDevKey, KeyInfo(buf)->Name);
            oa.ObjectName = &AudioDevKey;

            r = NtOpenKey(&hSubKey, NT_KEY_READ, &oa);
            if (r != STATUS_SUCCESS) {
                RtlFreeHeap(RtlGetProcessHeap(), 0, buf);
                return list;
            }

            r = NtQueryKey(hSubKey, KeyFullInformation, &AudioKeyInfo, sizeof(KEY_FULL_INFORMATION), &wtf);
            if (r != STATUS_SUCCESS) {
                continue;
            }

            for (j=0; j<AudioKeyInfo.SubKeys; j++) {
                unsigned char *strend;
                UNICODE_STRING AudioFilterKey;
                NtEnumerateKey(hSubKey, j, KeyBasicInformation, buf, 0x1000, &KeyInfoSizeNeeded);
                strend = ((unsigned char*)KeyInfo(buf)->Name + KeyInfo(buf)->NameLength);
                memcpy(strend, L"\\Control\0", 9*sizeof(WCHAR));

                oa.RootDirectory = hSubKey;
                RtlInitUnicodeString(&AudioFilterKey, KeyInfo(buf)->Name);
                oa.ObjectName = &AudioFilterKey;

                r = NtOpenKey(&hFilterKey, NT_KEY_READ, &oa);
                if (r == STATUS_SUCCESS) {
                    wchar_t *tmpbuf = NULL;
                    SIZE_T tmplen = cur_key_len + KeyInfo(buf)->NameLength;
                    (KeyInfo(buf)->Name[0]) = L'\\';
                    tmpbuf = RtlAllocateHeap(RtlGetProcessHeap(), 0, (tmplen+3)*sizeof(wchar_t));
                    tmplen >>= 1;

                    // Concatenate together the full path
                    memcpy(tmpbuf, cur_key, cur_key_len+6);
                    memcpy(tmpbuf+(cur_key_len >> 1)+3, KeyInfo(buf)->Name, KeyInfo(buf)->NameLength);
                    tmpbuf[tmplen+3] = L'\0';
                    devlist_append(&list, i, j, tmpbuf);
                    NtClose(hFilterKey);
                }
            }
            NtClose(hSubKey);
            RtlFreeHeap(RtlGetProcessHeap(), 0, buf);
        }
    }
    NtClose(hAudioKey);
    return list;
}

unsigned int NtWdmKSGetNumDevices(devlist_t *list)
{
  unsigned int n = 0;
  while (list != NULL) {
    list = list->next;
    n++;
  }
  return n;
}

BOOLEAN NtWdmKSIsDeviceInAliasList(devlist_t *list, UNICODE_STRING *devpath)
{ 
  while (list != NULL) { 
    devlist_t *cur_dev = list;
    list = list->next;
    if (RtlEqualUnicodeString(&(cur_dev->devpath), devpath, FALSE)) {
        return TRUE;
    }
  }
  return FALSE;
}

void NtWdmKSFreeAudioDeviceList(devlist_t *list)
{
    while (list != NULL) {
        devlist_t *cur_dev = list;
        list = list->next;
        RtlFreeHeap(RtlGetProcessHeap(), 0, cur_dev->devpath.Buffer);
        RtlFreeHeap(RtlGetProcessHeap(), 0, cur_dev);
    }
}

#ifdef TEST
#include <wincon.h>

static void dump_devlist(devlist_t *list)
{
  HANDLE errh = RtlGetConsoleHandle();
  wchar_t buf[1024];
  int len;
  DWORD done;

  while (list != NULL) {
    devlist_t *cur_dev = list;
    list = list->next;
    len = _snwprintf(buf, 1023, L"Found KS filter: device: %lu, filter index: %lu, device path: %ls\n", (cur_dev->idx >> 16), (cur_dev->idx & 0xFFFF),
                     cur_dev->devpath.Buffer);
    WriteConsoleW(errh, buf, len, &done, NULL);
    RtlFreeHeap(RtlGetProcessHeap(), 0, cur_dev->devpath.Buffer);
    RtlFreeHeap(RtlGetProcessHeap(), 0, cur_dev);
  }
}

int mainCRTStartup (void)
{
  HANDLE errh = RtlGetConsoleHandle();
  UNICODE_STRING KeyName;
  devlist_t *list;
  DWORD done;

  RtlInitUnicodeString(&KeyName, L"\\Registry\\Machine\\SYSTEM\\CurrentControlSet\\Control\\DeviceClasses\\{6994AD04-93EF-11D0-A3CC-00A0C9223196}");
  WriteConsoleW(errh, L"Enumerating KSCATEGORY_AUDIO devices...\n", 40, &done, NULL);
  list = NtWdmKSEnumerateAudioDevices(&KeyName);
  dump_devlist(list);

  RtlInitUnicodeString(&KeyName, L"\\Registry\\Machine\\SYSTEM\\CurrentControlSet\\Control\\DeviceClasses\\{65E8773E-8F56-11D0-A3B9-00A0C9223196}");
  WriteConsoleW(errh, L"Enumerating KSCATEGORY_RENDER devices...\n", 41, &done, NULL);
  list = NtWdmKSEnumerateAudioDevices(&KeyName);
  dump_devlist(list);

  RtlInitUnicodeString(&KeyName, L"\\Registry\\Machine\\SYSTEM\\CurrentControlSet\\Control\\DeviceClasses\\{65E8773D-8F56-11D0-A3B9-00A0C9223196}");
  WriteConsoleW(errh, L"Enumerating KSCATEGORY_CAPTURE devices...\n", 42, &done, NULL);
  list = NtWdmKSEnumerateAudioDevices(&KeyName);
  dump_devlist(list);

  RtlInitUnicodeString(&KeyName, L"\\Registry\\Machine\\SYSTEM\\CurrentControlSet\\Control\\DeviceClasses\\{EB115FFC-10C8-4964-831D-6DCB02E6F23F}");
  WriteConsoleW(errh, L"Enumerating KSCATEGORY_REALTIME devices...\n", 43, &done, NULL);
  list = NtWdmKSEnumerateAudioDevices(&KeyName);
  dump_devlist(list);

  NtTerminateProcess(((HANDLE)-1), ((list != NULL) ? 0 : -1));
  return ((list != NULL) ? 0 : -1);
}
#endif