Add agent-id support to distinguish collector instances

projects/cli/cli/cobaltstrike_connector/cobaltstrike_connector.py

@@ -73,6 +73,7 @@ async def setup_cobaltstrike_monitor(cobaltstrike_config: CobaltStrikeConfig, co
             config.cache_db_path,
             nemesis_client,
             cobaltstrike_config.project,
+            agent_id=cobaltstrike_config.agent_id,
             cobalt_strike=cobaltstrike_client,
         )
 

projects/cli/cli/cobaltstrike_connector/download_processor.py

@@ -42,6 +42,7 @@ def __init__(
         db_path: Path,
         nemesis: NemesisClient,
         project: str,
+        agent_id: str = "Cobalt Strike",
         cobalt_strike: CobaltStrikeClient | None = None,
     ):
         """Initialize with LevelDB path and optional components
@@ -63,6 +64,7 @@ def __init__(
             raise ValueError("NemesisClient is required")
 
         self.project = project
+        self.agent_id = agent_id
         self.client = nemesis
         self.cobalt_strike = cobalt_strike
         self.db = plyvel.DB(str(db_path), create_if_missing=True)
@@ -148,7 +150,7 @@ async def process_cobaltstrike_download(self, download: Download, beacon: Beacon
             source = f"host://{beacon.computer}" if beacon.computer else None
 
             metadata = FileMetadata(
-                agent_id="Cobalt Strike",
+                agent_id=self.agent_id,
                 source=source,
                 project=self.project,
                 timestamp=datetime.now(UTC),

projects/cli/cli/config.py

@@ -64,6 +64,7 @@ class NemesisConfig(BaseConfig):
 class MythicConfig(BaseConfig):
     url: StrictHttpUrl
     credential: PasswordCredential | TokenCredential
+    agent_id: str = Field(default="mythic", min_length=1, description="Agent identifier for uploaded Mythic files")
 
     @field_validator("credential")
     @classmethod
@@ -78,6 +79,7 @@ def validate_credential(cls, v):
 class OutflankConfig(BaseConfig):
     url: StrictHttpUrl
     credential: PasswordCredential
+    agent_id: str = Field(default="stage1", min_length=1, description="Agent identifier for uploaded Outflank files")
     downloads_dir_path: Path | None = Field(
         None,
         description="Optional: Path to Outflank C2's upload directory where files will be pulled from instead of the Outflank API",
@@ -95,6 +97,9 @@ def validate_path(cls, v):
 class CobaltStrikeConfig(BaseConfig):
     url: StrictHttpUrl
     credential: PasswordCredential
+    agent_id: str = Field(
+        default="Cobalt Strike", min_length=1, description="Agent identifier for uploaded Cobalt Strike files"
+    )
     project: str = Field(description="Project name for Nemesis file uploads")
     poll_interval_sec: Annotated[
         int, Field(gt=0, description="Polling interval of the Cobalt Strike API in seconds")

projects/cli/cli/mythic_connector/config.py

@@ -20,6 +20,7 @@ class TokenCredential:
 class MythicConfig:
     url: str
     credential: UsernamePasswordCredential | TokenCredential
+    agent_id: str
 
     @classmethod
     def from_dict(cls, data: dict) -> "MythicConfig":
@@ -33,7 +34,7 @@ def from_dict(cls, data: dict) -> "MythicConfig":
         else:
             credential = UsernamePasswordCredential(username=cred_data["username"], password=cred_data["password"])
 
-        return cls(url=url_value, credential=credential)
+        return cls(url=url_value, credential=credential, agent_id=data.get("agent_id", "mythic"))
 
 
 @dataclass
@@ -124,6 +125,7 @@ def from_dynaconf(cls, dynaconf: Dynaconf) -> "Settings":
         "mythic.credential.password", must_exist=True, when=Validator("mythic.credential.token", must_exist=False)
     ),
     Validator("mythic.credential.token", must_exist=False),
+    Validator("mythic.agent_id", default="mythic", is_type_of=str),
     # Nemesis validators
     Validator("nemesis.url", must_exist=True),
     Validator("nemesis.credential", must_exist=True),

projects/cli/cli/mythic_connector/handlers.py

@@ -69,7 +69,7 @@ def _build_metadata(self, file_meta: dict[str, Any]) -> dict[str, Any]:
             Formatted metadata dictionary
         """
         return {
-            "agent_id": file_meta["task"]["callback"]["agent_callback_id"],
+            "agent_id": self.cfg.mythic.agent_id,
             "agent_type": "mythic",
             "automated": True,
             "data_type": "file_data",
@@ -178,7 +178,7 @@ async def upload_to_nemesis(file_path: str) -> None:
                 source = f"host://{file_meta.get('host', 'unknown')}"
 
                 metadata = FileMetadata(
-                    agent_id="mythic",
+                    agent_id=self.cfg.mythic.agent_id,
                     source=source,
                     project=self.cfg.project,
                     timestamp=datetime.now(UTC),

projects/cli/cli/stage1_connector/download_processor.py

@@ -46,6 +46,7 @@ def __init__(
         db_path: Path,
         nemesis: NemesisClient,
         project: str,
+        agent_id: str = "stage1",
         outflank_downloads_dir_path: Path | None = None,
         outflank: OutflankC2Client | None = None,
     ):
@@ -78,6 +79,7 @@ def __init__(
 
         self.outflank_downloads_dir_path = outflank_downloads_dir_path
         self.project = project
+        self.agent_id = agent_id
         self.client = nemesis
         self.outflank = outflank
         self.db = plyvel.DB(str(db_path), create_if_missing=True)
@@ -163,7 +165,7 @@ async def process_outflank_download(self, download: Download, implant: Implant)
             source = f"host://{implant.hostname}" if implant.hostname else None
 
             metadata = FileMetadata(
-                agent_id="stage1",
+                agent_id=self.agent_id,
                 source=source,
                 project=self.project,
                 timestamp=datetime.now(UTC),

projects/cli/cli/stage1_connector/stage1_connector.py

@@ -84,6 +84,7 @@ async def setup_outflank_monitor(outflank_config: OutflankConfig, config: Config
                 config.cache_db_path,
                 nemesis_client,
                 project["name"],
+                agent_id=outflank_config.agent_id,
                 outflank_downloads_dir_path=outflank_config.downloads_dir_path,
             )
         else:
@@ -92,6 +93,7 @@ async def setup_outflank_monitor(outflank_config: OutflankConfig, config: Config
                 config.cache_db_path,
                 nemesis_client,
                 project["name"],
+                agent_id=outflank_config.agent_id,
                 outflank=outflank_client,
             )
 

projects/cli/settings_cobaltstrike.yaml

@@ -15,5 +15,6 @@ cobaltstrike:
     credential:
       username: "nemesis_bot"
       password: "cobaltstrike_password"
+    agent_id: "Cobalt Strike"
     project: "my-assessment"
     poll_interval_sec: 3

projects/cli/settings_mythic.yaml

@@ -2,6 +2,7 @@ project: "ASSESS-TEST"
 
 mythic:
   url: "https://mythic.local:7443"
+  agent_id: "mythic"
 
   # Password auth
   credential:

projects/cli/settings_outflank.yaml

@@ -15,6 +15,7 @@ outflank:
     credential:
       username: "nemesis_bot"
       password: "outflank_password"
+    agent_id: "stage1"
 
     # Optional:
     #   When set, stage1 downloads are retrieved from disk instead of stage1's API.

projects/cli/tests/test_config.py

@@ -129,13 +129,15 @@ def test_create_with_password_credential(self):
             credential=PasswordCredential(username="mythic_admin", password="pass"),
         )
         assert isinstance(cfg.credential, PasswordCredential)
+        assert cfg.agent_id == "mythic"
 
     def test_create_with_token_credential(self):
         cfg = MythicConfig(
             url=StrictHttpUrl("https://mythic.local:7443"),
             credential=TokenCredential(token="my-api-token"),
         )
         assert isinstance(cfg.credential, TokenCredential)
+        assert cfg.agent_id == "mythic"
 
     def test_validate_credential_from_dict_token(self):
         """Test that the field_validator correctly handles dict input with token."""
@@ -155,6 +157,13 @@ def test_validate_credential_from_dict_password(self):
         assert isinstance(cfg.credential, PasswordCredential)
         assert cfg.credential.username == "admin"
 
+    def test_custom_agent_id(self):
+        cfg = MythicConfig(
+            url=StrictHttpUrl("https://mythic.local:7443"),
+            credential={"username": "admin", "password": "secret"},
+            agent_id="mythic-initial",
+        )
+        assert cfg.agent_id == "mythic-initial"
 
 # --- OutflankConfig ---
 
@@ -165,6 +174,7 @@ def test_create_minimal(self):
             url=StrictHttpUrl("https://outflank.local"),
             credential=PasswordCredential(username="u", password="p"),
         )
+        assert cfg.agent_id == "stage1"
         assert cfg.downloads_dir_path is None
         assert cfg.poll_interval_sec == 3
 
@@ -195,6 +205,7 @@ def test_create(self):
             credential=PasswordCredential(username="u", password="p"),
             project="assessment-1",
         )
+        assert cfg.agent_id == "Cobalt Strike"
         assert cfg.project == "assessment-1"
         assert cfg.poll_interval_sec == 3
 
@@ -313,6 +324,7 @@ def test_cobaltstrike_section(self, cfg):
         assert str(cs.url) == "https://cobaltstrike.example.com:50443"
         assert cs.credential.username == "nemesis_bot"
         assert cs.credential.password == "cobaltstrike_password"
+        assert cs.agent_id == "Cobalt Strike"
         assert cs.project == "my-assessment"
         assert cs.poll_interval_sec == 3
 
@@ -347,6 +359,7 @@ def test_outflank_section(self, cfg):
         assert str(of.url) == "https://stage1.example.com"
         assert of.credential.username == "nemesis_bot"
         assert of.credential.password == "outflank_password"
+        assert of.agent_id == "stage1"
         assert of.downloads_dir_path is None
         assert of.poll_interval_sec == 3  # default
 

projects/cli/tests/test_sync.py

@@ -22,6 +22,7 @@ def _make_settings(
     mythic_cred=None,
     nemesis_url="https://nemesis.local:8080",
     nemesis_cred=None,
+    mythic_agent_id="mythic",
 ) -> Settings:
     """Create a Settings object for testing."""
     if mythic_cred is None:
@@ -31,7 +32,7 @@ def _make_settings(
 
     return Settings(
         project="TEST-PROJECT",
-        mythic=MythicConfig(url=mythic_url, credential=mythic_cred),
+        mythic=MythicConfig(url=mythic_url, credential=mythic_cred, agent_id=mythic_agent_id),
         nemesis=NemesisConfig(
             url=nemesis_url,
             credential=nemesis_cred,
@@ -278,6 +279,7 @@ def test_mythic_section(self):
         get_settings.cache_clear()
         cfg = get_settings(str(SETTINGS_DIR / "settings_mythic.yaml"))
         assert cfg.mythic.url == "https://mythic.local:7443"
+        assert cfg.mythic.agent_id == "mythic"
         assert isinstance(cfg.mythic.credential, UsernamePasswordCredential)
         assert cfg.mythic.credential.username == "a"
         assert cfg.mythic.credential.password == "a"