diff --git a/projects/cli/cli/cobaltstrike_connector/cobaltstrike_connector.py b/projects/cli/cli/cobaltstrike_connector/cobaltstrike_connector.py index 2fa53b97..01a221c5 100644 --- a/projects/cli/cli/cobaltstrike_connector/cobaltstrike_connector.py +++ b/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, ) diff --git a/projects/cli/cli/cobaltstrike_connector/download_processor.py b/projects/cli/cli/cobaltstrike_connector/download_processor.py index efe94508..18a3e492 100644 --- a/projects/cli/cli/cobaltstrike_connector/download_processor.py +++ b/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), diff --git a/projects/cli/cli/config.py b/projects/cli/cli/config.py index 0f0a5052..bc1dd93f 100644 --- a/projects/cli/cli/config.py +++ b/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") diff --git a/projects/cli/cli/mythic_connector/config.py b/projects/cli/cli/mythic_connector/config.py index 7d1ac24e..b2701b1c 100644 --- a/projects/cli/cli/mythic_connector/config.py +++ b/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), diff --git a/projects/cli/cli/mythic_connector/handlers.py b/projects/cli/cli/mythic_connector/handlers.py index 8e6ca95e..2c3cd926 100644 --- a/projects/cli/cli/mythic_connector/handlers.py +++ b/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), diff --git a/projects/cli/cli/stage1_connector/download_processor.py b/projects/cli/cli/stage1_connector/download_processor.py index 8f18a583..71db9816 100644 --- a/projects/cli/cli/stage1_connector/download_processor.py +++ b/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), diff --git a/projects/cli/cli/stage1_connector/stage1_connector.py b/projects/cli/cli/stage1_connector/stage1_connector.py index 5b4a5b2f..60d78f5b 100644 --- a/projects/cli/cli/stage1_connector/stage1_connector.py +++ b/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, ) diff --git a/projects/cli/settings_cobaltstrike.yaml b/projects/cli/settings_cobaltstrike.yaml index 2d009377..bddc5be9 100644 --- a/projects/cli/settings_cobaltstrike.yaml +++ b/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 diff --git a/projects/cli/settings_mythic.yaml b/projects/cli/settings_mythic.yaml index 7bceae0e..bddc4047 100644 --- a/projects/cli/settings_mythic.yaml +++ b/projects/cli/settings_mythic.yaml @@ -2,6 +2,7 @@ project: "ASSESS-TEST" mythic: url: "https://mythic.local:7443" + agent_id: "mythic" # Password auth credential: diff --git a/projects/cli/settings_outflank.yaml b/projects/cli/settings_outflank.yaml index 0e6f5c62..1c6294cf 100644 --- a/projects/cli/settings_outflank.yaml +++ b/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. diff --git a/projects/cli/tests/test_config.py b/projects/cli/tests/test_config.py index 609f32cc..72c0d49a 100644 --- a/projects/cli/tests/test_config.py +++ b/projects/cli/tests/test_config.py @@ -129,6 +129,7 @@ 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( @@ -136,6 +137,7 @@ def test_create_with_token_credential(self): 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 diff --git a/projects/cli/tests/test_sync.py b/projects/cli/tests/test_sync.py index 94c55d76..a3ad97a3 100644 --- a/projects/cli/tests/test_sync.py +++ b/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"