/**************************************************************************
 * Keeps track of memory allocation, and protection, thus we are able to 
 * locate self modifying code.
 * vmmap_area -> list which puts all pages with same Base, inside, is a list
 * of all pages with a different allocation base. In terms of Windows this
 * is returned by VirtualQuery in MEMORY_BASIC_INFORMATION:
 *
 * PVOID BaseAddress     <--- base address where region is allocated (when MEM_RESERVE
 *                                                                    is used on huge block)
 * PVOID AllocationBase  <--- actual commited memory
 *
 * All allocations with same BaseAddress go into vmmap_area, and are linked here
 * this helps in case of NtUnmapViewOfSection to easily remove memory and pages which we
 * don't need anymore
 *
 * Also when VirtualFree is called with size 0, this will free all pages from baseaddress
 ***************************************************************************/                                                                 
      
#include        "defs.h"

/***************************************************************
 * we need 4MB of memory to describe state of every page on a 
 * system. Page >> 12 is index into array where we keep lists
 * which will tell us protection of a page, basic blocks for
 * that page. Also protection is applied to check that page
 * didn't do self modify.
 ***************************************************************/
unsigned long   g_vmmap[0x100000];

static unsigned long   vmmap_windows_to_internal(__in unsigned long dwProtection, __in unsigned long currentProtection){
        unsigned long base_prot;
        
        base_prot = dwProtection & 0xFF;
         
        switch (base_prot)
        {
        case PAGE_EXECUTE_READ:
        case PAGE_EXECUTE:
        case PAGE_READONLY:
                return VMMAP_READ | VMMAP_EXEC;
        case PAGE_EXECUTE_READWRITE:
        case PAGE_READWRITE:
        case PAGE_EXECUTE_WRITECOPY:
                return VMMAP_READ | VMMAP_EXEC | VMMAP_WRITE;
        case PAGE_WRITECOPY:
                return VMMAP_READ | VMMAP_WRITE;
        case PAGE_NOACCESS:
                return 0;
        default:
                //DbgPrint(("%s -- base protection doesn't have translation : %.08X", __FUNCTION__, dwProtection));
                return 0;
        }
        
}

unsigned long vmmapRemoveRange(__in void *lpBuffer){
        MEMORY_BASIC_INFORMATION        mbi;
        unsigned long size              = 0;
        ULONG_PTR                       AllocationBase;
        ULONG_PTR                       base;
        
        ntVirtualQuery(lpBuffer, &mbi, sizeof(mbi));
        
        AllocationBase = (ULONG_PTR)mbi.AllocationBase;
        
        do{
                size += mbi.RegionSize;
                base = (ULONG_PTR)mbi.BaseAddress + mbi.RegionSize;
                ntVirtualQuery((PVOID)base, &mbi, sizeof(mbi));                
        }while (AllocationBase == (ULONG_PTR)mbi.AllocationBase);
        
        return size;
}

void    vmmapAdd(__in void *lpBuffer, __in unsigned long size, __in unsigned long dwProtection){
        unsigned long index;
        pvmmap_area parea;
        
        index = (ULONG_PTR)lpBuffer;                 
        while (index < (ULONG_PTR)lpBuffer + size){
                parea = vmmap_get_area((void *)index);
                if (dwProtection == 0){
                        DbgPrint(("%s -- for base : %.08X", __FUNCTION__, lpBuffer));
                        DbgPrint(("%s -- for size : %.08X", __FUNCTION__, size));
                        DbgPrint(("%s -- prot     : %.08X", __FUNCTION__, dwProtection));
                }
                
                parea->vmmap_protection = vmmap_windows_to_internal(dwProtection, 0);
                index += 0x1000;
        }
        
}

void    vmmapRemove(__in void *lpBuffer, __in unsigned long size){
        unsigned long index;
        unsigned long end;
        
        index = (unsigned long)lpBuffer >> 12;
        end   = (size >> 12) + index;
        
        while (index < end){
                g_vmmap[index] = 0;
                index++;
        }        
        
}

unsigned long vmmapGetProtection(__in void *lpBuffer){
        unsigned long   index;
        pvmmap_area     parea;
        MEMORY_BASIC_INFORMATION mbi;
        
        parea = vmmap_get_area(lpBuffer);
        if (parea->vmmap_protection == 0){
                ntVirtualQuery((void *)parea->vmmap_base, &mbi, sizeof(mbi));
                parea->vmmap_protection = vmmap_windows_to_internal(mbi.Protect, 0);        
        }
        return parea->vmmap_protection;
}


void    vmmapChangeProtection(__in void *lpBuffer, __in unsigned long size, __in unsigned long dwProtection){
        unsigned long index;
        unsigned long end;
        unsigned long prot;
        pvmmap_area   parea;
        
        
        index = (ULONG_PTR)lpBuffer;        
        prot = vmmap_windows_to_internal(dwProtection, 0);
        
        if (dwProtection == 0){
                DbgPrint(("%s -- for base : %.08X", __FUNCTION__, lpBuffer));
                DbgPrint(("%s -- for size : %.08X", __FUNCTION__, size));
                DbgPrint(("%s -- prot     : %.08X", __FUNCTION__, dwProtection));
        }
        
        while (index < (ULONG_PTR)lpBuffer + size){
                /**********************************************************************
                 * New protection doesn't have VMMAP_WRITE, but old one has VMMAP_WRITE 
                 * thus we mark page as VMMAP_WAS_WRITTEN 
                 * bbl builder knows to inspect this page...
                 ************************************************************/
                parea = vmmap_get_area((void *)index);
                if (parea->vmmap_protection & VMMAP_WRITE && !(prot & VMMAP_WRITE)){
                        parea->vmmap_protection = (unsigned char)prot | VMMAP_WAS_WRITE;                     
                }else{
                        parea->vmmap_protection = (unsigned char)prot;
                }
                index+= 0x1000;
        }        
        
}


void    vmmapRemoveWasWrite(__in void *lpBuffer){
        pvmmap_area   parea;
                
        parea = vmmap_get_area(lpBuffer);
        parea->vmmap_protection &= ~VMMAP_WAS_WRITE;       
}

pvmmap_area vmmap_get_area(__in void *ptr){
        unsigned long index;
        pvmmap_area     parea;
        MEMORY_BASIC_INFORMATION mbi;
        
        index = (unsigned long)ptr;        
        index = index >> 12;
        
        if (g_vmmap[index] == 0){
                ntVirtualQuery((void *)(index * 0x1000), &mbi, sizeof(mbi));
                parea = dlmalloc(sizeof(vmmap_area)); 
                memset(parea, 0, sizeof(vmmap_area));
                parea->vmmap_protection = vmmap_windows_to_internal(mbi.Protect, 0);
                if (mbi.Protect == 0 && mbi.State == MEM_COMMIT){
                        DbgPrint(("%s -- for base : %.08X", __FUNCTION__, index * 0x1000));
                        DbgPrint(("%s -- prot     : %.08X", __FUNCTION__, mbi.Protect));
                }
                parea->vmmap_base       = index * 0x1000;
                InitializeListHead(&parea->bbl_list);
                g_vmmap[index] = (unsigned long)parea;
        }else{
                parea = (pvmmap_area)g_vmmap[index];
        }
        
        return parea;                       
}

void    vmmap_delete_area(__in void *ptr){
        unsigned long index;
        pvmmap_area   parea;
        index = (unsigned long)ptr >> 12;
        
        if (g_vmmap[index] != 0){
                parea = (pvmmap_area)g_vmmap[index];
                if (parea->pbbl_array){
                        memset(parea->pbbl_array, 0, 0x1000 * 4);
                        dlfree(parea->pbbl_array);
                }
                memset((void *)g_vmmap[index], 0, sizeof(vmmap_area));
                dlfree((void *)g_vmmap[index]);
                g_vmmap[index] = 0;
        }
}
                

