\# Enter a description of the change.

Implemented CPU quotas and load quotas.

Rekall will now break after it exceeds the specified number of CPU
seconds. If a load quota is given, Rekall will reduce its CPU percent
utilization to the specified level by introducing sleep periods when
idle.

Review URL: https://codereview.appspot.com/313090043 .

_version.py

@@ -94,7 +94,10 @@ def tag_version_data(version_data, version_path="version.yaml"):
         pep440 += ".rc" + version_data["rc"]
 
     if version_data.get("dev", 0):
-        pep440 += ".dev" + str(version_data["dev"])
+        # A Development release comes _before_ the main release.
+        last = version_data["version"].rsplit(".", 1)
+        version_data["version"] = "%s.%s" % (last[0], int(last[1]) + 1)
+        pep440 = version_data["version"] + ".dev" + str(version_data["dev"])
 
     version_data["pep440"] = pep440
 

rekall-agent/_version.py

@@ -94,7 +94,10 @@ def tag_version_data(version_data, version_path="version.yaml"):
         pep440 += ".rc" + version_data["rc"]
 
     if version_data.get("dev", 0):
-        pep440 += ".dev" + str(version_data["dev"])
+        # A Development release comes _before_ the main release.
+        last = version_data["version"].rsplit(".", 1)
+        version_data["version"] = "%s.%s" % (last[0], int(last[1]) + 1)
+        pep440 = version_data["version"] + ".dev" + str(version_data["dev"])
 
     version_data["pep440"] = pep440
 

rekall-agent/setup.py

@@ -41,7 +41,7 @@ def find_data_files(source):
     return result
 
 install_requires = [
-    "rekall-core >= 1.6.0rc1, < 1.7",
+    "rekall-core >= 1.6.0, < 1.7",
     "requests==2.11.1",
     "httplib2==0.9.2",
     "oauth2client==3.0.0",

rekall-core/rekall/_version.py

@@ -94,7 +94,10 @@ def tag_version_data(version_data, version_path="version.yaml"):
         pep440 += ".rc" + version_data["rc"]
 
     if version_data.get("dev", 0):
-        pep440 += ".dev" + str(version_data["dev"])
+        # A Development release comes _before_ the main release.
+        last = version_data["version"].rsplit(".", 1)
+        version_data["version"] = "%s.%s" % (last[0], int(last[1]) + 1)
+        pep440 = version_data["version"] + ".dev" + str(version_data["dev"])
 
     version_data["pep440"] = pep440
 

rekall-core/rekall/plugins/windows/common.py

@@ -247,6 +247,11 @@ def __init__(self, tag=None, **kwargs):
         self.tag_offset = self.session.profile.get_obj_offset(
             "_POOL_HEADER", "PoolTag")
 
+        if self.tag_offset == None:
+            raise RuntimeError(
+                "Unable to get PoolTag offset in _POOL_HEADER. "
+                "Is the profile correct?")
+
     def skip(self, buffer_as, offset):
         return super(PoolTagCheck, self).skip(
             buffer_as, offset + self.tag_offset)

rekall-core/rekall/quotas.py

@@ -0,0 +1,101 @@
+# Rekall
+# Copyright (C) 2016 Michael Cohen <scudette@gmail.com>
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or (at
+# your option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
+#
+
+__author__ = "Michael Cohen <scudette@google.com>"
+
+"""This file implements CPU quota limiting on the Rekall session.
+
+The code works by wrapping a session object with progress handlers
+which check for the CPU quota.
+
+Note that when Rekall is used as a library, the caller must
+deliberately wrap its own session with this module.
+"""
+import time
+
+import psutil
+
+from rekall import config
+from rekall import plugin
+
+
+config.DeclareOption(
+    "--cpu_quota", type="IntParser", group="Quotas",
+    help="Number of allocated CPU seconds Rekall is allowed to consume. "
+    "If not set, unlimited CPU time can be used.")
+
+config.DeclareOption(
+    "--load_quota", type="IntParser", group="Quotas",
+    help="The target maximal process load level (in percent).")
+
+
+def wrap_session(session, cpu_quota=None, load_quota=None):
+    """Wraps the session limiting cpu quota."""
+    if load_quota is None:
+        load_quota = session.GetParameter("load_quota")
+
+    if cpu_quota is None:
+        cpu_quota = session.GetParameter("cpu_quota")
+
+    if cpu_quota == None and load_quota == None:
+        return session
+
+    # Store the process's current CPU utilization.
+    proc = psutil.Process()
+    cpu_times = proc.cpu_times()
+    start_time = cpu_times.user + cpu_times.system
+    state = dict(last=time.time(),
+                 start_time=start_time,
+                 proc=proc)
+
+    def quota_callback(*_, **__):
+        check_quota(state, cpu_quota, load_quota)
+
+    # Register our progress dispatcher.
+    session.progress.Register("quota", quota_callback)
+    return session
+
+
+def check_quota(state, cpu_quota, load_quota):
+    """A progress callback which checks quota is not exceeded."""
+    now = time.time()
+
+    # In order to not overwhelm psutil we throttle calls to once every
+    # few ms.
+    if now + 0.5 > state["last"]:
+        state["last"] = now
+        start_time = state["start_time"]
+        proc = state["proc"]
+        cpu_times = proc.cpu_times()
+        current = cpu_times.user + cpu_times.system
+        if cpu_quota and current > start_time + cpu_quota:
+            # CPU quota exceeded.
+            raise plugin.PluginError("CPU Quota exceeded (%s Seconds)." %
+                                     (current - start_time))
+
+        if load_quota:
+            while 1:
+                current_cpu_percent = proc.cpu_percent() * 100
+
+                # If our current CPU utilization exceeds the specified
+                # limits we sleep a bit.
+                if current_cpu_percent < load_quota:
+                    break
+
+                time.sleep(0.1)

rekall-core/rekall/rekal.py

@@ -33,6 +33,7 @@
 from rekall import constants
 from rekall import plugin
 from rekall import session
+from rekall import quotas
 
 from pkg_resources import iter_entry_points
 for entry_point in iter_entry_points(group='rekall.plugins', name=None):
@@ -94,6 +95,8 @@ def global_arg_cb(global_flags, _):
         argv=argv, global_arg_cb=global_arg_cb,
         user_session=user_session)
 
+    # Install any quotas the user requested.
+    user_session = quotas.wrap_session(user_session)
     try:
         # Run the plugin with plugin specific args.
         user_session.RunPlugin(plugin_cls, **config.RemoveGlobalOptions(flags))

rekall-gui/_version.py

@@ -94,7 +94,10 @@ def tag_version_data(version_data, version_path="version.yaml"):
         pep440 += ".rc" + version_data["rc"]
 
     if version_data.get("dev", 0):
-        pep440 += ".dev" + str(version_data["dev"])
+        # A Development release comes _before_ the main release.
+        last = version_data["version"].rsplit(".", 1)
+        version_data["version"] = "%s.%s" % (last[0], int(last[1]) + 1)
+        pep440 = version_data["version"] + ".dev" + str(version_data["dev"])
 
     version_data["pep440"] = pep440
 

tools/layout_expert/layout_expert/_version.py

@@ -94,7 +94,10 @@ def tag_version_data(version_data, version_path="version.yaml"):
         pep440 += ".rc" + version_data["rc"]
 
     if version_data.get("dev", 0):
-        pep440 += ".dev" + str(version_data["dev"])
+        # A Development release comes _before_ the main release.
+        last = version_data["version"].rsplit(".", 1)
+        version_data["version"] = "%s.%s" % (last[0], int(last[1]) + 1)
+        pep440 = version_data["version"] + ".dev" + str(version_data["dev"])
 
     version_data["pep440"] = pep440
 

version.py

@@ -87,7 +87,10 @@ def tag_version_data(version_data, version_path="version.yaml"):
         pep440 += ".rc" + version_data["rc"]
 
     if version_data.get("dev", 0):
-        pep440 += ".dev" + str(version_data["dev"])
+        # A Development release comes _before_ the main release.
+        last = version_data["version"].rsplit(".", 1)
+        version_data["version"] = "%s.%s" % (last[0], int(last[1]) + 1)
+        pep440 = version_data["version"] + ".dev" + str(version_data["dev"])
 
     version_data["pep440"] = pep440