rusthound_ce/
api.rs

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

rusthound_ce/
api.rs
