Source Code for Module rekall.quotas

  1  # Rekall 
  2  # Copyright (C) 2016 Michael Cohen <scudette@gmail.com> 
  3  # Copyright 2016 Google Inc. All Rights Reserved. 
  4  # 
  5  # This program is free software; you can redistribute it and/or modify 
  6  # it under the terms of the GNU General Public License as published by 
  7  # the Free Software Foundation; either version 2 of the License, or (at 
  8  # your option) any later version. 
  9  # 
 10  # This program is distributed in the hope that it will be useful, but 
 11  # WITHOUT ANY WARRANTY; without even the implied warranty of 
 12  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 
 13  # General Public License for more details. 
 14  # 
 15  # You should have received a copy of the GNU General Public License 
 16  # along with this program; if not, write to the Free Software 
 17  # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 
 18  # 
 19   
 20  __author__ = "Michael Cohen <scudette@google.com>" 
 21   
 22  """This file implements CPU quota limiting on the Rekall session. 
 23   
 24  The code works by wrapping a session object with progress handlers 
 25  which check for the CPU quota. 
 26   
 27  Note that when Rekall is used as a library, the caller must 
 28  deliberately wrap its own session with this module. 
 29  """ 
 30  import time 
 31   
 32  import psutil 
 33   
 34  from rekall import config 
 35  from rekall import plugin 
 36   
 37   
 38  config.DeclareOption( 
 39      "--cpu_quota", type="IntParser", group="Quotas", 
 40      help="Number of allocated CPU seconds Rekall is allowed to consume. " 
 41      "If not set, unlimited CPU time can be used.") 
 42   
 43  config.DeclareOption( 
 44      "--load_quota", type="IntParser", group="Quotas", 
 45      help="The target maximal process load level (in percent).") 
 46   
 47   
 48 -def wrap_session(session, cpu_quota=None, load_quota=None): 
 49      """Wraps the session limiting cpu quota.""" 
 50      if load_quota is None: 
 51          load_quota = session.GetParameter("load_quota") 
 52   
 53      if cpu_quota is None: 
 54          cpu_quota = session.GetParameter("cpu_quota") 
 55   
 56      if cpu_quota == None and load_quota == None: 
 57          return session 
 58   
 59      # Store the process's current CPU utilization. 
 60      proc = psutil.Process() 
 61      cpu_times = proc.cpu_times() 
 62      start_time = cpu_times.user + cpu_times.system 
 63      state = dict(last=time.time(), 
 64                   start_time=start_time, 
 65                   proc=proc) 
 66   
 67      def quota_callback(*_, **__): 
 68          check_quota(state, cpu_quota, load_quota) 
 69   
 70      # Register our progress dispatcher. 
 71      session.progress.Register("quota", quota_callback) 
 72      return session 
 73   
 74   
 75 -def check_quota(state, cpu_quota, load_quota): 
 76      """A progress callback which checks quota is not exceeded.""" 
 77      now = time.time() 
 78   
 79      # In order to not overwhelm psutil we throttle calls to once every 
 80      # few ms. 
 81      if now + 0.5 > state["last"]: 
 82          state["last"] = now 
 83          start_time = state["start_time"] 
 84          proc = state["proc"] 
 85          cpu_times = proc.cpu_times() 
 86          current = cpu_times.user + cpu_times.system 
 87          if cpu_quota and current > start_time + cpu_quota: 
 88              # CPU quota exceeded. 
 89              raise plugin.PluginError("CPU Quota exceeded (%s Seconds)." % 
 90                                       (current - start_time)) 
 91   
 92          if load_quota: 
 93              while 1: 
 94                  current_cpu_percent = proc.cpu_percent() * 100 
 95   
 96                  # If our current CPU utilization exceeds the specified 
 97                  # limits we sleep a bit. 
 98                  if current_cpu_percent < load_quota: 
 99                      break 
100   
101                  time.sleep(0.1) 
102