Merge pull request #71 from NullArray/animation

fixes an issue with `stop_animation`

Author: Ekultek

api_calls/censys.py

@@ -1,10 +1,9 @@
 import requests
+import threading
 
+import lib.settings
+from lib.output import error
 from lib.errors import AutoSploitAPIConnectionError
-from lib.output import (
-    error,
-    info
-)
 from lib.settings import (
     HOST_FILE,
     API_URLS,
@@ -28,9 +27,9 @@ def censys(self):
         """
         connect to the Censys API and pull all IP addresses from the provided query
         """
-        info("searching Censys with given query '{}'".format(self.query))
         discovered_censys_hosts = set()
         try:
+            lib.settings.start_animation("searching Censys with given query '{}'".format(self.query))
             req = requests.post(API_URLS["censys"], auth=(self.id, self.token), json={"query": self.query})
             json_data = req.json()
             for item in json_data["results"]:

api_calls/shodan.py

@@ -2,11 +2,9 @@
 
 import requests
 
+from lib.settings import start_animation
+from lib.output import error
 from lib.errors import AutoSploitAPIConnectionError
-from lib.output import (
-    error,
-    info
-)
 from lib.settings import (
     API_URLS,
     HOST_FILE,
@@ -30,7 +28,7 @@ def shodan(self):
         """
         connect to the API and grab all IP addresses associated with the provided query
         """
-        info("searching Shodan with given query '{}'".format(self.query))
+        start_animation("search Shodan with given query '{}'".format(self.query))
         discovered_shodan_hosts = set()
         try:
             req = requests.get(API_URLS["shodan"].format(query=self.query, token=self.token))

api_calls/zoomeye.py

@@ -4,11 +4,9 @@
 
 import requests
 
+from lib.settings import start_animation
 from lib.errors import AutoSploitAPIConnectionError
-from lib.output import (
-    error,
-    info
-)
+from lib.output import error
 from lib.settings import (
     API_URLS,
     HOST_FILE,
@@ -59,7 +57,7 @@ def zoomeye(self):
         connect to the API and pull all the IP addresses that are associated with the
         given query
         """
-        info("searching ZoomEye with given query '{}'".format(self.query))
+        start_animation("searching ZoomEye with given query '{}'".format(self.query))
         discovered_zoomeye_hosts = set()
         try:
             token = self.__get_auth()

autosploit.py

@@ -22,11 +22,8 @@
 import shlex
 import pickle
 import threading
-import subprocess
 import censysSearch
 import shodan
-# idk if you're going to need this since retrying is a decorator (see line 410)
-# from retrying import retry
 
 from lib.jsonize import load_exploits
 from lib.cmdline.cmd import AutoSploitParser

lib/settings.py

@@ -4,7 +4,7 @@
 import socket
 import getpass
 import tempfile
-import subprocess
+# import subprocess
 
 import psutil
 
@@ -73,7 +73,10 @@ def write_to_file(data_to_write, filename, mode="a+"):
     """
     write data to a specified file, if it exists, ask to overwrite
     """
+    global stop_animation
+
     if os.path.exists(filename):
+        stop_animation = True
         is_append = lib.output.prompt("would you like to (a)ppend or (o)verwrite the file")
         if is_append == "o":
             mode = "w"
@@ -162,7 +165,6 @@ def animation(text):
     single threaded so that it will not screw with the
     current running process
     """
-    # TODO:/ this will not stop when stop animation is True
     global stop_animation
     i = 0
     while not stop_animation:
@@ -177,5 +179,11 @@ def animation(text):
         sys.stdout.flush()
         i += 1
         time.sleep(0.1)
-    else:
-        print("\n")
+
+
+def start_animation(text):
+    import threading
+
+    t = threading.Thread(target=animation, args=(text,))
+    t.daemon = True
+    t.start()

lib/term/terminal.py

@@ -1,6 +1,5 @@
 import os
 import sys
-# import threading
 
 import lib.settings
 import lib.output
@@ -10,9 +9,6 @@
 import api_calls.censys
 
 
-stop_animation = False
-
-
 class AutoSploitTerminal(object):
 
     """
@@ -109,16 +105,14 @@ def gather_hosts(self, query, given_choice=None):
 
         option 2 must be provided
         """
-        global stop_animation
-
         choice_dict = {
             1: api_calls.shodan.ShodanAPIHook,
             2: api_calls.zoomeye.ZoomEyeAPIHook,
             3: api_calls.censys.CensysAPIHook
         }
         searching = False
         if given_choice is None:
-            lib.output.info("please choose an API to gather from (choosing two "
+            lib.output.info("please choose an API to gather from (choosing two or more "
                             "separate by comma IE; 1,2)")
             for i, api in enumerate(lib.settings.API_URLS.keys(), start=1):
                 print("{}. {}".format(i, api.title()))
@@ -128,23 +122,14 @@ def gather_hosts(self, query, given_choice=None):
         while not searching:
             try:
                 choice = int(choice)
-                # t = threading.Thread(
-                #     target=lib.settings.animation,
-                #     args=("performing lookup for provided query '{}'".format(query),)
-                # )
-                # t.daemon = True
-                # t.start()
                 if choice == 1:
                     choice_dict[choice](self.tokens["shodan"][0], query).shodan()
-                    # stop_animation = True
                     break
                 elif choice == 2:
                     choice_dict[choice](query).zoomeye()
-                    # stop_animation = True
                     break
                 elif choice == 3:
                     choice_dict[choice](self.tokens["censys"][1], self.tokens["censys"][0], query).censys()
-                    # stop_animation = True
                     break
                 else:
                     lib.output.warning("invalid option provided")