Package Peach :: Package Agent :: Module memory
[hide private]

Source Code for Module Peach.Agent.memory

  1   
  2  ''' 
  3  Agent that monitors the amount of memory a process is utilizing.  This is  
  4  useful for detecting memory leaks within the fuzzing target 
  5   
  6  @author: Blake Frantz (blakefrantz@gmail.com) 
  7  @version: $Id: memory.py 963 2008-06-11 20:29:46Z meddingt $ 
  8  ''' 
  9   
 10  # 
 11  # Copyright (c) 2008 Michael Eddington 
 12  # 
 13  # Permission is hereby granted, free of charge, to any person obtaining a copy  
 14  # of this software and associated documentation files (the "Software"), to deal 
 15  # in the Software without restriction, including without limitation the rights  
 16  # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell  
 17  # copies of the Software, and to permit persons to whom the Software is  
 18  # furnished to do so, subject to the following conditions: 
 19  # 
 20  # The above copyright notice and this permission notice shall be included in     
 21  # all copies or substantial portions of the Software. 
 22  # 
 23  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  
 24  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  
 25  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE  
 26  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER  
 27  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
 28  # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
 29  # SOFTWARE. 
 30  # 
 31   
 32  # Authors: 
 33  #   Blake Frantz (blakefrantz@gmail.com) 
 34   
 35  # $Id: memory.py 963 2008-06-11 20:29:46Z meddingt $ 
 36   
 37   
 38  try: 
 39          import sys 
 40          import ctypes 
 41          import win32pdhutil 
 42          import win32api 
 43   
 44          sys.path.append("..") 
 45          sys.path.append("../..") 
 46   
 47          from Peach.agent import Monitor 
 48   
 49          PROCESS_VM_READ = 0x0010 
 50          PROCESS_QUERY_INFORMATION = 0x0400 
 51   
 52          DWORD = ctypes.c_ulong 
 53          SIZE_T = ctypes.c_ulong 
 54   
 55          MAX_PROCESSES = 1024 
 56          MAX_PATH = 1024 
 57   
 58          Psapi = ctypes.windll.Psapi 
 59          Kernel32 = ctypes.windll.Kernel32 
 60   
61 - class PROCESS_MEMORY_COUNTERS_EX(ctypes.Structure):
62 _fields_ = [("cb", DWORD), 63 ("PageFaultCount", DWORD), 64 ("PeakWorkingSetSize", SIZE_T), 65 ("WorkingSetSize", SIZE_T), 66 ("QuotaPeakPagedPoolUsage", SIZE_T), 67 ("QuotaPagedPoolUsage", SIZE_T), 68 ("QuotaPeakNonPagedPoolUsage", SIZE_T), 69 ("QuotaNonPagedPoolUsage", SIZE_T), 70 ("PagefileUsage", SIZE_T), 71 ("PeakPagefileUsage", SIZE_T), 72 ("PrivateUsage", SIZE_T), 73 ]
74
75 - class Memory(Monitor):
76 ''' 77 Agent that monitors the amount of memory a process is utilizing. This is 78 useful for detecting memory leaks within the fuzzing target 79 ''' 80
81 - def __init__(self, args):
82 ''' 83 Constructor. Arguments are supplied via the Peach XML 84 file. 85 86 @type args: Dictionary 87 @param args: Dictionary of parameters 88 ''' 89 90 try: 91 92 # Our name for this monitor 93 self._name = "Memory Monitor" 94 self._pid = None 95 self._processName = None 96 self._hProcess = None 97 self._internalError = False 98 self._memoryInfo = None 99 self._threshold = None 100 self._detectedFault = False 101 self._stopOnFault = False 102 103 # Report an error if no MemoryLimit and/or neither pid nor processName is defined 104 105 while 1: 106 107 if args.has_key('StopOnFault'): 108 self._stopOnFault = str(args["StopOnFault"]) 109 110 if args.has_key('MemoryLimit'): 111 self._memoryLimit = int(args['MemoryLimit']) 112 print "Memory: Memory Limit = %d" % self._memoryLimit 113 else: 114 print "Memory: No memory limit specified" 115 self._internalError = True 116 break 117 118 if args.has_key('Pid'): 119 self._pid = int(args['Pid']) 120 print "Memory: Pid = %d" % self._pid 121 122 if args.has_key('ProcessName'): 123 self._processName = str(args['ProcessName']) 124 print "Memory: Process Name = %s" % self._processName 125 126 if self._pid == None and self._processName == None: 127 print "Memory: No pid or process name provided" 128 self._internalError = True 129 break 130 131 break 132 133 except: 134 print "Memory: Caught Exception" 135 raise
136
137 - def _OpenProcess(self, pid = None):
138 139 if pid != None: 140 return Kernel32.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, 0, pid) 141 else: 142 return None
143
144 - def _CloseProcess(self, handle = None):
145 146 if handle != None: 147 Kernel32.CloseHandle(handle)
148
149 - def _GetProcessIdByName(self, name):
150 ''' 151 Try and get pid for a process by name. 152 ''' 153 154 try: 155 win32pdhutil.GetPerformanceAttributes('Process','ID Process', name) 156 except: 157 sys.stdout.write("Memory: Unable to locate process [%s]\n" % name) 158 raise 159 160 pids = win32pdhutil.FindPerformanceAttributesByName(name) 161 162 # If _my_ pid in there, remove it 163 try: 164 pids.remove(win32api.GetCurrentProcessId()) 165 except ValueError: 166 pass 167 168 return pids[0]
169
170 - def _GetProcessMemoryInfo(self, handle = None):
171 172 if handle == None: 173 return None 174 175 psmemCounters = PROCESS_MEMORY_COUNTERS_EX() 176 cb = DWORD(ctypes.sizeof(psmemCounters)) 177 b = Psapi.GetProcessMemoryInfo(handle, ctypes.byref(psmemCounters), cb) 178 179 if not b: 180 return None 181 182 dict = {} 183 184 for k, t in psmemCounters._fields_: 185 dict[k] = getattr(psmemCounters, k) 186 187 return dict
188
189 - def OnTestStarting(self):
190 ''' 191 Called right before start of test case or variation 192 ''' 193 194 # if only a process name was passed in, derive the pid from it 195 if self._processName != None: 196 self._pid = self._GetProcessIdByName(self._processName) 197 198 if self._pid == None: 199 print "Memory: OnTestStarting: Could not resolve pid" 200 self._internalError = True 201 return 202 203 self._hProcess = self._OpenProcess(self._pid) 204 205 if self._hProcess == None: 206 print "Memory: Could not open target process" 207 self._internalError = True 208 return 209 210 print "OnTestStarting: Process handle = %d" % self._hProcess 211 212 if self._hProcess == None: 213 print "Memory: Could not open target process" 214 self._internalError = True 215 return 216 217 self._memoryInfo = self._GetProcessMemoryInfo(self._hProcess) 218 219 if self._memoryInfo == None: 220 print "Memory: Could not acquire memory info" 221 self._internalError = True 222 return 223 else: 224 print "Memory Used = %d" % self._memoryInfo['PrivateUsage'] 225 226 if self._memoryInfo['PrivateUsage'] > self._memoryLimit: 227 self._detectedFault = True 228 229 if self._stopOnFault == "True": 230 print "Memory: Stopping on fault" 231 self._internalError = True
232
233 - def OnTestFinished(self):
234 ''' 235 Called right after a test case or varation 236 ''' 237 self._CloseProcess(self._hProcess) 238 self._hProcess = None
239
240 - def GetMonitorData(self):
241 ''' 242 Get any monitored data from a test case. 243 ''' 244 return {'MemoryUsed.txt': str(self._memoryInfo['PrivateUsage'])}
245
246 - def DetectedFault(self):
247 ''' 248 Check if a fault was detected. 249 ''' 250 return self._detectedFault
251
252 - def OnFault(self):
253 ''' 254 Called when a fault was detected. 255 ''' 256 pass
257
258 - def OnShutdown(self):
259 ''' 260 Called when Agent is shutting down, typically at end 261 of a test run or when a Stop-Run occurs 262 ''' 263 self._CloseProcess(self._handle)
264
265 - def StopRun(self):
266 ''' 267 Return True to force test run to fail. This 268 should return True if an unrecoverable error 269 occurs. 270 ''' 271 return self._internalError
272 273 #end 274 275 if __name__ == "__main__": 276 d = {} 277 d["MemoryLimit"] = "5000000" 278 d["ProcessName"] = "CrashableServer" 279 a = Memory(d) 280 a.OnTestStarting() 281 282 print a.DetectedFault() 283 a.OnTestFinished() 284 a.OnTestStarting() 285 286 print a.DetectedFault() 287 a.OnTestFinished() 288 except: 289 pass 290