From 9f620d4f337c861b456376dcd36cd907b395a1ac Mon Sep 17 00:00:00 2001
From: xyzkab <13573557+xyzkab@users.noreply.github.com>
Date: Wed, 8 Jan 2020 23:42:52 +0700
Subject: [PATCH] Update handlers: add default and load_handler method
 execution

---
 core/ssrf.py | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/core/ssrf.py b/core/ssrf.py
index 65dcea0..6ff1e39 100644
--- a/core/ssrf.py
+++ b/core/ssrf.py
@@ -7,6 +7,7 @@
 
 class SSRF(object):
     modules   = set()
+    handler   = None
     requester = None
 
     def __init__(self, args):
@@ -15,10 +16,14 @@ def __init__(self, args):
         self.load_modules()
 
         # Start a reverse shell handler
-        if args.handler:
-            handler = Handler(args.handler)
+        if args.handler == "1":
+            handler = Handler(args.lport)
             handler.start()
-        
+        else:
+            self.load_handler(args.handler)
+            handler = self.handler.exploit(args.lport)
+            handler.start()
+
         # Init a requester
         self.requester = Requester(args.reqfile, args.useragent, args.ssl)
 
@@ -40,19 +45,21 @@ def __init__(self, args):
 
         # Handling a shell
         while args.handler:
-            if handler.connected == True:
-                cmd = input("Shell> $ ")
-                if cmd == "exit":
-                    handler.kill()
-                    print("BYE !")
-                    exit()
-                handler.send_command(cmd+"\n\n")
-            else:
-                time.sleep(5)
+            handler.listen_command()
+            time.sleep(5)
 
     def load_modules(self):
         for index,name in enumerate(os.listdir("./modules")):
             location = os.path.join("./modules", name)
             if ".py" in location:
                 mymodule = SourceFileLoader(name, location).load_module()
-                self.modules.add(mymodule)
\ No newline at end of file
+                self.modules.add(mymodule)
+
+    def load_handler(self, name):
+        handler_file = "{}.py".format(name)
+        try:
+            location = os.path.join("./handlers", handler_file)
+            self.handler = SourceFileLoader(handler_file, location).load_module()
+        except Exception as e:
+            logging.error("Invalid no such handler: {}".format(name))
+            exit(1)