#!/usr/bin/python
# 2008/05/25
# Memory Testing Tool
# Copyright 2008, Robin H. Johnson <robbat2@gentoo.org>
# Written for bertl @ #vserver

import mmap, os, re, sys, random
from struct import pack,unpack
import array

file = sys.argv[1]
blocksize = 128
phase_pass1_permute = 1
phase_pass1_read = 1
phase_pass1_xor = 0
phase_pass1_write = 0
phase_pass2_permute = 1
phase_pass2_read = 1
phase_pass2_xor = 1
phase_pass2_write = 1

#if blocksize > mmap.PAGESIZE:
#	print "Blocksize of %d is larger than pagesize of %d" % (blocksize, mmap.PAGESIZE)
#	sys.exit(-2)

st = os.stat(file)
filesize = st.st_size
print "Filesize of %s is %d" % (file, filesize)

if filesize <= 0:
	print "Your file must have a non-zero size!"
	sys.exit(-2)

mapsize = filesize
f = open(file, "r+b")
m = mmap.mmap(f.fileno(), mapsize, prot=mmap.PROT_READ|mmap.PROT_WRITE)

blockcount = filesize/blocksize + 1
blocklist = range(0,blockcount-1)

x = [0] * blockcount
if phase_pass1_permute:
	random.shuffle(blocklist)
if phase_pass1_write or phase_pass1_read:
	for i,b in enumerate(blocklist):
		pos = b*blocksize
		if phase_pass1_read:
			x = m[pos:pos+blocksize]
		else:
			x = ['\0'] * blocksize
		if phase_pass1_write:
			if phase_pass1_xor:
				s = [ord(r) ^ random.randint(0,255) for r in x]
			else:
				s = [ord(r) for r in x]
			s = array.array('B',s).tostring()
			m[pos:pos+blocksize] = s

if phase_pass2_permute:
	random.shuffle(blocklist)
if phase_pass2_write or phase_pass2_read:
	for i,b in enumerate(blocklist):
		pos = b*blocksize
		if phase_pass2_read:
			x = m[pos:pos+blocksize]
		else:
			x = ['\0'] * blocksize
		if phase_pass2_write:
			if phase_pass2_xor:
				s = [ord(r) ^ random.randint(0,255) for r in x]
			else:
				s = [ord(r) for r in x]
			s = array.array('B',s).tostring()
			m[pos:pos+blocksize] = s

del m,f