Merge tag 'pull-sploit-error-handling' of https://github.com/Dusoleil/lib-des-gnux

Refactor exception handling and cleanup in main.py

Refactor exception handling to be simpler and easier to read/maintain
Manually call garbage collector after exec to handle some weird python
behavior.

* tag 'pull-sploit-error-handling' of https://github.com/Dusoleil/lib-des-gnux:
  Manually run garbage collection after exec
  Clean up exception handling in main.py

1 files changed, 24 insertions, 22 deletions

diff --git a/sploit/main.py b/sploit/main.py
index 6404786..77ef49e 100644
--- a/sploit/main.py
+++ b/sploit/main.py
@@ -1,6 +1,7 @@
 import argparse
 import tempfile
 import traceback
+import gc
 
 from sploit.comm import *
 
@@ -14,18 +15,15 @@ def main():
                         help='target program to exploit')
     args = parser.parse_args()
 
-    try:
-        if(len(args.target)>0):
-            if(args.daemon):
-                print("Target Given. Ignoring Daemon Flag...")
-            target(args.script, args.target)
+    if(len(args.target)>0):
+        if(args.daemon):
+            print("Target Given. Ignoring Daemon Flag...")
+        target(args.script, args.target)
+    else:
+        if(args.daemon):
+            daemon(args.script)
         else:
-            if(args.daemon):
-                daemon(args.script)
-            else:
-                pipe(args.script)
-    except KeyboardInterrupt:
-        pass
+            pipe(args.script)
 
 def daemon(script):
     print("Running in Pipe Daemon Mode...")
@@ -35,12 +33,7 @@ def daemon(script):
                 p = Pipes(tmpdir)
             except KeyboardInterrupt:
                 break
-            try:
-                runscript(script, Comm(p));
-            except KeyboardInterrupt:
-                pass
-            except:
-                traceback.print_exc()
+            runscript(script, Comm(p));
             del p
 
 def pipe(script):
@@ -52,8 +45,17 @@ def target(script, target):
     runscript(script, Comm(Process(target)));
 
 def runscript(script, comm):
-    print("Running Script...")
-    code = compile(open(script).read(), script, 'exec')
-    exec(code, {'io': comm})
-    print("Script Finished!")
-    comm.readall()
+    try:
+        print("Running Script...")
+        code = compile(open(script).read(), script, 'exec')
+        exec(code, {'io': comm})
+        print("Script Finished!")
+        comm.readall()
+        return
+    except KeyboardInterrupt:
+        pass
+    except:
+        traceback.print_exc()
+    finally:
+        gc.collect()
+    print("Script Ended Early!")