sploit: cache results of external commands

rather than cacheing ELF instantiations, just cache the results of
external commands

Signed-off-by: dusoleil <howcansocksbereal@gmail.com>

5 files changed, 25 insertions, 15 deletions

diff --git a/sploit/__init__.py b/sploit/__init__.py
index 137cd35..f5a82fc 100644
--- a/sploit/__init__.py
+++ b/sploit/__init__.py
@@ -5,5 +5,6 @@ from sploit import (
         mem,
         payload,
         until,
+        util,
         rev,
 )
diff --git a/sploit/rev/elf.py b/sploit/rev/elf.py
index d9edd40..1957c15 100644
--- a/sploit/rev/elf.py
+++ b/sploit/rev/elf.py
@@ -1,16 +1,6 @@
 from sploit.rev import ldd, r2
 
-__ELF_CACHE__ = {}
-
-def ELF(path):
-    if path in __ELF_CACHE__:
-        return __ELF_CACHE__[path]
-    else:
-        elf = __ELF__(path)
-        __ELF_CACHE__[path] = elf
-        return elf
-
-class __ELF__:
+class ELF:
     def __init__(self, path):
         self.path = path
         self.sym = r2.get_elf_symbols(self.path)
diff --git a/sploit/rev/ldd.py b/sploit/rev/ldd.py
index 60306f1..d162207 100644
--- a/sploit/rev/ldd.py
+++ b/sploit/rev/ldd.py
@@ -1,9 +1,10 @@
+from sploit.util import run_cmd_cached
+
 import re
-from subprocess import run
 from collections import namedtuple as nt
 
 def get_libraries(elf):
-    out = run(['ldd',elf],capture_output=True).stdout.decode('utf-8').split('\n')[:-1]
+    out = run_cmd_cached(['ldd',elf])
     out = [re.split(r'\s+',lib)[1:] for lib in out]
     Lib = nt("Lib", "name path addr")
     out = {l[0]:Lib(l[0],l[0] if l[0][0]=='/' else l[2] if l[1]=='=>' else None,l[-1]) for l in out}
diff --git a/sploit/rev/r2.py b/sploit/rev/r2.py
index af0fe24..c7a8a65 100644
--- a/sploit/rev/r2.py
+++ b/sploit/rev/r2.py
@@ -1,12 +1,12 @@
 from sploit.mem import Symtbl
 from sploit.arch import arch
+from sploit.util import run_cmd_cached
 
 import re
-from subprocess import run
 from collections import namedtuple as nt
 
 def run_cmd(binary,cmd):
-    return run(['r2','-q','-c',cmd,'-e','scr.color=false',binary],capture_output=True).stdout.decode('utf-8').split('\n')[:-1]
+    return run_cmd_cached(['r2','-q','-c',cmd,'-e','scr.color=false',binary])
 
 def get_elf_symbols(elf):
     out = {}
diff --git a/sploit/util.py b/sploit/util.py
new file mode 100644
index 0000000..b0572a0
--- /dev/null
+++ b/sploit/util.py
@@ -0,0 +1,18 @@
+from subprocess import run
+
+def run_cmd(cmd):
+    return run(cmd,capture_output=True).stdout.decode('utf-8').split('\n')[:-1]
+
+__RUN_CACHE__ = {}
+def run_cmd_cached(cmd):
+    key = ''.join(cmd)
+    if key in __RUN_CACHE__:
+        print('cache hit')
+        return __RUN_CACHE__[key]
+    else:
+        print('cache miss')
+        result = run_cmd(cmd)
+        __RUN_CACHE__[key] = result
+        return result
+
+