Fixed Rekall colors in Linux. Changed the --nocolors flag to be a cho…

…ice colors=yes no or auto.

Can view moo plugin with
rekal moo --color yes | less -RS

[email protected]

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

rekall/plugins/tools/ipython.py

@@ -147,7 +147,7 @@ class SessionMod(plugin.Command):
 
     Any session parameters can be set here. For example:
 
-    smod nocolors=True, paging_limit=10, pager="less"
+    smod colors="no", paging_limit=10, pager="less"
 
     """
     __name = "smod"

rekall/plugins/tools/mspdb.py

@@ -203,13 +203,13 @@ def FetchPDBFile(self, pdb_filename, guid):
                         # This should already be installed on windows systems.
                         subprocess.check_call(
                             ["expand", compressed_output_file, output_file],
-                            cwd=self.dump_dir)
+                            cwd=temp_dir)
                     else:
                         # In Linux we just hope the cabextract program was
                         # installed.
                         subprocess.check_call(
                             ["cabextract", compressed_output_file],
-                            cwd=self.dump_dir)
+                            cwd=temp_dir)
 
                 except subprocess.CalledProcessError:
                     raise RuntimeError(

rekall/plugins/tools/profile_tool.py

@@ -744,7 +744,7 @@ def fetch_and_parse(self, module_name=None, guid=None):
         profile_name = "{0}/GUID/{1}".format(module_name, guid)
 
         # Get the first repository to write to.
-        repository = self.session.repository_managers.values()[0]
+        repository = self.session.repository_managers[0][1]
         if module_name != "nt":
             data = self._fetch_and_parse(module_name, guid)
             return repository.StoreData(profile_name, data)
@@ -761,8 +761,8 @@ def fetch_and_parse(self, module_name=None, guid=None):
                     "profile repository.", profile_name)
 
                 return repository.StoreData(profile_name, data)
-            except IOError:
-                pass
+            except IOError, e:
+                logging.error("Error: %s", e)
 
         raise IOError("Profile not found")
 

rekall/session.py

@@ -392,15 +392,12 @@ def __init__(self, **kwargs):
         # on this object.
         self.plugins = PluginContainer(self)
 
-        # These are the IO managers that are used to fetch profiles from the
-        # profile repository.
-        self.repository_managers = {}
-
         # When the session switches process context we store various things in
         # this cache, so we can restore the context quickly. The cache is
         # indexed by the current process_context which can be found from
         # session.GetParameter("process_context").
         self.context_cache = {}
+        self._repository_managers = []
 
         # Store user configurable attributes here. These will be read/written to
         # the configuration file.
@@ -411,6 +408,25 @@ def __init__(self, **kwargs):
             for k, v in kwargs.items():
                 self.state.Set(k, v)
 
+    @property
+    def repository_managers(self):
+        """The IO managers that are used to fetch profiles from the profile
+        repository.
+
+        """
+        if self._repository_managers:
+            return self._repository_managers
+
+        # The profile path is specified in search order.
+        repository_path = (self.state.Get("repository_path") or
+                           self.state.Get("profile_path") or [])
+
+        for path in repository_path:
+            self._repository_managers.append(
+                (path, io_manager.Factory(path, session=self)))
+
+        return self._repository_managers
+
     def __enter__(self):
         # Allow us to update the state context manager.
         self.state.__enter__()
@@ -684,19 +700,9 @@ def LoadProfile(self, name, use_cache=True):
 
         # Traverse the profile path until one works.
         if not result:
-            # The profile path is specified in search order.
-            repository_path = (self.state.Get("repository_path") or
-                               self.state.Get("profile_path") or [])
-
             # Add the last supported repository as the last fallback path.
-            for path in repository_path:
+            for path, manager in self.repository_managers:
                 try:
-                    if path not in self.repository_managers:
-                        self.repository_managers[path] = io_manager.Factory(
-                            path, session=self)
-
-                    manager = self.repository_managers[path]
-
                     # The inventory allows us to fail fetching the profile
                     # quickly - without making the round trip.
                     if not manager.CheckInventory(name):

rekall/ui/text.py

@@ -52,8 +52,10 @@
     help="The number of output lines before we invoke the pager.")
 
 config.DeclareOption(
-    "--nocolors", default=False, type="Boolean", group="Interface",
-    help="If set suppress outputting colors.")
+    "--colors", default="auto", type="Choices",
+    choices=["auto", "yes", "no"],
+    group="Interface", help="Color control. If set to auto only output "
+    "colors when connected to a terminal.")
 
 
 HIGHLIGHT_SCHEME = dict(
@@ -152,7 +154,7 @@ def __init__(self, session=None, term_fd=None):
 
         self.colorizer = Colorizer(
             self.term_fd,
-            nocolor=self.session.GetParameter("nocolors"),
+            color=self.session.GetParameter("colors"),
         )
 
     def __enter__(self):
@@ -258,29 +260,32 @@ class Colorizer(object):
 
     terminal_capable = False
 
-    def __init__(self, stream, nocolor=False):
+    def __init__(self, stream, color="auto"):
         """Initialize a colorizer.
 
         Args:
           stream: The stream to write to.
 
-          nocolor: If True we suppress using colors, even if the output stream
+          color: If "no" we suppress using colors, even if the output stream
              can support them.
         """
         if stream is None:
             stream = sys.stdout
 
-        if nocolor:
+        # We currently do not support Win32 colors.
+        if curses is None or color == "no":
             self.terminal_capable = False
-            return
 
-        try:
-            if curses and stream.isatty():
-                curses.setupterm()
+        elif color == "yes":
+            self.terminal_capable = True
 
-                self.terminal_capable = True
-        except AttributeError:
-            pass
+        elif color == "auto":
+            try:
+                if curses and stream.isatty():
+                    curses.setupterm()
+                    self.terminal_capable = True
+            except AttributeError:
+                pass
 
     def tparm(self, capabilities, *args):
         """A simplified version of tigetstr without terminal delays."""
@@ -858,14 +863,16 @@ def highlight_line(self, line, offset, last_highlight):
                 if fg is not None:
                     if isinstance(fg, basestring):
                         fg = self.colorizer.COLOR_MAP[fg]
+
                     escape_seq += self.colorizer.tparm(
-                        ["setf", "setaf"], fg)
+                        ["setaf", "setf"], fg)
 
                 if bg is not None:
                     if isinstance(bg, basestring):
                         bg = self.colorizer.COLOR_MAP[bg]
+
                     escape_seq += self.colorizer.tparm(
-                        ["setb", "setab"], bg)
+                        ["setab", "setb"], bg)
 
                 if bold:
                     escape_seq += self.colorizer.tparm(["bold"])
@@ -1228,7 +1235,7 @@ def __init__(self, tablesep=" ", output=None, mode="a+b", fd=None,
 
         self.colorizer = Colorizer(
             self.fd,
-            nocolor=self.session.GetParameter("nocolors"))
+            color=self.session.GetParameter("colors"))
 
     def section(self, name=None, width=50):
         if name is None: