increase coverage by optimizing conftest and removed unused migration…

… scripts. add tag filters to credentials (BC-SECURITY#728)

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 -   Upgrade Pydantic to v2 (@Vinnybod)
 -   Update common FastAPI Dependencies to use 'Annotated' types for simpler code (@Vinnybod)
+-   Add tags search to credentials endpoints (@Vinnybod)
+-   Remove unused migration scripts (@Vinnybod)
+-   Simplify TestClient setup (@Vinnybod)
 
 ## [5.8.0] - 2023-11-06
 

empire/server/api/app.py

@@ -66,7 +66,9 @@ def load_starkiller(v2App, ip, port):
         log.info(f"Starkiller served at http://localhost:{port}/index.html")
 
 
-def initialize(secure: bool = False, ip: str = "0.0.0.0", port: int = 1337):
+def initialize(
+    secure: bool = False, ip: str = "0.0.0.0", port: int = 1337, run: bool = True
+):
     # Not pretty but allows us to use main_menu by delaying the import
     from empire.server.api.v2.agent import agent_api, agent_file_api, agent_task_api
     from empire.server.api.v2.bypass import bypass_api
@@ -151,23 +153,26 @@ def shutdown_event():
 
     cert_path = os.path.abspath("./empire/server/data/")
 
-    if not secure:
-        uvicorn.run(
-            v2App,
-            host=ip,
-            port=port,
-            log_config=None,
-            lifespan="on",
-            # log_level="info",
-        )
-    else:
-        uvicorn.run(
-            v2App,
-            host=ip,
-            port=port,
-            log_config=None,
-            lifespan="on",
-            ssl_keyfile=f"{cert_path}/empire-priv.key",
-            ssl_certfile=f"{cert_path}/empire-chain.pem",
-            # log_level="info",
-        )
+    if run:
+        if not secure:
+            uvicorn.run(
+                v2App,
+                host=ip,
+                port=port,
+                log_config=None,
+                lifespan="on",
+                # log_level="info",
+            )
+        else:
+            uvicorn.run(
+                v2App,
+                host=ip,
+                port=port,
+                log_config=None,
+                lifespan="on",
+                ssl_keyfile=f"{cert_path}/empire-priv.key",
+                ssl_certfile=f"{cert_path}/empire-chain.pem",
+                # log_level="info",
+            )
+
+    return v2App

empire/server/api/v2/credential/credential_api.py

@@ -1,4 +1,4 @@
-from fastapi import Depends, HTTPException
+from fastapi import Depends, HTTPException, Query
 from starlette.responses import Response
 from starlette.status import HTTP_204_NO_CONTENT
 
@@ -14,6 +14,7 @@
 from empire.server.api.v2.shared_dependencies import CurrentSession
 from empire.server.api.v2.shared_dto import BadRequestResponse, NotFoundResponse
 from empire.server.api.v2.tag import tag_api
+from empire.server.api.v2.tag.tag_dto import TagStr
 from empire.server.core.db import models
 from empire.server.server import main
 
@@ -54,11 +55,12 @@ async def read_credentials(
     db: CurrentSession,
     search: str | None = None,
     credtype: str | None = None,
+    tags: list[TagStr] | None = Query(None),
 ):
     credentials = list(
         map(
             lambda x: domain_to_dto_credential(x),
-            credential_service.get_all(db, search, credtype),
+            credential_service.get_all(db, search, credtype, tags),
         )
     )
 

empire/server/core/credential_service.py

@@ -10,7 +10,9 @@ def __init__(self, main_menu):
         self.main_menu = main_menu
 
     @staticmethod
-    def get_all(db: Session, search: str = None, credtype: str = None):
+    def get_all(
+        db: Session, search: str = None, credtype: str = None, tags: list[str] = None
+    ):
         query = db.query(models.Credential)
 
         if search:
@@ -23,6 +25,15 @@ def get_all(db: Session, search: str = None, credtype: str = None):
                 )
             )
 
+        if tags:
+            tags_split = [tag.split(":", 1) for tag in tags]
+            query = query.join(models.Credential.tags).filter(
+                and_(
+                    models.Tag.name.in_([tag[0] for tag in tags_split]),
+                    models.Tag.value.in_([tag[1] for tag in tags_split]),
+                )
+            )
+
         if credtype:
             query = query.filter(models.Credential.credtype == credtype)
 

empire/test/conftest.py

@@ -6,7 +6,6 @@
 from pathlib import Path
 
 import pytest
-from fastapi import FastAPI
 from starlette.testclient import TestClient
 
 from empire.client.src.utils.data_util import get_random_string
@@ -50,48 +49,14 @@ def client():
     args = arguments.parent_parser.parse_args()
 
     import empire.server.server
+    from empire.server.api.app import initialize
     from empire.server.common.empire import MainMenu
 
     empire.server.server.main = MainMenu(args)
 
-    from empire.server.api.v2.agent import agent_api, agent_file_api, agent_task_api
-    from empire.server.api.v2.bypass import bypass_api
-    from empire.server.api.v2.credential import credential_api
-    from empire.server.api.v2.download import download_api
-    from empire.server.api.v2.host import host_api, process_api
-    from empire.server.api.v2.listener import listener_api, listener_template_api
-    from empire.server.api.v2.meta import meta_api
-    from empire.server.api.v2.module import module_api
-    from empire.server.api.v2.obfuscation import obfuscation_api
-    from empire.server.api.v2.plugin import plugin_api, plugin_task_api
-    from empire.server.api.v2.profile import profile_api
-    from empire.server.api.v2.stager import stager_api, stager_template_api
-    from empire.server.api.v2.tag import tag_api
-    from empire.server.api.v2.user import user_api
-
-    v2App = FastAPI()
-    v2App.include_router(listener_template_api.router)
-    v2App.include_router(listener_api.router)
-    v2App.include_router(stager_template_api.router)
-    v2App.include_router(stager_api.router)
-    v2App.include_router(agent_task_api.router)
-    v2App.include_router(agent_file_api.router)
-    v2App.include_router(agent_api.router)
-    v2App.include_router(module_api.router)
-    v2App.include_router(bypass_api.router)
-    v2App.include_router(obfuscation_api.router)
-    v2App.include_router(profile_api.router)
-    v2App.include_router(plugin_api.router)
-    v2App.include_router(plugin_task_api.router)
-    v2App.include_router(credential_api.router)
-    v2App.include_router(host_api.router)
-    v2App.include_router(user_api.router)
-    v2App.include_router(process_api.router)
-    v2App.include_router(download_api.router)
-    v2App.include_router(meta_api.router)
-    v2App.include_router(tag_api.router)
-
-    yield TestClient(v2App)
+    app = initialize(ip="localhost", run=False)
+
+    yield TestClient(app)
 
     from empire.server.server import main