
#include <ntddk.h>
#include "version.h"
#include "vad-def.h"
#include "vad-secure.h"


static ULONG VersionCode = VERSION_UNKNOWN;

static VOID VSSecureEntryToSecureInfo(PLIST_ENTRY SecureEntry, PVAD_SECURE_INFO_ENTRY SecureInfo)
{
  PMMSECURE_ENTRY_WXP EntryWXP = CONTAINING_RECORD(SecureEntry, MMSECURE_ENTRY_WXP, List);
  PMMSECURE_ENTRY_W2K3 EntryW2K3 = CONTAINING_RECORD(SecureEntry, MMSECURE_ENTRY_W2K3, List);
  PMMSECURE_ENTRY_VISTA EntryVista = CONTAINING_RECORD(SecureEntry, MMSECURE_ENTRY_VISTA, List);
  PMMSECURE_ENTRY_W7 EntryW7 = CONTAINING_RECORD(SecureEntry, MMSECURE_ENTRY_W7, List);
  KdPrint(("vad-secure.c: VSSecureEntryToSecureInfo(SecureEntry=0x%p; SecureInfo=0x%p)\n", SecureEntry, SecureInfo));

  switch (VersionCode) {
    case VERSION_WXP:
      SecureInfo->StartingVpn = EntryWXP->StartingVpn;
      SecureInfo->EndingVpn = EntryWXP->EndingVpn;
      SecureInfo->ReadOnly = EntryWXP->VadFlags.ReadOnly;
      break;
	case VERSION_W2K3:
      SecureInfo->StartingVpn = EntryW2K3->StartingVpn;
      SecureInfo->EndingVpn = EntryW2K3->EndingVpn;
      SecureInfo->ReadOnly = EntryW2K3->VadFlags.ReadOnly;
      break;
    case VERSION_VISTA:
      SecureInfo->StartingVpn = EntryVista->Range.StartVpn & ~1;
      SecureInfo->EndingVpn = EntryVista->Range.EndVpn;
      SecureInfo->ReadOnly = EntryVista->Range.StartVpn & 1;
      break;
    case VERSION_W7:
      SecureInfo->StartingVpn = EntryW7->Range.StartVpn & ~1;
      SecureInfo->EndingVpn = EntryW7->Range.EndVpn;
      SecureInfo->ReadOnly = EntryW7->Range.StartVpn & 1;
      break;
	default:
      KdPrint(("vad-secure.c: VSSecureEntryToSecureInfo(): ERROR: This OS version is not supported\n"));
      ASSERT(FALSE);
      break;
  }

  KdPrint(("vad-secure.c: VSSecureEntryToSecureInfo(-):StartAddress=0x%p, EndAddress=0x%p, ReadOnly=%u\n", (ULONG_PTR)SecureInfo->StartingVpn, (ULONG_PTR)SecureInfo->EndingVpn, SecureInfo->ReadOnly));
  return;
}


static VOID VSAddressListToSecureEntry(PMMADDRESS_LIST AddressList, PVAD_SECURE_INFO_ENTRY SecureEntry)
{
  KdPrint(("vad-secure.c: VSAddressListToSecureEntry(AddressList=0x%p; SecureEtry=0x%p)\n", AddressList, SecureEntry));

  SecureEntry->StartingVpn = AddressList->StartVpn;
  SecureEntry->EndingVpn = AddressList->EndVpn;
  SecureEntry->ReadOnly = AddressList->StartVpn & 1;

  KdPrint(("vad-secure.c: VSAddressListToSecureEntry(-)\n"));
  return;
}

static VOID VSSecureEntryListToSecureInfoList(PLIST_ENTRY ListHead, PVAD_SECURE_INFO Info)
{
  PLIST_ENTRY Tmp = ListHead->Blink;
  ULONG Count = 0;
  PVAD_SECURE_INFO_ENTRY InfoEntry = &Info->Entries[0];
  KdPrint(("vad-secure.c: VSSecureEntryListToSecureInfoList(ListHead=0x%p; Info=0x%p)\n", ListHead, Info));
  
  while (Tmp != ListHead && Count < Info->EntryCount) {
    VSSecureEntryToSecureInfo(Tmp, InfoEntry);
    ++Count;
    ++InfoEntry;
    Tmp = Tmp->Blink;
  }

  KdPrint(("vad-secure.c: VSSecureEntryListToSecureInfoList(-)\n"));
  return;
}


BOOLEAN VSIsVadSecured(PVOID Vad)
{
  MMVAD_FLAGS2 Flags2;
  BOOLEAN Ret = FALSE;
  KdPrint(("vad-secure.c: VSIsVadSecured(Vad=0x%p)\n", Vad));

  if (VDGetVadStructureType(Vad) == VadStructureLong) {
	Flags2 = VDGetVadFlags2(Vad);
	Ret = (Flags2.OneSecured == 1 || Flags2.MultipleSecured == 1);
  }

  KdPrint(("vad-secure.c: VSIsVadSecured(-):%u\n", Ret));
  return Ret;
}

ULONG VSGetVadSecuritySize(PVOID Vad)
{
  PLIST_ENTRY Entry = NULL;
  PLIST_ENTRY SecureList = NULL;
  MMVAD_FLAGS2 Flags2;
  ULONG Ret = 0;
  KdPrint(("vad-secure.c: VSGetVadSecuritySize(Vad=0x%p)\n", Vad));

  if (VSIsVadSecured(Vad)) {
    Flags2 = VDGetVadFlags2(Vad);
    if (Flags2.MultipleSecured == 1) {
      SecureList = (PLIST_ENTRY)VDGetVADSecurity(Vad);
      if (SecureList != NULL) {
        Entry = SecureList->Blink;
        while (Entry != SecureList) {
          Ret += sizeof(VAD_SECURE_INFO_ENTRY);
          Entry = Entry->Blink;
        }
      }
	} else Ret = sizeof(VAD_SECURE_INFO_ENTRY);
  } else KdPrint(("vad-secure.c: VSGetVadSecuritySize(): ERROR: This VAD is not secured\n"));

  KdPrint(("vad-secure.c: VSGetVadSecuritySize(-):%u\n", Ret));
  return Ret;
}


NTSTATUS VSGetVadSecurity(PVOID Vad, PVAD_SECURE_INFO SecureInfo)
{
  PVAD_SECURE_INFO_ENTRY SecureEntry = NULL;
  ULONG InfoSize = 0;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  MMVAD_FLAGS2 Flags2;
  KdPrint(("vad-secure.c: VSGetVadSecurity(Vad=0x%p; SecureInfo=0x%p)\n", Vad, SecureInfo));
  
  if (VSIsVadSecured(Vad)) {
    Flags2 = VDGetVadFlags2(Vad); 
    InfoSize = VSGetVadSecuritySize(Vad);
    SecureInfo->EntryCount = InfoSize / sizeof(VAD_SECURE_INFO_ENTRY);
    SecureEntry = &SecureInfo->Entries[0];
    Status = STATUS_SUCCESS;
    if (Flags2.MultipleSecured == 1) {
      PLIST_ENTRY VadSecureList = (PLIST_ENTRY)VDGetVADSecurity(Vad);

	  VSSecureEntryListToSecureInfoList(VadSecureList, SecureInfo);
    } else {
      PMMADDRESS_LIST VadAddressList = (PMMADDRESS_LIST)VDGetVADSecurity(Vad);
      
	  VSAddressListToSecureEntry(VadAddressList, SecureEntry);
    }
  } else KdPrint(("vad-secure.c: VSGetVadSecurity(): ERROR: This is not a secured VAD\n"));

#ifdef DBG
  {
    LONG i = 0;

    for (i = 0; i < SecureInfo->EntryCount; ++i)
      KdPrint(("SECUREENTRY: Start=0x%p; End=0x%p; ReadOnly=%u\n", 
          (ULONG_PTR)SecureInfo->Entries[i].StartingVpn,  
          (ULONG_PTR)SecureInfo->Entries[i].EndingVpn,
          SecureInfo->Entries[i].ReadOnly));
  }
#endif

  KdPrint(("vad-secure.c: VSGetVadSecurity(-):0x%x\n", Status));
  return Status;
}


NTSTATUS VSModuleInit(VOID)
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("vad-secure.c: VSModuleInit()\n"));

  VersionCode = GetVersionCode();
  Status = STATUS_SUCCESS;

  KdPrint(("vad-secure.c: VSModuleInit(-):0x%x\n", Status));
  return Status;
}


VOID VSModuleFinit(VOID)
{
  KdPrint(("vad-secure.c: VSModuleFinit()\n"));

  KdPrint(("vad-secure.c: VSModuleFinit(-)\n"));
  return;
}
