Clean up checksec code, use --file in all cases since Binjitsu supports it now

pwndbg/commands/checksec.py

@@ -1,24 +1,27 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+import os
+import subprocess
+
 import gdb
 import pwndbg.commands
-
-import subprocess
+import pwndbg.which
 
 @pwndbg.commands.Command
-def checksec():
+def checksec(file=None):
     '''
     Prints out the binary security settings. Attempts to call the binjitsu
     checksec first, and then falls back to checksec.sh.
     '''
-    local_path = pwndbg.file.get_file(pwndbg.proc.exe)
-    try:
-        subprocess.call(['checksec', local_path])
-    except:
-        try:
-            subprocess.call(['checksec.sh', '--file', local_path])
-        except:
-            print(pwndbg.color.red(
-                'An error occurred when calling checksec. ' \
-                'Make sure the checksec binary is in your PATH.'
-            ))
+    local_path = file or pwndbg.file.get_file(pwndbg.proc.exe)
+
+    if not local_path:
+        print('No file is selected')
+
+    for program in ['checksec', 'checksec.sh']:
+        program = pwndbg.which.which(program)
+
+        if program:
+            return subprocess.call([program, '--file', local_path])
+    else:
+        print('Could not find checksec or checksec.sh in $PATH.')

pwndbg/which.py

@@ -0,0 +1,82 @@
+# This license covers everything within this project, except for a few pieces
+# of code that we either did not write ourselves or which we derived from code
+# that we did not write ourselves. These few pieces have their license specified
+# in a header, or by a file called LICENSE.txt, which will explain exactly what
+# it covers. The few relevant pieces of code are all contained inside these
+# directories:
+#
+# - pwnlib/constants/
+# - pwnlib/data/
+#
+#
+# Copyright (c) 2015 Gallopsled and Zach Riggle
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+import os
+import stat
+
+def which(name, all = False):
+    """which(name, flags = os.X_OK, all = False) -> str or str set
+
+    Works as the system command ``which``; searches $PATH for ``name`` and
+    returns a full path if found.
+
+    If `all` is :const:`True` the set of all found locations is returned, else
+    the first occurence or :const:`None` is returned.
+
+    Arguments:
+      `name` (str): The file to search for.
+      `all` (bool):  Whether to return all locations where `name` was found.
+
+    Returns:
+      If `all` is :const:`True` the set of all locations where `name` was found,
+      else the first location or :const:`None` if not found.
+
+    Example:
+      >>> which('sh')
+      '/bin/sh'
+    """
+    # If name is a path, do not attempt to resolve it.
+    if os.path.sep in name:
+        return name
+
+    isroot = os.getuid() == 0
+    out = set()
+    try:
+        path = os.environ['PATH']
+    except KeyError:
+        log.exception('Environment variable $PATH is not set')
+    for p in path.split(os.pathsep):
+        p = os.path.join(p, name)
+        if os.access(p, os.X_OK):
+            st = os.stat(p)
+            if not stat.S_ISREG(st.st_mode):
+                continue
+            # work around this issue: https://bugs.python.org/issue9311
+            if isroot and not \
+              st.st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH):
+                continue
+            if all:
+                out.add(p)
+            else:
+                return p
+    if all:
+        return out
+    else:
+        return None