rusthound_ce/
api.rs

1use std::{collections::HashMap, error::Error};
2
3use indicatif::ProgressBar;
4use ldap3::SearchEntry;
5use rayon::prelude::*;
6use std::path::PathBuf;
7
8const CACHE_DIR: &str = ".rusthound-cache";
9const CACHE_FILE: &str = "ldap.bin";
10
11use crate::{
12    args::Options, banner::progress_bar, enums::{PARSER_MOD_RE1, PARSER_MOD_RE2, Type, get_type}, json::checker::check_all_result, 
13    objects::{
14        aiaca::AIACA,
15        certtemplate::CertTemplate,
16        common::parse_unknown,
17        computer::Computer,
18        container::Container,
19        domain::Domain,
20        enterpriseca::EnterpriseCA,
21        fsp::Fsp,
22        gpo::Gpo,
23        group::Group,
24        inssuancepolicie::IssuancePolicie,
25        ntauthstore::NtAuthStore,
26        ou::Ou,
27        rootca::RootCA,
28        trust::Trust,
29        user::User,
30        schema::Schema,
31    },
32    transport::ldap::{LdapSearchEntry, collect_from_ldap_into},
33    storage::{DiskStorageReader, DiskStorage, EntrySource},
34    modules::run_modules,
35    json::maker::make_result,
36};
37
38#[derive(Default)]
39pub struct ADResults {
40    pub users: Vec<User>,
41    pub groups: Vec<Group>,
42    pub computers: Vec<Computer>,
43    pub ous: Vec<Ou>,
44    pub domains: Vec<Domain>,
45    pub gpos: Vec<Gpo>,
46    pub fsps: Vec<Fsp>,
47    pub containers: Vec<Container>,
48    pub trusts: Vec<Trust>,
49    pub ntauthstores: Vec<NtAuthStore>,
50    pub aiacas: Vec<AIACA>,
51    pub rootcas: Vec<RootCA>,
52    pub enterprisecas: Vec<EnterpriseCA>,
53    pub certtemplates: Vec<CertTemplate>,
54    pub issuancepolicies: Vec<IssuancePolicie>,
55    pub mappings: DomainMappings,
56}
57
58#[derive(Default)]
59pub struct DomainMappings {
60    /// DN to SID
61    pub dn_sid: HashMap<String, String>,
62    ///  DN to Type
63    pub sid_type: HashMap<String, String>,
64    /// FQDN to SID
65    pub fqdn_sid: HashMap<String, String>,
66    /// fqdn to an ip address
67    pub fqdn_ip: HashMap<String, String>,
68    /// schema guid map
69    pub schema_guid_map: HashMap<String, String>,
70}
71
72impl ADResults {
73    pub fn new() -> Self {
74        Self::default()
75    }
76}
77
78pub async fn prepare_results_from_source<S: EntrySource>(
79    source: S,
80    options: &Options,
81    total_objects: Option<usize>,
82) -> Result<ADResults, Box<dyn std::error::Error>> {
83    let mut ad_results = parse_result_type_from_source(options, source, total_objects)?;
84    run_checker(options, &mut ad_results)?;
85    Ok(ad_results)
86}
87
88/// Like [`prepare_results_from_source`], but reads directly from a disk cache
89/// (`--cache` / `--resume`). The heavy bincode decode of each record is done on
90/// worker threads instead of the single reader thread, which is the difference
91/// between a mostly-idle CPU and a saturated one on large caches.
92pub async fn prepare_results_from_disk(
93    reader: DiskStorageReader<LdapSearchEntry>,
94    options: &Options,
95    total_objects: Option<usize>,
96) -> Result<ADResults, Box<dyn std::error::Error>> {
97    let mut ad_results = parse_result_type_from_disk(options, reader, total_objects)?;
98    run_checker(options, &mut ad_results)?;
99    Ok(ad_results)
100}
101
102/// Post-parse pass: replace and add missing values.
103fn run_checker(options: &Options, ad_results: &mut ADResults) -> Result<(), Box<dyn std::error::Error>> {
104    check_all_result(
105        options,
106        &mut ad_results.users,
107        &mut ad_results.groups,
108        &mut ad_results.computers,
109        &mut ad_results.ous,
110        &mut ad_results.domains,
111        &mut ad_results.gpos,
112        &mut ad_results.fsps,
113        &mut ad_results.containers,
114        &mut ad_results.trusts,
115        &mut ad_results.ntauthstores,
116        &mut ad_results.aiacas,
117        &mut ad_results.rootcas,
118        &mut ad_results.enterprisecas,
119        &mut ad_results.certtemplates,
120        &mut ad_results.issuancepolicies,
121        &ad_results.mappings.dn_sid,
122        &ad_results.mappings.sid_type,
123        &ad_results.mappings.fqdn_sid,
124        &ad_results.mappings.fqdn_ip,
125    )
126}
127
128/// Number of bulk entries buffered before being parsed as one parallel batch.
129///
130/// The batch is split across all CPU cores by rayon, then merged sequentially.
131/// A larger value amortizes scheduling overhead; keeping it bounded caps the
132/// extra memory to one batch of `SearchEntry`s regardless of domain size.
133const PARSE_BATCH: usize = 16_384;
134
135/// Map contributions produced by parsing a single object.
136///
137/// Each `parse()` only ever *inserts* the object's own DN→SID / SID→Type
138/// (and, for computers, FQDN→SID / FQDN→IP) entries — it never reads the maps.
139/// That makes bulk parsing data-parallel: every thread fills its own
140/// `LocalMaps`, which are merged into the global maps afterwards.
141#[derive(Default)]
142struct LocalMaps {
143    dn_sid: HashMap<String, String>,
144    sid_type: HashMap<String, String>,
145    fqdn_sid: HashMap<String, String>,
146    fqdn_ip: HashMap<String, String>,
147}
148
149/// A parsed object, tagged by collection so the merge step can route it into
150/// the right vector. `Domain` and `Schema` are parsed sequentially (they must
151/// run before the bulk), so they never appear here.
152enum Parsed {
153    User(Box<User>),
154    Group(Box<Group>),
155    Computer(Box<Computer>),
156    Ou(Box<Ou>),
157    Gpo(Box<Gpo>),
158    Fsp(Box<Fsp>),
159    Container(Box<Container>),
160    Trust(Box<Trust>),
161    NtAuthStore(Box<NtAuthStore>),
162    AIACA(Box<AIACA>),
163    RootCA(Box<RootCA>),
164    EnterpriseCA(Box<EnterpriseCA>),
165    CertTemplate(Box<CertTemplate>),
166    IssuancePolicie(Box<IssuancePolicie>),
167    /// Filtered-out container, unknown/other object — nothing to collect.
168    Skip,
169}
170
171/// Parse one bulk entry. Runs on a rayon worker thread, so it takes the shared
172/// `schema_guid_map`/`domain_sid` by read-only reference and returns its own
173/// local map contributions.
174fn parse_one(
175    entry: SearchEntry,
176    domain: &str,
177    domain_sid: &str,
178    schema_guid_map: &HashMap<String, String>,
179) -> Result<(Parsed, LocalMaps), Box<dyn Error>> {
180    let mut m = LocalMaps::default();
181    let atype = get_type(&entry).unwrap_or(Type::Unknown);
182
183    let parsed = match atype {
184        Type::User => {
185            let mut user = User::new();
186            user.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
187            Parsed::User(Box::new(user))
188        }
189        Type::Group => {
190            let mut group = Group::new();
191            group.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
192            Parsed::Group(Box::new(group))
193        }
194        Type::Computer => {
195            let mut computer = Computer::new();
196            computer.parse(
197                entry, domain,
198                &mut m.dn_sid, &mut m.sid_type, &mut m.fqdn_sid, &mut m.fqdn_ip,
199                domain_sid, schema_guid_map,
200            )?;
201            Parsed::Computer(Box::new(computer))
202        }
203        Type::Ou => {
204            let mut ou = Ou::new();
205            ou.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
206            Parsed::Ou(Box::new(ou))
207        }
208        Type::Gpo => {
209            let mut gpo = Gpo::new();
210            gpo.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
211            Parsed::Gpo(Box::new(gpo))
212        }
213        Type::ForeignSecurityPrincipal => {
214            let mut security_principal = Fsp::new();
215            security_principal.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid)?;
216            Parsed::Fsp(Box::new(security_principal))
217        }
218        Type::Container => {
219            if PARSER_MOD_RE1.is_match(&entry.dn.to_uppercase())
220                || PARSER_MOD_RE2.is_match(&entry.dn.to_uppercase())
221            {
222                Parsed::Skip
223            } else {
224                let mut container = Container::new();
225                container.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
226                Parsed::Container(Box::new(container))
227            }
228        }
229        Type::Trust => {
230            let mut trust = Trust::new();
231            trust.parse(entry, domain)?;
232            Parsed::Trust(Box::new(trust))
233        }
234        Type::NtAutStore => {
235            let mut nt_auth_store = NtAuthStore::new();
236            nt_auth_store.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
237            Parsed::NtAuthStore(Box::new(nt_auth_store))
238        }
239        Type::AIACA => {
240            let mut aiaca = AIACA::new();
241            aiaca.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
242            Parsed::AIACA(Box::new(aiaca))
243        }
244        Type::RootCA => {
245            let mut root_ca = RootCA::new();
246            root_ca.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
247            Parsed::RootCA(Box::new(root_ca))
248        }
249        Type::EnterpriseCA => {
250            let mut enterprise_ca = EnterpriseCA::new();
251            enterprise_ca.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
252            Parsed::EnterpriseCA(Box::new(enterprise_ca))
253        }
254        Type::CertTemplate => {
255            let mut cert_template = CertTemplate::new();
256            cert_template.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
257            Parsed::CertTemplate(Box::new(cert_template))
258        }
259        Type::IssuancePolicie => {
260            let mut issuance_policie = IssuancePolicie::new();
261            issuance_policie.parse(entry, domain, &mut m.dn_sid, &mut m.sid_type, domain_sid, schema_guid_map)?;
262            Parsed::IssuancePolicie(Box::new(issuance_policie))
263        }
264        // Handled sequentially before the bulk; should not reach here.
265        Type::Domain | Type::Schema => Parsed::Skip,
266        Type::Unknown => {
267            let _unknown = parse_unknown(entry, domain);
268            Parsed::Skip
269        }
270    };
271
272    Ok((parsed, m))
273}
274
275/// Parse a buffered batch of bulk entries in parallel and merge the results
276/// (objects + map contributions) into `results`, preserving input order.
277#[allow(clippy::too_many_arguments)]
278fn parse_batch(
279    buffer: &mut Vec<SearchEntry>,
280    results: &mut ADResults,
281    domain: &str,
282    domain_sid: &str,
283    count: &mut usize,
284    total: Option<usize>,
285    pb: &ProgressBar,
286) -> Result<(), Box<dyn Error>> {
287    if buffer.is_empty() {
288        return Ok(());
289    }
290
291    let entries = std::mem::take(buffer);
292
293    // Serialize across all CPU cores. Errors are stringified inside the worker
294    // (Box<dyn Error> isn't Send) so fail-fast semantics are preserved.
295    let parsed: Vec<Result<(Parsed, LocalMaps), String>> = {
296        let schema_guid_map = &results.mappings.schema_guid_map;
297        entries
298            .into_par_iter()
299            .map(|entry| parse_one(entry, domain, domain_sid, schema_guid_map).map_err(|e| e.to_string()))
300            .collect()
301    };
302
303    // Merge sequentially, in input order, so output ordering matches the
304    // single-threaded collector exactly.
305    for item in parsed {
306        let (obj, maps) = item.map_err(|e| -> Box<dyn Error> { e.into() })?;
307
308        results.mappings.dn_sid.extend(maps.dn_sid);
309        results.mappings.sid_type.extend(maps.sid_type);
310        results.mappings.fqdn_sid.extend(maps.fqdn_sid);
311        results.mappings.fqdn_ip.extend(maps.fqdn_ip);
312
313        match obj {
314            Parsed::User(o) => results.users.push(*o),
315            Parsed::Group(o) => results.groups.push(*o),
316            Parsed::Computer(o) => results.computers.push(*o),
317            Parsed::Ou(o) => results.ous.push(*o),
318            Parsed::Gpo(o) => results.gpos.push(*o),
319            Parsed::Fsp(o) => results.fsps.push(*o),
320            Parsed::Container(o) => results.containers.push(*o),
321            Parsed::Trust(o) => results.trusts.push(*o),
322            Parsed::NtAuthStore(o) => results.ntauthstores.push(*o),
323            Parsed::AIACA(o) => results.aiacas.push(*o),
324            Parsed::RootCA(o) => results.rootcas.push(*o),
325            Parsed::EnterpriseCA(o) => results.enterprisecas.push(*o),
326            Parsed::CertTemplate(o) => results.certtemplates.push(*o),
327            Parsed::IssuancePolicie(o) => results.issuancepolicies.push(*o),
328            Parsed::Skip => {}
329        }
330
331        update_progress(count, total, pb)?;
332    }
333
334    Ok(())
335}
336
337/// Advance the parsing progress bar.
338fn update_progress(count: &mut usize, total: Option<usize>, pb: &ProgressBar) -> Result<(), Box<dyn Error>> {
339    if let Some(total) = total {
340        *count += 1;
341        // Percentage (%) = 100 x partial value / total value
342        let percentage = 100 * *count / total;
343        progress_bar(
344            pb.to_owned(),
345            "Parsing LDAP objects".to_string(),
346            percentage.try_into()?,
347            "%".to_string(),
348        );
349    }
350    Ok(())
351}
352
353// for `total_objects`, the total number of objects may not be known if the ldap query was never run
354// (e.g run was resumed from cached results)
355//
356// Parsing strategy:
357//   * Schema and Domain objects are parsed sequentially as they stream in. The
358//     collector guarantees they come first (see the naming-context ordering in
359//     `ldap.rs`), and the rest of the parsing reads `schema_guid_map` /
360//     `domain_sid` read-only, so they must be complete beforehand.
361//   * Every other object is buffered and parsed in parallel batches across all
362//     CPU cores, which is the hot path on large domains.
363pub fn parse_result_type_from_source(
364    common_args: &Options,
365    source: impl EntrySource,
366    total_objects: Option<usize>,
367) -> Result<ADResults, Box<dyn Error>> {
368    let mut results = ADResults::default();
369    // Domain name
370    let domain = &common_args.domain;
371
372    // Needed for progress bar stats
373    let pb = ProgressBar::new(1);
374    let mut count = 0usize;
375    let total = total_objects;
376    let mut domain_sid: String = "DOMAIN_SID".to_owned();
377
378    log::info!("Starting the LDAP objects parsing...");
379
380    let mut buffer: Vec<SearchEntry> = Vec::with_capacity(PARSE_BATCH);
381
382    for entry in source.into_entry_iter() {
383        let entry: SearchEntry = entry?.into();
384        route_entry(entry, &mut results, &mut buffer, domain, &mut domain_sid, &mut count, total, &pb)?;
385    }
386
387    // Parse whatever remains in the buffer.
388    parse_batch(&mut buffer, &mut results, domain, &domain_sid, &mut count, total, &pb)?;
389
390    pb.finish_and_clear();
391    log::info!("Parsing LDAP objects finished!");
392    Ok(results)
393}
394
395/// Route a single decoded entry: schema and domain are parsed sequentially
396/// (they must complete before the bulk), everything else is buffered for
397/// parallel batch parsing.
398#[allow(clippy::too_many_arguments)]
399fn route_entry(
400    entry: SearchEntry,
401    results: &mut ADResults,
402    buffer: &mut Vec<SearchEntry>,
403    domain: &str,
404    domain_sid: &mut String,
405    count: &mut usize,
406    total: Option<usize>,
407    pb: &ProgressBar,
408) -> Result<(), Box<dyn Error>> {
409    let atype = get_type(&entry).unwrap_or(Type::Unknown);
410    match atype {
411        Type::Schema => {
412            // Sequential: builds schema_guid_map, read by every ACE parse.
413            let schema = Schema::new();
414            schema.parse(entry, &mut results.mappings.schema_guid_map)?;
415            update_progress(count, total, pb)?;
416        }
417        Type::Domain => {
418            // Flush any already-buffered bulk objects first, so they keep the
419            // domain_sid they were collected under (matches the single-threaded
420            // ordering). In practice the domain object precedes the bulk, so the
421            // buffer is empty here.
422            parse_batch(buffer, results, domain, domain_sid.as_str(), count, total, pb)?;
423
424            let mut domain_object = Domain::new();
425            let domain_sid_from_domain = domain_object.parse(
426                entry,
427                domain,
428                &mut results.mappings.dn_sid,
429                &mut results.mappings.sid_type,
430                &results.mappings.schema_guid_map,
431            )?;
432            // Update only if domain_sid is valid
433            if domain_sid_from_domain != "DOMAIN_SID" && !domain_sid_from_domain.is_empty() {
434                *domain_sid = domain_sid_from_domain;
435            }
436            if !domain_object.object_identifier().is_empty() {
437                results.domains.push(domain_object);
438            }
439            update_progress(count, total, pb)?;
440        }
441        _ => {
442            buffer.push(entry);
443            if buffer.len() >= PARSE_BATCH {
444                parse_batch(buffer, results, domain, domain_sid.as_str(), count, total, pb)?;
445            }
446        }
447    }
448    Ok(())
449}
450
451// Parse directly from a disk cache, decoding bincode records on worker threads.
452//
453// The reader thread only does cheap length-prefixed reads; the CPU-heavy
454// bincode decode happens inside a rayon batch, alongside the object parsing.
455// This keeps all cores busy on `--cache` / `--resume`, where the single-threaded
456// decode was previously the bottleneck.
457pub fn parse_result_type_from_disk(
458    common_args: &Options,
459    mut reader: DiskStorageReader<LdapSearchEntry>,
460    total_objects: Option<usize>,
461) -> Result<ADResults, Box<dyn Error>> {
462    let mut results = ADResults::default();
463    let domain = &common_args.domain;
464
465    let pb = ProgressBar::new(1);
466    let mut count = 0usize;
467    let total = total_objects;
468    let mut domain_sid: String = "DOMAIN_SID".to_owned();
469
470    log::info!("Starting the LDAP objects parsing...");
471
472    let mut raw_buf: Vec<Vec<u8>> = Vec::with_capacity(PARSE_BATCH);
473    let mut buffer: Vec<SearchEntry> = Vec::with_capacity(PARSE_BATCH);
474
475    loop {
476        match reader.next_raw() {
477            Some(Ok(blob)) => {
478                raw_buf.push(blob);
479                if raw_buf.len() >= PARSE_BATCH {
480                    decode_and_route(&mut raw_buf, &mut results, &mut buffer, domain, &mut domain_sid, &mut count, total, &pb)?;
481                }
482            }
483            Some(Err(e)) => return Err(e.into()),
484            None => break,
485        }
486    }
487
488    // Decode and route any remaining raw records, then parse the last bulk batch.
489    decode_and_route(&mut raw_buf, &mut results, &mut buffer, domain, &mut domain_sid, &mut count, total, &pb)?;
490    parse_batch(&mut buffer, &mut results, domain, &domain_sid, &mut count, total, &pb)?;
491
492    pb.finish_and_clear();
493    log::info!("Parsing LDAP objects finished!");
494    Ok(results)
495}
496
497/// Decode a batch of raw bincode records in parallel, then route each decoded
498/// entry (preserving stream order, so schema/domain are handled before bulk).
499#[allow(clippy::too_many_arguments)]
500fn decode_and_route(
501    raw_buf: &mut Vec<Vec<u8>>,
502    results: &mut ADResults,
503    buffer: &mut Vec<SearchEntry>,
504    domain: &str,
505    domain_sid: &mut String,
506    count: &mut usize,
507    total: Option<usize>,
508    pb: &ProgressBar,
509) -> Result<(), Box<dyn Error>> {
510    if raw_buf.is_empty() {
511        return Ok(());
512    }
513
514    let blobs = std::mem::take(raw_buf);
515
516    // The CPU-heavy bincode decode, spread across all cores.
517    let decoded: Vec<Result<SearchEntry, String>> = blobs
518        .into_par_iter()
519        .map(|data| {
520            bincode::decode_from_slice::<LdapSearchEntry, _>(&data, bincode::config::standard())
521                .map(|(entry, _)| SearchEntry::from(entry))
522                .map_err(|e| format!("Failed to decode item: {e:?}"))
523        })
524        .collect();
525
526    for decoded_entry in decoded {
527        let entry = decoded_entry.map_err(|e| -> Box<dyn Error> { e.into() })?;
528        route_entry(entry, results, buffer, domain, domain_sid, count, total, pb)?;
529    }
530
531    Ok(())
532}
533
534fn cache_path(domain: &str) -> PathBuf {
535    PathBuf::from(CACHE_DIR).join(domain).join(CACHE_FILE)
536}
537
538/// Run the complete RustHound-CE workflow over an already-authenticated LDAP
539/// session (get one from `transport::ldap::ldap_auth`, or bring your own):
540/// collect -> parse -> modules -> JSON/zip. Returns the output path.
541///
542/// Storage follows `options`: `resume` parses an existing disk cache without
543/// collecting; `cache` streams to `.rusthound-cache/<domain>/ldap.bin`;
544/// otherwise entries are kept in memory. Modules run per `collection_method`.
545///
546/// Never calls `process::exit`; never unbinds/drops the session.
547pub async fn run_collection(
548    ldap: &mut ldap3::Ldap,
549    options: &Options,
550) -> Result<String, Box<dyn std::error::Error>> {
551    let mut results = if options.resume {
552        let reader = DiskStorageReader::from_path(cache_path(&options.domain))?;
553        prepare_results_from_disk(reader, options, None).await?
554    } else if options.cache {
555        let path = cache_path(&options.domain);
556        if let Some(parent) = path.parent() {
557            std::fs::create_dir_all(parent)?;
558        }
559        let mut writer = DiskStorage::new_with_capacity(path, options.cache_buffer_size)?;
560        let total = collect_from_ldap_into(ldap, &options.ldap_filter, &mut writer).await?;
561        prepare_results_from_disk(writer.into_reader()?, options, Some(total)).await?
562    } else {
563        let mut entries = Vec::new();
564        let total = collect_from_ldap_into(ldap, &options.ldap_filter, &mut entries).await?;
565        prepare_results_from_source(entries, options, Some(total)).await?
566    };
567
568    run_modules(options, &mut results).await?;
569    make_result(options, results)
570}
rusthound_ce/api.rs

rusthound_ce/
api.rs
