Merge pull request #29 from chvancooten/feature/wip-refactor-and-fix

Feature/wip refactor and fix

NimPlant.py

@@ -11,10 +11,10 @@
 
 import os
 import random
+import sys
 import time
 import toml
-from pathlib import Path
-from client.dist.srdi.ShellcodeRDI import *
+from client.dist.srdi.ShellcodeRDI import ConvertToShellcode, HashFunctionName
 
 
 def print_banner():
@@ -62,8 +62,8 @@ def print_usage():
     )
 
 
-def getXorKey(force_new=False):
-    if os.path.isfile(".xorkey") and force_new == False:
+def get_xor_key(force_new=False):
+    if os.path.isfile(".xorkey") and not force_new:
         file = open(".xorkey", "r")
         xor_key = int(file.read())
     else:
@@ -72,8 +72,8 @@ def getXorKey(force_new=False):
             "NOTE: Make sure the '.xorkey' file matches if you run the server elsewhere!"
         )
         xor_key = random.randint(0, 2147483647)
-        file = open(".xorkey", "w")
-        file.write(str(xor_key))
+        with open(".xorkey", "w") as file:
+            file.write(str(xor_key))
 
     return xor_key
 
@@ -123,10 +123,10 @@ def compile_nim_debug(binary_type, _):
 
 def compile_nim(binary_type, xor_key, debug=False):
     # Parse config for certain compile-time tasks
-    configPath = os.path.abspath(
+    config_path = os.path.abspath(
         os.path.join(os.path.dirname(sys.argv[0]), "config.toml")
     )
-    config = toml.load(configPath)
+    config = toml.load(config_path)
 
     # Enable Ekko sleep mask if defined in config.toml, but only for self-contained executables
     sleep_mask_enabled = config["nimplant"]["sleepMask"]
@@ -224,16 +224,16 @@ def compile_nim(binary_type, xor_key, debug=False):
                 binary = "all"
 
             if "rotatekey" in sys.argv:
-                xor_key = getXorKey(True)
+                xor_key = get_xor_key(True)
             else:
-                xor_key = getXorKey()
+                xor_key = get_xor_key()
 
             compile_implant(implant, binary, xor_key)
 
             print("Done compiling! You can find compiled binaries in 'client/bin/'.")
 
         elif sys.argv[1] == "server":
-            xor_key = getXorKey()
+            xor_key = get_xor_key()
             from server.server import main
 
             try:

server/api/server.py

@@ -1,16 +1,24 @@
-from ..util.commands import getCommands, handleCommand
-from ..util.config import config
-from ..util.crypto import randString
-from ..util.func import exitServer
-from ..util.nimplant import np_server
+import os
+from threading import Thread
 
+import flask
 from flask_cors import CORS
 from gevent.pywsgi import WSGIServer
-from server.util.db import *
-from threading import Thread
 from werkzeug.utils import secure_filename
-import flask
-import os
+
+from server.util.commands import get_commands, handle_command
+from server.util.config import config
+from server.util.crypto import random_string
+from server.util.func import exit_server
+from server.util.nimplant import np_server
+from server.util.db import (
+    db_get_nimplant_console,
+    db_get_nimplant_details,
+    db_get_nimplant_info,
+    db_get_server_console,
+    db_get_server_info,
+)
+
 
 # Parse server configuration
 server_ip = config["server"]["ip"]
@@ -23,22 +31,25 @@
     static_folder="../web/static",
     template_folder="../web",
 )
-app.secret_key = randString(32)
+app.secret_key = random_string(32)
+
 
 # Define the API server
 def api_server():
     # Get available commands
     @app.route("/api/commands", methods=["GET"])
-    def get_commands():
-        return flask.jsonify(getCommands()), 200
+    def get_command_list():
+        return flask.jsonify(get_commands()), 200
 
     # Get download information
     @app.route("/api/downloads", methods=["GET"])
     def get_downloads():
         try:
-            downloadsPath = os.path.abspath(f"server/downloads/server-{np_server.guid}")
+            downloads_path = os.path.abspath(
+                f"server/downloads/server-{np_server.guid}"
+            )
             res = []
-            items = os.scandir(downloadsPath)
+            items = os.scandir(downloads_path)
             for item in items:
                 if item.is_dir() and item.name.startswith("nimplant-"):
                     downloads = os.scandir(item.path)
@@ -62,19 +73,19 @@ def get_downloads():
     @app.route("/api/downloads/<nimplant_guid>/<filename>", methods=["GET"])
     def get_download(nimplant_guid, filename):
         try:
-            downloadsPath = os.path.abspath(
+            downloads_path = os.path.abspath(
                 f"server/downloads/server-{np_server.guid}/nimplant-{nimplant_guid}"
             )
             return flask.send_from_directory(
-                downloadsPath, filename, as_attachment=True
+                downloads_path, filename, as_attachment=True
             )
         except FileNotFoundError:
             return flask.jsonify("File not found"), 404
 
     # Get server configuration
     @app.route("/api/server", methods=["GET"])
     def get_server_info():
-        return flask.jsonify(dbGetServerInfo(np_server.guid)), 200
+        return flask.jsonify(db_get_server_info(np_server.guid)), 200
 
     # Get the last X lines of console history
     @app.route("/api/server/console", methods=["GET"])
@@ -85,12 +96,12 @@ def get_server_console(lines="1000", offset="0"):
         if not lines.isnumeric() or not offset.isnumeric():
             return flask.jsonify("Invalid parameters"), 400
 
-        return flask.jsonify(dbGetServerConsole(np_server.guid, lines, offset)), 200
+        return flask.jsonify(db_get_server_console(np_server.guid, lines, offset)), 200
 
     # Exit the server
     @app.route("/api/server/exit", methods=["POST"])
     def post_exit_server():
-        Thread(target=exitServer).start()
+        Thread(target=exit_server).start()
         return flask.jsonify("Exiting server..."), 200
 
     # Upload a file to the server's "uploads" folder
@@ -117,13 +128,13 @@ def post_upload():
     # Get all active nimplants with basic information
     @app.route("/api/nimplants", methods=["GET"])
     def get_nimplants():
-        return flask.jsonify(dbGetNimplantInfo(np_server.guid)), 200
+        return flask.jsonify(db_get_nimplant_info(np_server.guid)), 200
 
     # Get a specific nimplant with its details
     @app.route("/api/nimplants/<guid>", methods=["GET"])
     def get_nimplant(guid):
-        if np_server.getNimplantByGuid(guid):
-            return flask.jsonify(dbGetNimplantDetails(guid)), 200
+        if np_server.get_nimplant_by_guid(guid):
+            return flask.jsonify(db_get_nimplant_details(guid)), 200
         else:
             return flask.jsonify("Invalid Nimplant GUID"), 404
 
@@ -136,31 +147,31 @@ def get_nimplant_console(guid, lines="1000", offset="0"):
         if not lines.isnumeric() or not offset.isnumeric():
             return flask.jsonify("Invalid parameters"), 400
 
-        if np_server.getNimplantByGuid(guid):
-            return flask.jsonify(dbGetNimplantConsole(guid, lines, offset)), 200
+        if np_server.get_nimplant_by_guid(guid):
+            return flask.jsonify(db_get_nimplant_console(guid, lines, offset)), 200
         else:
             return flask.jsonify("Invalid Nimplant GUID"), 404
 
     # Issue a command to a specific nimplant
     @app.route("/api/nimplants/<guid>/command", methods=["POST"])
     def post_nimplant_command(guid):
-        np = np_server.getNimplantByGuid(guid)
+        np = np_server.get_nimplant_by_guid(guid)
         data = flask.request.json
         command = data["command"]
 
         if np and command:
-            handleCommand(command, np)
+            handle_command(command, np)
             return flask.jsonify(f"Command queued: {command}"), 200
         else:
             return flask.jsonify("Invalid Nimplant GUID or command"), 404
 
     # Exit a specific nimplant
     @app.route("/api/nimplants/<guid>/exit", methods=["POST"])
     def post_nimplant_exit(guid):
-        np = np_server.getNimplantByGuid(guid)
+        np = np_server.get_nimplant_by_guid(guid)
 
         if np:
-            handleCommand("kill", np)
+            handle_command("kill", np)
             return flask.jsonify("Instructed Nimplant to exit"), 200
         else:
             return flask.jsonify("Invalid Nimplant GUID"), 404
@@ -183,7 +194,7 @@ def nimplantdetails():
         return flask.render_template("nimplants/details.html")
 
     @app.route("/<path:path>")
-    def catch_all(path):
+    def catch_all(_path):
         return flask.render_template("404.html")
 
     @app.errorhandler(Exception)

server/server.py

@@ -10,57 +10,74 @@
 import threading
 import time
 
-from .api.server import api_server, server_ip, server_port
-from .util.db import initDb, dbInitNewServer, dbPreviousServerSameConfig
-from .util.func import nimplantPrint, periodicNimplantChecks
-from .util.listener import *
-from .util.nimplant import *
-from .util.input import *
+from server.api.server import (
+    api_server,
+    server_ip,
+    server_port,
+)
+from server.util.db import (
+    initialize_database,
+    db_initialize_server,
+    db_is_previous_server_same_config,
+)
+from server.util.func import (
+    exit_server_console,
+    nimplant_print,
+    periodic_nimplant_checks,
+)
+from server.util.listener import (
+    flask_listener,
+    listener_type,
+    listener_ip,
+    listener_port,
+)
+from server.util.nimplant import np_server
+from server.util.input import prompt_user_for_command
 
 
 def main(xor_key=459457925, name=""):
     # Initialize the SQLite database
-    initDb()
+    initialize_database()
 
     # Restore the previous server session if config remains unchanged
     # Otherwise, initialize a new server session
-    if dbPreviousServerSameConfig(np_server, xor_key):
-        nimplantPrint("Existing server session found, restoring...")
-        np_server.restoreServerFromDb()
+    if db_is_previous_server_same_config(np_server, xor_key):
+        nimplant_print("Existing server session found, restoring...")
+        np_server.restore_from_db()
     else:
-        np_server.initNewServer(name, xor_key)
-        dbInitNewServer(np_server)
+        np_server.initialize(name, xor_key)
+        db_initialize_server(np_server)
 
     # Start daemonized Flask server for API communications
     t1 = threading.Thread(name="Listener", target=api_server)
-    t1.setDaemon(True)
+    t1.daemon = True
     t1.start()
-    nimplantPrint(f"Started management server on http://{server_ip}:{server_port}.")
+    nimplant_print(f"Started management server on http://{server_ip}:{server_port}.")
 
     # Start another thread for NimPlant listener
-    t2 = threading.Thread(name="Listener", target=flaskListener, args=(xor_key,))
-    t2.setDaemon(True)
+    t2 = threading.Thread(name="Listener", target=flask_listener, args=(xor_key,))
+    t2.daemon = True
     t2.start()
-    nimplantPrint(
-        f"Started NimPlant listener on {listenerType.lower()}://{listenerIp}:{listenerPort}. CTRL-C to cancel waiting for NimPlants."
+    nimplant_print(
+        f"Started NimPlant listener on {listener_type.lower()}://{listener_ip}:{listener_port}. CTRL-C to cancel waiting for NimPlants."
     )
 
     # Start another thread to periodically check if nimplants checked in on time
-    t3 = threading.Thread(name="Listener", target=periodicNimplantChecks)
-    t3.setDaemon(True)
+    t3 = threading.Thread(name="Listener", target=periodic_nimplant_checks)
+    t3.daemon = True
     t3.start()
 
     # Run the console as the main thread
     while True:
         try:
-            if np_server.isActiveNimplantSelected():
-                promptUserForCommand()
-            elif np_server.containsActiveNimplants():
-                np_server.selectNextActiveNimplant()
+            if np_server.is_active_nimplant_selected():
+                prompt_user_for_command()
+            elif np_server.has_active_nimplants():
+                np_server.get_next_active_nimplant()
             else:
                 pass
 
             time.sleep(0.5)
 
         except KeyboardInterrupt:
-            exitServerConsole()
+            exit_server_console()