Socketio reimplemented for 5.0 (#285)

changelog

@@ -4,6 +4,9 @@
         - Breaking Changes
             - The old REST API is removed and replaced with the v2 API
             - The old WebSocket API is removed and replaced with the v2 API
+            - socketport is removed since socketio runs on the same port as the API
+            - AFTER_AGENT_STAGE2_HOOK is removed and AFTER_AGENT_CHECKIN_HOOK now runs after stage2 of checkin
+            - Removed last seen time for a user since it could cause db locking issues
         - Other changes
             - Starkiller is integrated
             - v2 API
@@ -14,6 +17,11 @@
             - JWT auth instead of basic auth
             - server-side storage of stagers
             - a new listener object is created for each listener instead of using a shared state
+            - socketio emit is now async
+            - hooks can be sync or async functions
+            - new hook for AFTER_LISTENER_CREATED_HOOK
+            - listener, agent, and task hooks now triggered via hooks
+            -
 
 1/24/2022
 ------------

empire/arguments.py

@@ -58,9 +58,6 @@
     nargs=1,
     help="Port to run the Empire RESTful API on. Defaults to 1337",
 )
-rest_group.add_argument(
-    "--socketport", type=int, nargs=1, help="Port to run socketio on. Defaults to 5000"
-)
 
 args = parent_parser.parse_args()
 

empire/server/common/agents.py

@@ -58,7 +58,6 @@
 import base64
 import json
 import os
-import sqlite3
 import string
 import threading
 
@@ -78,7 +77,7 @@
 # Empire imports
 from empire.server.database.models import TaskingStatus
 
-from . import encryption, events, helpers, messages, packets
+from . import encryption, helpers, messages, packets
 
 
 class Agents(object):
@@ -195,8 +194,6 @@ def add_agent(
 
         with SessionLocal.begin() as db:
             db.add(agent)
-            # Todo do we need to send the db session with the hook?
-            hooks.run_hooks(hooks.AFTER_AGENT_CHECKIN_HOOK, agent)
 
         # dispatch this event
         message = "[*] New agent {} checked in".format(sessionID)
@@ -1141,13 +1138,6 @@ def handle_agent_staging(
                             listener=listenerName,
                         )
 
-                        if self.mainMenu.socketio:
-                            self.mainMenu.socketio.emit(
-                                "agents/new",
-                                self.get_agent_for_socket(sessionID),
-                                broadcast=True,
-                            )
-
                         clientSessionKey = (
                             self.mainMenu.agents.get_agent_session_key_db(sessionID)
                         )
@@ -1226,13 +1216,6 @@ def handle_agent_staging(
                         listener=listenerName,
                     )
 
-                    if self.mainMenu.socketio:
-                        self.mainMenu.socketio.emit(
-                            "agents/new",
-                            self.get_agent_for_socket(sessionID),
-                            broadcast=True,
-                        )
-
                     # step 4 of negotiation -> server returns HMAC(AESn(nonce+PUBs))
                     data = "%s%s" % (nonce, serverPub.publicKey)
                     encryptedMsg = encryption.aes_encrypt_then_hmac(stagingKey, data)
@@ -1368,11 +1351,8 @@ def handle_agent_staging(
             dispatcher.send(signal, sender="agents/{}".format(sessionID))
 
             agent = self.mainMenu.agents.get_agent_for_socket(sessionID)
-            if self.mainMenu.socketio:
-                self.mainMenu.socketio.emit("agents/stage2", agent, broadcast=True)
-
             hooks.run_hooks(
-                hooks.AFTER_AGENT_STAGE2_HOOK,
+                hooks.AFTER_AGENT_CHECKIN_HOOK,
                 self.get_agent_from_name_or_session_id(sessionID),
             )
 
@@ -1620,7 +1600,7 @@ def handle_agent_response(self, sessionID, encData, update_lastseen=False):
                 # signal that this agent returned results
                 message = "[*] Agent {} returned results.".format(sessionID)
                 signal = json.dumps({"print": False, "message": message})
-                dispatcher.send(signal, sender="agents/{}".format(sessionID))
+                # dispatcher.send(signal, sender="agents/{}".format(sessionID))
 
             # return a 200/valid
             return "VALID"
@@ -1658,7 +1638,18 @@ def process_agent_packet(
                 "event_type": "result",
             }
         )
-        dispatcher.send(signal, sender="agents/{}".format(session_id))
+        # dispatcher.send(signal, sender="agents/{}".format(session_id))
+
+        tasking = (
+            db.query(models.Tasking)
+            .filter(
+                and_(
+                    models.Tasking.id == task_id,
+                    models.Tasking.agent_id == session_id,
+                )
+            )
+            .first()
+        )
 
         # insert task results into the database, if it's not a file
         if (
@@ -1667,17 +1658,6 @@ def process_agent_packet(
             not in ["TASK_DOWNLOAD", "TASK_CMD_JOB_SAVE", "TASK_CMD_WAIT_SAVE"]
             and data is not None
         ):
-            # Update result with data
-            tasking = (
-                db.query(models.Tasking)
-                .filter(
-                    and_(
-                        models.Tasking.id == task_id,
-                        models.Tasking.agent_id == session_id,
-                    )
-                )
-                .first()
-            )
             # add keystrokes to database
             if "function Get-Keystrokes" in tasking.input:
                 key_log_task_id = tasking.id
@@ -1702,31 +1682,6 @@ def process_agent_packet(
 
             db.flush()
 
-            hooks.run_hooks(hooks.AFTER_TASKING_RESULT_HOOK, tasking)
-
-            if (
-                self.mainMenu.socketio
-                and "function Get-Keystrokes" not in tasking.input
-            ):
-                result_string = tasking.output
-                if isinstance(result_string, bytes):
-                    result_string = tasking.output.decode("UTF-8")
-
-                self.mainMenu.socketio.emit(
-                    f"agents/{session_id}/task",
-                    {
-                        "taskID": tasking.id,
-                        "command": tasking.input,
-                        "results": result_string,
-                        "user_id": tasking.user_id,
-                        "created_at": tasking.created_at,
-                        "updated_at": tasking.updated_at,
-                        "username": tasking.user.username,
-                        "agent": tasking.agent_id,
-                    },
-                    broadcast=True,
-                )
-
         # TODO: for heavy traffic packets, check these first (i.e. SOCKS?)
         #       so this logic is skipped
 
@@ -1852,18 +1807,6 @@ def process_agent_packet(
                 file_data = helpers.decode_base64(data.encode("UTF-8"))
                 name = self.get_agent_name_db(session_id)
 
-                # this whole big block could probably be cleaned up so we don't have to redo this tasking lookup.
-                tasking = (
-                    db.query(models.Tasking)
-                    .filter(
-                        and_(
-                            models.Tasking.id == task_id,
-                            models.Tasking.agent_id == session_id,
-                        )
-                    )
-                    .first()
-                )
-                # TODO VR: Need to handle all the other tasking types that are creating their own db sessions.
                 if index == "0":
                     self.save_file(name, path, file_data, filesize, tasking, db)
                 else:
@@ -1967,32 +1910,13 @@ def process_agent_packet(
             msg = "[+] Output saved to .%s" % (final_save_path)
             self.save_agent_log(session_id, msg)
 
-            # Retrieve tasking data
-            tasking: models.Tasking = (
-                db.query(models.Tasking)
-                .filter(
-                    and_(
-                        models.Tasking.id == task_id,
-                        models.Tasking.agent_id == session_id,
-                    )
-                )
-                .first()
-            )
-
             # attach file to tasking
             download = models.Download(
                 location=final_save_path, size=os.path.getsize(final_save_path)
             )
             db.add(download)
             db.flush()
             tasking.downloads.append(download)
-            # todo vr: what is this used for.
-            # Send server notification for saving file
-            # self.mainMenu.socketio.emit(f'agents/{session_id}/task', {
-            #     'taskID': tasking.id, 'command': tasking.input,
-            #     'results': msg, 'user_id': tasking.user_id,
-            #     'created_at': tasking.created_at, 'updated_at': tasking.updated_at,
-            #     'username': tasking.user.username, 'agent': tasking.agent}, broadcast=True)
 
         elif response_name == "TASK_CMD_JOB":
             # check if this is the powershell keylogging task, if so, write output to file instead of screen
@@ -2153,3 +2077,5 @@ def process_agent_packet(
                     "[!] Unknown response %s from %s" % (response_name, session_id)
                 )
             )
+
+        hooks.run_hooks(hooks.AFTER_TASKING_RESULT_HOOK, tasking)

empire/server/common/empire.py

@@ -9,11 +9,13 @@
 """
 from __future__ import absolute_import, print_function
 
+import asyncio
 import json
 import os
 import threading
 import time
 from builtins import input, str
+from socket import SocketIO
 from typing import Optional
 
 from prompt_toolkit import HTML, PromptSession
@@ -121,6 +123,7 @@ def __init__(self, args=None):
         self.autoRuns = {}
         self.directory = {}
 
+        self.get_directories()
         message = "[*] Empire starting up..."
         signal = json.dumps({"print": True, "message": message})
         # dispatcher.send(signal, sender="empire")
@@ -181,9 +184,11 @@ def plugin_socketio_message(self, plugin_name, msg):
         if self.args.debug is not None:
             print(helpers.color(msg))
         if self.socketio:
-            self.socketio.emit(
-                f"plugins/{plugin_name}/notifications",
-                {"message": msg, "plugin_name": plugin_name},
+            asyncio.run(
+                self.socketio.emit(
+                    f"plugins/{plugin_name}/notifications",
+                    {"message": msg, "plugin_name": plugin_name},
+                )
             )
 
     def shutdown(self):

empire/server/common/hooks.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Callable, Dict
 
 from empire.server.common import helpers
@@ -12,6 +13,10 @@ class Hooks(object):
     Potential future addition: Filters. Add a filter to an event to do some synchronous modification to the data.
     """
 
+    # This event is triggered after the creation of a listener.
+    # Its arguments are (listener: models.Listener)
+    AFTER_LISTENER_CREATED_HOOK = "after_listener_created_hook"
+
     # This event is triggered after the tasking is written to the database.
     # Its arguments are (tasking: models.Tasking)
     AFTER_TASKING_HOOK = "after_tasking_hook"
@@ -26,14 +31,10 @@ class Hooks(object):
     # Its arguments are (tasking: models.Tasking) where tasking is the db record.
     AFTER_TASKING_RESULT_HOOK = "after_tasking_result_hook"
 
-    # This event is triggered after the agent has checked in and a record written to the database.
-    # It has one argument (agent: models.Agent)
-    AFTER_AGENT_CHECKIN_HOOK = "after_agent_checkin_hook"
-
     # This event is triggered after the agent has completed the stage2 of the checkin process,
     # and the sysinfo has been written to the database.
     # It has one argument (agent: models.Agent)
-    AFTER_AGENT_STAGE2_HOOK = "after_agent_stage2_hook"
+    AFTER_AGENT_CHECKIN_HOOK = "after_agent_checkin_hook"
 
     def __init__(self):
         self.hooks: Dict[str, Dict[str, Callable]] = {}
@@ -77,6 +78,7 @@ def unregister_filter(self, name: str, event: str = None):
         if name in self.filters.get(event, {}):
             self.filters[event].pop(name)
 
+    # todo can this be made async?
     def run_hooks(self, event: str, *args):
         """
         Run all hooks for a hook type.
@@ -86,7 +88,18 @@ def run_hooks(self, event: str, *args):
             return
         for hook in self.hooks.get(event, {}).values():
             try:
-                hook(*args)
+                if asyncio.iscoroutinefunction(hook):
+                    try:  # https://stackoverflow.com/a/61331974/
+                        loop = asyncio.get_running_loop()
+                    except RuntimeError:
+                        loop = None
+
+                    if loop and loop.is_running():
+                        loop.create_task(hook(*args))
+                    else:
+                        asyncio.run(hook(*args))
+                else:
+                    hook(*args)
             except Exception as e:
                 print(helpers.color(f"[!] Hook {hook} failed: {e}"))