1pub mod actions;
3
4pub use actions::LdapClientProtocol;
5
6use crate::llm::actions::client_trait::{Client, ClientActionResult};
7use anyhow::{Context, Result};
8use std::net::SocketAddr;
9use std::sync::Arc;
10use tokio::sync::mpsc;
11use tracing::{debug, error, info, trace, warn};
12
13use crate::client::ldap::actions::{
14 LDAP_CLIENT_BIND_RESPONSE_EVENT, LDAP_CLIENT_CONNECTED_EVENT,
15 LDAP_CLIENT_MODIFY_RESPONSE_EVENT, LDAP_CLIENT_SEARCH_RESULTS_EVENT,
16};
17use crate::client::llm_budget::call_llm_for_client;
18use crate::llm::ollama_client::OllamaClient;
19use crate::llm::ClientLlmResult;
20use crate::protocol::Event;
21use crate::state::app_state::AppState;
22use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
23use crate::state::{ClientId, ClientStatus};
24
25use ldap3::{LdapConn, Mod, Scope};
26
27enum Applied {
32 Ran { detail: String, event: Event },
35 Disconnect,
37 Nothing(&'static str),
39}
40
41const MAX_FOLLOWUP_DEPTH: usize = 4;
49
50pub struct LdapClient;
52
53impl LdapClient {
54 pub async fn connect_with_llm_actions(
56 remote_addr: String,
57 llm_client: OllamaClient,
58 app_state: Arc<AppState>,
59 status_tx: mpsc::UnboundedSender<String>,
60 client_id: ClientId,
61 ) -> Result<SocketAddr> {
62 info!("LDAP client {} connecting to {}", client_id, remote_addr);
63
64 let ldap_url = if remote_addr.starts_with("ldap://") || remote_addr.starts_with("ldaps://")
66 {
67 remote_addr.clone()
68 } else {
69 format!("ldap://{}", remote_addr)
70 };
71
72 let ldap = tokio::task::spawn_blocking(move || LdapConn::new(&ldap_url))
74 .await
75 .context("Failed to spawn LDAP connection task")??;
76
77 let socket_addr: SocketAddr = remote_addr
81 .split("://")
82 .last()
83 .unwrap_or(&remote_addr)
84 .parse()
85 .context(format!(
86 "Failed to parse socket address from {}",
87 remote_addr
88 ))?;
89
90 info!("LDAP client {} connected to {}", client_id, socket_addr);
91
92 app_state
94 .update_client_status(client_id, ClientStatus::Connected)
95 .await;
96 let _ = status_tx.send(format!("[CLIENT] LDAP client {} connected", client_id));
97 let _ = status_tx.send("__UPDATE_UI__".to_string());
98
99 let ldap = Arc::new(tokio::sync::Mutex::new(ldap));
102 let protocol = Arc::new(LdapClientProtocol::new());
103
104 let command_rx =
108 crate::client::command_support::register_command_channel(&app_state, client_id).await;
109 let cmd_task = tokio::spawn(Self::command_loop(
110 command_rx,
111 protocol.clone(),
112 Arc::clone(&ldap),
113 client_id,
114 llm_client.clone(),
115 app_state.clone(),
116 status_tx.clone(),
117 ));
118 app_state.register_client_task(client_id, cmd_task).await;
119
120 let protocol_clone = protocol.clone();
122 let ldap_clone = Arc::clone(&ldap);
123 let app_state_clone = Arc::clone(&app_state);
124 let status_tx_clone = status_tx.clone();
125
126 let task_registrar = app_state.clone();
130 let task_handle = tokio::spawn(async move {
131 if let Some(instruction) = app_state_clone.get_instruction_for_client(client_id).await {
132 let event = Event::new(
133 &LDAP_CLIENT_CONNECTED_EVENT,
134 serde_json::json!({
135 "remote_addr": socket_addr.to_string(),
136 }),
137 );
138
139 let memory = app_state_clone
140 .get_memory_for_client(client_id)
141 .await
142 .unwrap_or_default();
143
144 match call_llm_for_client(
145 &llm_client,
146 &app_state_clone,
147 client_id.to_string(),
148 &instruction,
149 &memory,
150 Some(&event),
151 protocol_clone.as_ref(),
152 &status_tx_clone,
153 )
154 .await
155 {
156 Ok(ClientLlmResult {
157 actions,
158 memory_updates,
159 }) => {
160 if let Some(mem) = memory_updates {
162 app_state_clone.set_memory_for_client(client_id, mem).await;
163 }
164
165 for action in actions {
167 if let Err(e) = Self::execute_ldap_action(
168 action,
169 &ldap_clone,
170 &protocol_clone,
171 &llm_client,
172 &app_state_clone,
173 &status_tx_clone,
174 client_id,
175 &instruction,
176 0,
177 )
178 .await
179 {
180 error!("Failed to execute LDAP action: {}", e);
181 }
182 }
183 }
184 Err(e) => {
185 error!("LLM error for LDAP client {}: {}", client_id, e);
186 }
187 }
188 }
189 });
190 task_registrar
191 .register_client_task(client_id, task_handle)
192 .await;
193
194 Ok(socket_addr)
195 }
196
197 #[allow(clippy::too_many_arguments)]
205 async fn command_loop(
206 mut command_rx: mpsc::Receiver<ClientCommand>,
207 protocol: Arc<LdapClientProtocol>,
208 ldap: Arc<tokio::sync::Mutex<LdapConn>>,
209 client_id: ClientId,
210 llm_client: OllamaClient,
211 app_state: Arc<AppState>,
212 status_tx: mpsc::UnboundedSender<String>,
213 ) {
214 use crate::llm::actions::protocol_trait::Protocol;
215 use crate::state::AccessLogOwner;
216
217 while let Some(command) = command_rx.recv().await {
218 let action = command.action.clone();
219
220 let mut follow_up: Option<Event> = None;
221 let outcome = match protocol.execute_action(action.clone()) {
222 Err(e) => Ok(ClientSendOutcome::Rejected {
223 error: e.to_string(),
224 }),
225 Ok(result) => match Self::apply_action(result, &ldap, client_id).await {
226 Err(e) => Err(e),
227 Ok(Applied::Disconnect) => Ok(ClientSendOutcome::Disconnected),
228 Ok(Applied::Nothing(what)) => Ok(ClientSendOutcome::Executed {
229 detail: what.to_string(),
230 }),
231 Ok(Applied::Ran { detail, event }) => {
232 follow_up = Some(event);
233 Ok(ClientSendOutcome::Executed { detail })
234 }
235 },
236 };
237
238 let outcome_json = match &outcome {
239 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
240 Err(e) => serde_json::json!({"error": e.to_string()}),
241 };
242 app_state
243 .record_access_log(
244 AccessLogOwner::Client(client_id.as_u32()),
245 protocol.protocol_name(),
246 None,
247 "injected_action",
248 action,
249 vec![outcome_json],
250 )
251 .await;
252
253 let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
254 if let Err(e) = &outcome {
255 error!("LDAP client {} injected action failed: {}", client_id, e);
256 let _ = status_tx.send(format!(
257 "[WARN] Client {} injected action failed: {}",
258 client_id, e
259 ));
260 }
261 let _ = status_tx.send("__UPDATE_UI__".to_string());
262 crate::client::command_support::reply(command, outcome);
263
264 if disconnect {
265 Self::mark_disconnected(client_id, &app_state, &status_tx).await;
266 break;
267 }
268
269 if let Some(event) = follow_up {
272 let instruction = app_state
273 .get_instruction_for_client(client_id)
274 .await
275 .unwrap_or_default();
276 if let Err(e) = Self::call_llm_with_event(
277 &llm_client,
278 &app_state,
279 &status_tx,
280 client_id,
281 &instruction,
282 &protocol,
283 &ldap,
284 event,
285 0,
286 )
287 .await
288 {
289 error!("LDAP client {} response event failed: {}", client_id, e);
290 }
291 }
292 }
293
294 app_state.remove_client_handle(client_id).await;
297 let _ = status_tx.send("__UPDATE_UI__".to_string());
298 }
299
300 #[allow(clippy::too_many_arguments)]
305 fn execute_ldap_action<'a>(
306 action: serde_json::Value,
307 ldap: &'a Arc<tokio::sync::Mutex<LdapConn>>,
308 protocol: &'a Arc<LdapClientProtocol>,
309 llm_client: &'a OllamaClient,
310 app_state: &'a Arc<AppState>,
311 status_tx: &'a mpsc::UnboundedSender<String>,
312 client_id: ClientId,
313 instruction: &'a str,
314 depth: usize,
315 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>> {
316 Box::pin(async move {
317 match Self::apply_action(protocol.execute_action(action)?, ldap, client_id).await? {
318 Applied::Ran { event, .. } => {
319 Self::call_llm_with_event(
320 llm_client,
321 app_state,
322 status_tx,
323 client_id,
324 instruction,
325 protocol,
326 ldap,
327 event,
328 depth,
329 )
330 .await?;
331 }
332 Applied::Disconnect => {
333 info!("LDAP client {} disconnecting", client_id);
334 Self::mark_disconnected(client_id, app_state, status_tx).await;
335 }
336 Applied::Nothing(what) => {
337 trace!("LDAP client {} action had no effect: {}", client_id, what);
338 }
339 }
340
341 Ok(())
342 })
343 }
344
345 async fn apply_action(
348 result: ClientActionResult,
349 ldap: &Arc<tokio::sync::Mutex<LdapConn>>,
350 client_id: ClientId,
351 ) -> Result<Applied> {
352 match result {
353 ClientActionResult::Custom { name, data } => match name.as_str() {
354 "ldap_bind" => {
355 let dn = data
356 .get("dn")
357 .and_then(|v| v.as_str())
358 .context("Missing 'dn' in bind action")?;
359 let password = data
360 .get("password")
361 .and_then(|v| v.as_str())
362 .context("Missing 'password' in bind action")?;
363
364 debug!("LDAP client {} binding as {}", client_id, dn);
365
366 let dn_owned = dn.to_string();
367 let password_owned = password.to_string();
368 let ldap_clone = Arc::clone(ldap);
369
370 let bind_result = tokio::task::spawn_blocking(move || {
372 let mut ldap_guard = ldap_clone.blocking_lock();
373 ldap_guard.simple_bind(&dn_owned, &password_owned)
374 })
375 .await
376 .context("Failed to spawn bind task")??;
377
378 let (success, message) = match bind_result.success() {
379 Ok(_) => (true, "Bind successful".to_string()),
380 Err(e) => (false, format!("Bind failed: {:?}", e)),
381 };
382
383 info!("LDAP client {} bind result: {}", client_id, message);
384
385 Ok(Applied::Ran {
386 detail: format!("bind as {}: {}", dn, message),
387 event: Event::new(
388 &LDAP_CLIENT_BIND_RESPONSE_EVENT,
389 serde_json::json!({
390 "dn": dn,
391 "success": success,
392 "message": message,
393 }),
394 ),
395 })
396 }
397 "ldap_search" => {
398 let base_dn = data
399 .get("base_dn")
400 .and_then(|v| v.as_str())
401 .context("Missing 'base_dn' in search action")?;
402 let filter = data
403 .get("filter")
404 .and_then(|v| v.as_str())
405 .context("Missing 'filter' in search action")?;
406 let attributes: Vec<String> = data
407 .get("attributes")
408 .and_then(|v| v.as_array())
409 .map(|arr| {
410 arr.iter()
411 .filter_map(|v| v.as_str().map(|s| s.to_string()))
412 .collect()
413 })
414 .unwrap_or_default();
415 let scope_str = data
416 .get("scope")
417 .and_then(|v| v.as_str())
418 .unwrap_or("subtree");
419
420 let scope = match scope_str {
421 "base" => Scope::Base,
422 "one" => Scope::OneLevel,
423 _ => Scope::Subtree,
424 };
425
426 debug!(
427 "LDAP client {} searching: base={}, filter={}, scope={:?}",
428 client_id, base_dn, filter, scope
429 );
430
431 let base_dn_owned = base_dn.to_string();
432 let filter_owned = filter.to_string();
433 let attrs_owned: Vec<String> = if attributes.is_empty() {
434 vec!["*".to_string()]
435 } else {
436 attributes.clone()
437 };
438
439 let ldap_clone = Arc::clone(ldap);
440
441 let search_result = tokio::task::spawn_blocking(move || {
443 let mut ldap_guard = ldap_clone.blocking_lock();
444 let attrs: Vec<&str> = attrs_owned.iter().map(|s| s.as_str()).collect();
445 ldap_guard.search(&base_dn_owned, scope, &filter_owned, attrs)
446 })
447 .await
448 .context("Failed to spawn search task")??;
449
450 let entries = match search_result.success() {
456 Ok((entries, _)) => entries,
457 Err(e) => {
458 let message = format!("Search failed: {:?}", e);
459 info!("LDAP client {} search result: {}", client_id, message);
460 return Ok(Applied::Ran {
461 detail: format!(
462 "search base={} filter={}: {}",
463 base_dn, filter, message
464 ),
465 event: Event::new(
466 &LDAP_CLIENT_SEARCH_RESULTS_EVENT,
467 serde_json::json!({
468 "base_dn": base_dn,
469 "filter": filter,
470 "success": false,
471 "message": message,
472 "entries": [],
473 "count": 0,
474 }),
475 ),
476 });
477 }
478 };
479
480 let mut json_entries = Vec::new();
482 for entry in entries {
483 use ldap3::SearchEntry;
484 let search_entry = SearchEntry::construct(entry);
485 let mut attrs_map = serde_json::Map::new();
486 for (attr_name, attr_values) in search_entry.attrs {
487 attrs_map.insert(attr_name, serde_json::json!(attr_values));
488 }
489 json_entries.push(serde_json::json!({
490 "dn": search_entry.dn,
491 "attributes": attrs_map,
492 }));
493 }
494
495 let count = json_entries.len();
496 info!(
497 "LDAP client {} search returned {} entries",
498 client_id, count
499 );
500
501 Ok(Applied::Ran {
502 detail: format!(
503 "search base={} filter={} returned {} entries",
504 base_dn, filter, count
505 ),
506 event: Event::new(
507 &LDAP_CLIENT_SEARCH_RESULTS_EVENT,
508 serde_json::json!({
509 "base_dn": base_dn,
510 "filter": filter,
511 "success": true,
512 "entries": json_entries,
513 "count": count,
514 }),
515 ),
516 })
517 }
518 "ldap_add" => {
519 let dn = data
520 .get("dn")
521 .and_then(|v| v.as_str())
522 .context("Missing 'dn' in add action")?;
523 let attributes = data
524 .get("attributes")
525 .context("Missing 'attributes' in add action")?;
526
527 debug!("LDAP client {} adding entry: {}", client_id, dn);
528
529 let attrs_obj = attributes
535 .as_object()
536 .context("'attributes' in add action must be an object")?;
537 let mut attrs_vec = Vec::new();
538 for (attr_name, attr_values) in attrs_obj {
539 let values: std::collections::HashSet<String> = match attr_values {
540 serde_json::Value::String(v) => std::iter::once(v.clone()).collect(),
541 serde_json::Value::Array(arr) => arr
542 .iter()
543 .map(|v| {
544 v.as_str().map(str::to_string).with_context(|| {
545 format!("attribute '{attr_name}' has a non-string value")
546 })
547 })
548 .collect::<Result<_>>()?,
549 _ => anyhow::bail!(
550 "attribute '{attr_name}' must be a string or an array of strings"
551 ),
552 };
553 attrs_vec.push((attr_name.clone(), values));
554 }
555
556 let dn_owned = dn.to_string();
557 let ldap_clone = Arc::clone(ldap);
558
559 let add_result = tokio::task::spawn_blocking(move || {
561 let mut ldap_guard = ldap_clone.blocking_lock();
562 ldap_guard.add(&dn_owned, attrs_vec)
563 })
564 .await
565 .context("Failed to spawn add task")??;
566
567 let (success, message) = match add_result.success() {
568 Ok(_) => (true, "Entry added successfully".to_string()),
569 Err(e) => (false, format!("Add failed: {:?}", e)),
570 };
571
572 info!("LDAP client {} add result: {}", client_id, message);
573
574 Ok(Applied::Ran {
575 detail: format!("add {}: {}", dn, message),
576 event: Event::new(
577 &LDAP_CLIENT_MODIFY_RESPONSE_EVENT,
578 serde_json::json!({
579 "operation": "add",
580 "dn": dn,
581 "success": success,
582 "message": message,
583 }),
584 ),
585 })
586 }
587 "ldap_modify" => {
588 let dn_owned = data
589 .get("dn")
590 .and_then(|v| v.as_str())
591 .context("Missing 'dn' in modify action")?
592 .to_string();
593 let operation_owned = data
594 .get("operation")
595 .and_then(|v| v.as_str())
596 .context("Missing 'operation' in modify action")?
597 .to_string();
598 let attribute_owned = data
599 .get("attribute")
600 .and_then(|v| v.as_str())
601 .context("Missing 'attribute' in modify action")?
602 .to_string();
603 let values_owned: Vec<String> = data
604 .get("values")
605 .and_then(|v| v.as_array())
606 .map(|arr| {
607 arr.iter()
608 .filter_map(|v| v.as_str().map(|s| s.to_string()))
609 .collect()
610 })
611 .unwrap_or_default();
612
613 debug!(
614 "LDAP client {} modifying entry: {} ({} {})",
615 client_id, dn_owned, operation_owned, attribute_owned
616 );
617
618 let dn_for_detail = dn_owned.clone();
619 let ldap_clone = Arc::clone(ldap);
620
621 let modify_result = tokio::task::spawn_blocking(move || {
624 let mut ldap_guard = ldap_clone.blocking_lock();
625
626 let values_refs: std::collections::HashSet<&str> =
628 values_owned.iter().map(|s| s.as_str()).collect();
629
630 let mod_op = match operation_owned.as_str() {
632 "add" => Mod::Add(attribute_owned.as_str(), values_refs.clone()),
633 "delete" => Mod::Delete(attribute_owned.as_str(), values_refs.clone()),
634 "replace" => Mod::Replace(attribute_owned.as_str(), values_refs),
635 _ => {
636 return Err(anyhow::anyhow!(
637 "Invalid operation: {}",
638 operation_owned
639 ))
640 }
641 };
642
643 ldap_guard
644 .modify(&dn_owned, vec![mod_op])
645 .context("Failed to modify entry")
646 })
647 .await
648 .context("Failed to spawn modify task")??;
649
650 let (success, message) = match modify_result.success() {
651 Ok(_) => (true, "Entry modified successfully".to_string()),
652 Err(e) => (false, format!("Modify failed: {:?}", e)),
653 };
654
655 info!("LDAP client {} modify result: {}", client_id, message);
656
657 Ok(Applied::Ran {
658 event: Event::new(
659 &LDAP_CLIENT_MODIFY_RESPONSE_EVENT,
660 serde_json::json!({
661 "operation": "modify",
662 "dn": dn_for_detail,
663 "success": success,
664 "message": message,
665 }),
666 ),
667 detail: format!("modify {}: {}", dn_for_detail, message),
668 })
669 }
670 "ldap_delete" => {
671 let dn = data
672 .get("dn")
673 .and_then(|v| v.as_str())
674 .context("Missing 'dn' in delete action")?;
675
676 debug!("LDAP client {} deleting entry: {}", client_id, dn);
677
678 let dn_owned = dn.to_string();
679 let ldap_clone = Arc::clone(ldap);
680
681 let delete_result = tokio::task::spawn_blocking(move || {
683 let mut ldap_guard = ldap_clone.blocking_lock();
684 ldap_guard.delete(&dn_owned)
685 })
686 .await
687 .context("Failed to spawn delete task")??;
688
689 let (success, message) = match delete_result.success() {
690 Ok(_) => (true, "Entry deleted successfully".to_string()),
691 Err(e) => (false, format!("Delete failed: {:?}", e)),
692 };
693
694 info!("LDAP client {} delete result: {}", client_id, message);
695
696 Ok(Applied::Ran {
697 detail: format!("delete {}: {}", dn, message),
698 event: Event::new(
699 &LDAP_CLIENT_MODIFY_RESPONSE_EVENT,
700 serde_json::json!({
701 "operation": "delete",
702 "dn": dn,
703 "success": success,
704 "message": message,
705 }),
706 ),
707 })
708 }
709 other => {
710 trace!("Unknown LDAP custom action: {}", other);
711 Ok(Applied::Nothing("custom result not handled by this client"))
712 }
713 },
714 ClientActionResult::Disconnect => {
715 let ldap_clone = Arc::clone(ldap);
718 let _ = tokio::task::spawn_blocking(move || {
719 let mut ldap_guard = ldap_clone.blocking_lock();
720 ldap_guard.unbind()
721 })
722 .await;
723 Ok(Applied::Disconnect)
724 }
725 ClientActionResult::WaitForMore => Ok(Applied::Nothing("wait_for_more")),
726 ClientActionResult::NoAction => Ok(Applied::Nothing("no_action")),
727 ClientActionResult::SendData(_) => Ok(Applied::Nothing(
728 "send_data is not meaningful for an ldap3 connection",
729 )),
730 ClientActionResult::Multiple(_) => Ok(Applied::Nothing(
731 "multiple results not handled by this client",
732 )),
733 }
734 }
735
736 #[allow(clippy::too_many_arguments)]
738 async fn call_llm_with_event(
739 llm_client: &OllamaClient,
740 app_state: &Arc<AppState>,
741 status_tx: &mpsc::UnboundedSender<String>,
742 client_id: ClientId,
743 instruction: &str,
744 protocol: &Arc<LdapClientProtocol>,
745 ldap: &Arc<tokio::sync::Mutex<LdapConn>>,
746 event: Event,
747 depth: usize,
748 ) -> Result<()> {
749 let memory = app_state
750 .get_memory_for_client(client_id)
751 .await
752 .unwrap_or_default();
753
754 match call_llm_for_client(
755 llm_client,
756 app_state,
757 client_id.to_string(),
758 instruction,
759 &memory,
760 Some(&event),
761 protocol.as_ref(),
762 status_tx,
763 )
764 .await
765 {
766 Ok(ClientLlmResult {
767 actions,
768 memory_updates,
769 }) => {
770 if let Some(mem) = memory_updates {
772 app_state.set_memory_for_client(client_id, mem).await;
773 }
774
775 if !actions.is_empty() && depth >= MAX_FOLLOWUP_DEPTH {
776 warn!(
777 "LDAP client {} reached the follow-up depth bound ({}); dropping {} \
778 action(s) rather than looping",
779 client_id,
780 MAX_FOLLOWUP_DEPTH,
781 actions.len()
782 );
783 return Ok(());
784 }
785
786 for action in actions {
788 if let Err(e) = Self::execute_ldap_action(
789 action,
790 ldap,
791 protocol,
792 llm_client,
793 app_state,
794 status_tx,
795 client_id,
796 instruction,
797 depth + 1,
798 )
799 .await
800 {
801 error!("Failed to execute LDAP action: {}", e);
802 }
803 }
804 }
805 Err(e) => {
806 error!("LLM error for LDAP client {}: {}", client_id, e);
807 }
808 }
809
810 Ok(())
811 }
812
813 async fn mark_disconnected(
814 client_id: ClientId,
815 app_state: &Arc<AppState>,
816 status_tx: &mpsc::UnboundedSender<String>,
817 ) {
818 app_state
819 .update_client_status(client_id, ClientStatus::Disconnected)
820 .await;
821 let _ = status_tx.send(format!("[CLIENT] LDAP client {} disconnected", client_id));
822 let _ = status_tx.send("__UPDATE_UI__".to_string());
823 }
824}