Skip to content

Commit 7ffd027

Browse files
Multi transport builds (#1517)
* Add multi-callback * Add DSN configuration * Force lowercase * Auto grab pubkey with multi-callback using 0th element. * Update v1 * Update callback rotation * fmt * Remove agent callback_uris and idx in favor of config's * Remove callback_uri from transport init replaced by config * remove callback_uri direct ref * fmt * Cleanup warnings * Address code review feedback - Remove unused imports and variables in agent.rs - Switch to and_then syntax for cleaner Option chaining - Refactor parse_dsn to return Transport struct directly - Convert default parameters to const values - Bubble up errors instead of returning defaults - Use named enum variants instead of numeric constants in tests - Update all test cases to use TransportType enum 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Hulto <[email protected]> * fmt * Fix test_parse_callback_interval_valid test The test was failing because parse_callback_interval() now returns anyhow::Result<u64> instead of u64. Updated the test to properly unwrap the Result before comparison. Co-authored-by: Hulto <[email protected]> * Update enum name * Fix tests * Fix tests * Setup tests * fix? --------- Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Hulto <[email protected]>
1 parent 2d11111 commit 7ffd027

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

53 files changed

+1660
-911
lines changed

implants/imix/src/agent.rs

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -88,17 +88,7 @@ impl<T: Transport + 'static> Agent<T> {
8888
* Callback once using the configured client to claim new tasks and report available output.
8989
*/
9090
pub async fn callback(&mut self) -> Result<()> {
91-
//TODO: De-dupe this fields - probably just need to pass the config not the callback URI.
92-
self.t = T::new(
93-
self.cfg
94-
.info
95-
.clone()
96-
.context("failed to get info")?
97-
.active_transport
98-
.context("failed to get transport")?
99-
.uri,
100-
self.cfg.clone(),
101-
)?;
91+
self.t = T::new(self.cfg.clone())?;
10292
self.claim_tasks(self.t.clone()).await?;
10393
self.report(self.t.clone()).await?;
10494
self.t = T::init(); // re-init to make sure no active connections during sleep
@@ -130,13 +120,23 @@ impl<T: Transport + 'static> Agent<T> {
130120
return Ok(());
131121
}
132122

133-
let interval = match self.cfg.info.clone() {
134-
Some(b) => Ok(b
135-
.active_transport
136-
.context("failed to get transport")?
137-
.interval),
138-
None => Err(anyhow::anyhow!("beacon info is missing from agent")),
139-
}?;
123+
let interval = self
124+
.cfg
125+
.info
126+
.clone()
127+
.context("beacon info is missing from agent")
128+
.and_then(|b| {
129+
b.available_transports
130+
.context("failed to get available transports")
131+
})
132+
.and_then(|available_transports| {
133+
available_transports
134+
.transports
135+
.get(available_transports.active_index as usize)
136+
.cloned()
137+
.context("active transport index out of bounds")
138+
})
139+
.map(|active_transport| active_transport.interval)?;
140140
let delay = match interval.checked_sub(start.elapsed().as_secs()) {
141141
Some(secs) => Duration::from_secs(secs),
142142
None => Duration::from_secs(0),

implants/imix/src/run.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ pub async fn handle_main() {
2525
let retry_interval = cfg
2626
.info
2727
.as_ref()
28-
.and_then(|beacon| beacon.active_transport.as_ref())
28+
.and_then(|beacon| beacon.available_transports.as_ref())
29+
.and_then(|available| available.transports.get(available.active_index as usize))
2930
.map(|transport| transport.interval)
3031
.unwrap_or(DEFAULT_RETRY_INTERVAL);
3132

0 commit comments

Comments
 (0)