/*
 * scanphysmem
 *
 * Yoann Guillot, 2005-10-15
 *
 * retrieves the virtual memory map from a physical memory dump
 * for windows processes, by searching PE headers in the memory
 * and looking for a PGD such that the PE would be loaded at its preferred
 * base address
 *
 * usage: ./scan <dump>
 */
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <signal.h>

typedef unsigned long u32;

const void *ptr;	// mmap(physdump)
u32 sz;			// physdump length

volatile int showadvance = 0;
void sigh(int sig __attribute__((unused))) {
	showadvance = 1;
	alarm(1);
}

const char const *ptemsk[4] = { "sys ro", "sys rw", "usr ro", "usr rw" };
#define msk2s(m) (ptemsk[((m) >> 1) & 3])
void dump_map(const u32 off)
{
	const u32 *pgd = ptr + off;
	const u32 *pt;
	u32 i, j;

	printf("Virtual memory map from PGD at offset 0x%.8lx:\n", off);

#define NOSTART (u32)-1
	u32 start = NOSTART, end = 0, dmask = 0, tmask = 0, zone4M = 0;
	
#define PRINT_NOW \
	do { if (start != NOSTART) {		\
		if (zone4M) \
			printf("\t%.8lx-%.8lx (%.4lX - %s)\n", start, end, dmask, msk2s(dmask)); \
		else \
			printf("\t%.8lx-%.8lx (%.3lX %.3lX - %s, %s)\n", start, end, dmask, tmask, msk2s(dmask), msk2s(tmask)); \
		start = NOSTART;		\
	}} while(0)
	
	for (i = 0 ; i < 1024 ; i++) {
		// notpresent ?
		if (!(pgd[i] & 1)) {
			PRINT_NOW;
			continue;
		}
		
		// 4M page ?
		if (pgd[i] & 0x80) {
			if ((start != NOSTART) && zone4M && ((pgd[i] & 0x1FFF) == dmask)) {
				end = ((i+1) << 22) - 1;
				continue;
			}
			
			PRINT_NOW;
			start = (i << 22);
			end = ((i+1) << 22) - 1;
			dmask = pgd[i] & 0x1FFF;
			zone4M = 1;
			continue;
		}
		
		// 4k page
		if (!zone4M || ((pgd[i] & 0xFFF) != dmask))
			PRINT_NOW;
		
		pt = ptr + (pgd[i] & 0xFFFFF000);
		for (j=0 ; j<1024 ; j++) {
			if (!(pt[j] & 1)) {
				PRINT_NOW;
				continue;
			}
			if (start != NOSTART && (pt[j] & 0xFFF) == tmask) {
				end = (i << 22) + ((j+1) << 12) - 1;
				continue;
			}
			PRINT_NOW;
			start = (i << 22) + (j << 12);
			end = (i << 22) + ((j+1) << 12) - 1;
			dmask = pgd[i] & 0xFFF;
			tmask = pt[j] & 0xFFF;
			zone4M = 0;
		}
	}
	PRINT_NOW;
	printf("\n");
}

// scans the memory for page offsets such that isvalid returns 1,
// put these offsets in a table allocated here. Returns the size of the table.
u32 fill_pageoffsets(u32 **t, int (*isvalid)(u32))
{
	u32 off;
	u32 cur = 0, max = 1024;
	*t = (u32 *)malloc(4*max);
	if (!*t) {
		perror("malloc");
		return 0;
	}

	signal(SIGALRM, sigh);
	alarm(1);
	for (off = 0 ; off < sz ; off += 4096) {
		if (showadvance) {
			// timer thing
			showadvance = 0;
			fprintf(stderr, "\r          \r%d%%",
					(int)(((float)off)*100/(float)sz));
		}

		if (isvalid(off)) {
			(*t)[cur++] = off;
			if (cur == max) {
				max += 1024;
				u32 *old = *t;
				*t = realloc(*t, 4*max);
				if (!*t) {
					*t = old;
					perror("realloc");
					return cur;
				}
			}
		}
	}
	alarm(0);
	fprintf(stderr, "\r          \r");
	return cur;
}

// returns the prefered base address of the PE
static u32 getPEbase(u32 off)
{
	const u32 *ul;
	
	ul = ptr + off + 0x3c;
	if (*ul > 4092 - (4 + 0x14 + 0x1c))
		// not in current page
		return 0x00400000;
	
	ul = ptr + off + *ul + 4 + 0x14 + 0x1c;
	return *ul;
}
	
/* valid header PE = MZ + PE */
int isvalid_PEhdr(u32 off)
{
	const unsigned short *mz;
	const u32 *ul;
	
	mz = ptr + off;
	if (*mz != 0x5a4d)
		return 0;
	
	ul = ptr + off + 0x3c;
	if (*ul > 4096 - 4)
		// out of the current page
		return 1;
	
	ul = ptr + off + *ul;
	if (*ul == 0x00004550) {
//		printf("found pe at %.8lX, prefbase %.8lX\n", off, getPEbase(off));
		return 1;
	}

	return 0;
}

// ensure that each page table pointed by the pgd is inside physmem
int valid_pgd(u32 off)
{
	const u32 *pgd = ptr + off;
	u32 pde;
	int i;
	
	for (i=0 ; i<1024 ; i++) {
		pde = pgd[i];
		if (!(pde & 1) || (pde & 0x80))
			continue;
		if ((pde & 0xfffff000) > sz-4096) {
//			printf("invalid pde at %.8lX: %.8lX\n", (u32)&pgd[i], pde);
			return 0;
		}
	}
	return 1;
}

u32 *pehdr = 0, *pebase = 0, pehdr_tlen;
// is there at offset a page directory pointing to a PE ?
int isvalid_PD2PE(u32 off)
{
	const u32 *pd = ptr + off;
	u32 i;

	for (i=0 ; i<pehdr_tlen ; i++) {
		// for each known PE:
		register u32 base;
		register u32 pde;
		const u32 *pt;
		
		// retrieve its base address
		base = pebase[i];

		// inspects the pde of this address
		pde = pd[ (base >> 22) & 0x3ff ];
		
		if (!(pde & 1))
			// not present
			continue;
		
		if (pde & 0x80) {
			// 4M page
			if (((pde & 0xffc00000) == (pehdr[i] & 0xffc00000)) &&
	    (((pehdr[i] >> 12) & 0x3ff) == ((base >> 12) & 0x3ff)))
				return valid_pgd(off);
			continue;
		}

		// 4k page
		pde &= 0xfffff000;
		if (pde > sz-4096)
			continue;
		
		pt = ptr + pde;
		pde = pt[ (base >> 12) & 0x3ff ];
		if ((pde & 1) && ((pde & 0xfffff000) == pehdr[i]))
			return valid_pgd(off);
	}
	return 0;
}

/* search PE headers, then page dir pointing to them */
void scan_mem()
{
	u32 *pdt = 0, pdt_tlen;
	u32 i;

	printf("scanning memory for pe headers\n");
	pehdr_tlen = fill_pageoffsets(&pehdr, isvalid_PEhdr);
	printf("found %lu pe headers\n", pehdr_tlen);

	pebase = (u32*)malloc(4*pehdr_tlen);
	if (!pebase) {
		perror("malloc");
		free(pehdr);
		return;
	}
	for (i=0 ; i<pehdr_tlen ; i++)
		pebase[i] = getPEbase(pehdr[i]);

	printf("scanning memory for pde pointing to them\n");
	pdt_tlen = fill_pageoffsets(&pdt, isvalid_PD2PE);
	printf("found %lu page directories\n", pdt_tlen);
	
	free(pehdr);

	for (i=0 ; i<pdt_tlen ; i++) {
		dump_map(pdt[i]);
	}
	
	free(pdt);
}

/*
// checks the validity of a pte: 0 = invalid, 1 = normal, 2 = not present
inline int valid_pte(const u32 pte)
{
	if (!(pte & 1))
		return 2;
	if ((pte & 0xfffff000) < sz)
		return 1;
	return 0;
}

int valid_pt(const u32 off)
{
	const u32 *pt = ptr + off;
	int present = 0;
	int i;
	
	for (i = 0 ; i < 1024 ; i++) {
		switch (valid_pte(pt[i])) {
			case 0:
				return -1;
			case 1:
				present++;
		}
	}
	return present;
}

void scan_pde(u32 *ptt, u32 pt_max)
{
	const u32 *pgd;
	u32 pde;
	u32 i, j, k;
	int empty;
	
	// test each page table to see if it is a page directory
	for (i = 0 ; i < pt_max ; i++) {
		empty = 1;
		
		pgd = ptr + ptt[i];

		// validate each pde
		for (j = 0 ; j < 1024 ; j++) {
			pde = pgd[j];
			
			if (!(pde & 1)) {
				// not present
				continue;
			}
			
			if (pde & 0x80) {
				// 4M table
				if (pde & 0x3FE000)
					break;
				empty = 0;
				continue;
			}
			
			if (pde & 0x40)
				break;
			
			// 4k table: check if it references another known pte
			pde &= 0xFFFFF000;
			for (k = 0 ; k < pt_max ; k++) {
				if (pde == ptt[k]) {
					empty = 0;
					break;
				}
			}
			if (k == pt_max)
				break;
		}
		if (!empty && j == 1024) {
			printf("found possible pde at %.8lX\n", ptt[i]);
			dump_map(ptt[i]);
		}
	}	
}

void scan_mem()
{
	u32 off;
	int type;
	int count = 0;

	u32 ptt_max = 1024;
	u32 ptt_cur = 0;
	u32 *ptt = (u32 *)malloc(pttable_max*4);
	if (!ptt) {
		perror("malloc");
		return;
	}

	signal(SIGALRM, sigh);
	alarm(1);
	
	for (off = 0 ; off < sz ; off += 4096) {
		if (showadvance) {
			showadvance = 0;
			fprintf(stderr, "\r          \r%d%%", (int)(((float)off)*100/(float)sz));
		}
		
		type = valid_pt(off);
		if (type >= 0) {
			count++;
			
			ptt[ptt_cur++] = off;
			
			if (ptt_cur >= ptt_max) {
				ptt_max += 1024;
				ptt = realloc(ptt, 4*ptt_max);
				if (!ptt) {
					perror("realloc");
					return;
				}
			}
		}
	}
	alarm(0);
	fprintf(stderr, "\r          \r");
	
	printf("found %d potential page tables\n", count);

	scan_pde(ptt, ptt_cur);
	
	free(ptt);
}
*/
int main(int argc, char **argv)
{
	if (argc < 2) {
		fprintf(stderr, "usage: %s <physdump file>\n", argv[0]);
		exit(1);
	}

	int fd = open(argv[1], O_RDONLY);
	if (fd == -1) {
		perror("open");
		exit(1);
	}
	
	sz = 8*1024*1024;
	struct stat stat;
	if (!fstat(fd, &stat)) {
		sz = stat.st_size;
		printf("dump size: %.8lX (%luMo)\n", sz, sz>>20);
	}
	
	ptr = mmap(0, sz, PROT_READ, MAP_PRIVATE, fd, 0);
	if (!ptr) {
		close(fd);
		perror("mmap");
		exit(1);
	}

	scan_mem();

	munmap((void *)ptr, sz);
	close(fd);
	printf("done\n");
	return 0;
}