ntdsextract2 1.4.20

use anyhow::{bail, Result};
use clap::Subcommand;
use dfir_windows_types::Guid;
use sddl::{AccessMaskFlag, Contains, Sid};
use serde::Serialize;
use std::{
    collections::{hash_map::Entry, HashMap, HashSet},
    str::FromStr,
};
use strum::Display;

use crate::{
    cli::OutputFormat,
    ntds::{DataTableRecord, InterestingAcl},
    win32_types::Rdn,
    CDatabase,
};

#[derive(Subcommand, Clone, Display, Eq, PartialEq)]
pub enum Command {
    /// list all available extended rights
    List {
        /// Output format
        #[clap(value_enum, short('F'), long("format"), default_value_t = OutputFormat::Csv)]
        format: OutputFormat,
    },

    /// list all objects (and subjects) with some specific extended right
    Objects {
        /// name or GUID of the extended right to search for
        extended_right: String,
    },
}

impl Command {
    pub fn run<'info, 'db>(&self, database: &CDatabase<'info, 'db>) -> Result<()> {
        match self {
            Command::List { format } => self.list_extended_rights(database, format),
            Command::Objects { extended_right } => {
                self.list_extended_right_subjects(database, extended_right)
            }
        }
    }

    fn list_extended_rights<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        format: &OutputFormat,
    ) -> Result<()> {
        #[derive(Serialize)]
        struct ExtendedRight {
            object_name: String,
            rights_guid: Guid,
            object_guid: Guid,
        }

        if let Ok(extended_rights) = database
            .data_table()
            .object_tree()
            .get_special_records()?
            .extended_rights_container()
            .children()
            .lock()
        {
            let mut writer = format.create_writer()?;
            for right in extended_rights.iter() {
                let object = database
                    .entry(crate::EntryId::Id(*right.record_ptr().ds_record_id()))?
                    .expect("illegal record id in object tree");

                writer.write(ExtendedRight {
                    object_name: object.att_object_name()?.to_string(),
                    rights_guid: object
                        .att_rights_guid_opt()?
                        .expect("extended right without rights GUID"),
                    object_guid: object.att_object_guid()?,
                })?;
            }
            drop(writer);
        } else {
            bail!("unable to acquire lock to extended rights container")
        }

        Ok(())
    }

    fn list_extended_right_subjects<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        extended_right: &str,
    ) -> std::result::Result<(), anyhow::Error> {
        let format = OutputFormat::Json;

        #[derive(Serialize, Eq, PartialEq, Hash, Clone)]
        struct ObjectInfo {
            name: Option<Rdn>,
            guid: Option<Guid>,
            sid: Option<String>,
        }

        let right_guid = self.resolve_extended_right(database, extended_right)?;
        let sd_map = self.find_sd_with_extended_right(database, Some(right_guid))?;

        let mut extended_rights_mapping: HashMap<String, HashSet<ObjectInfo>> = HashMap::new();

        for (sd_id, acl) in sd_map.iter() {
            let mut objects = HashSet::new();
            for obj in self.find_objects_with_sd(database, *sd_id)? {
                let name = obj.att_object_name_opt()?;
                let guid = obj.att_object_guid_opt()?;
                let sid = obj.att_object_sid_opt()?.map(|sid| sid.to_string());
                let obj_info = ObjectInfo { name, guid, sid };
                objects.insert(obj_info);
            }

            for sid in acl.sids().iter().map(Sid::to_string) {
                match extended_rights_mapping.entry(sid) {
                    Entry::Occupied(mut sid_entry) => {
                        for object_info in objects.iter() {
                            (*sid_entry.get_mut()).insert(object_info.clone());
                        }
                    }
                    Entry::Vacant(sid_entry) => {
                        sid_entry.insert(objects.clone());
                    }
                }
            }
        }

        let mut writer = format.create_writer()?;
        writer.write(extended_rights_mapping)?;
        drop(writer);

        Ok(())
    }

    fn resolve_extended_right<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        extended_right: &str,
    ) -> std::result::Result<Guid, anyhow::Error> {
        if let Ok(guid) = Guid::from_str(extended_right) {
            Ok(guid)
        } else if let Ok(extended_rights) = database
            .data_table()
            .object_tree()
            .get_special_records()?
            .extended_rights_container()
            .children()
            .lock()
        {
            for right in extended_rights.iter() {
                let object = database
                    .entry(crate::EntryId::Id(*right.record_ptr().ds_record_id()))?
                    .expect("illegal record id in object tree");
                if object.att_object_name()?.as_str() == extended_right {
                    return Ok(object
                        .att_rights_guid_opt()?
                        .expect("extended right without rights GUID"));
                }
            }
            bail!("unknown extended right: '{extended_right}'");
        } else {
            bail!("unable to acquire lock to Extended Rights container")
        }
    }

    /// If `right_guid` is `None`, then search for ACEs with the
    /// CONTROL_ACCESS flag set, but without rights guid (which means
    /// that all extended rights are permitted)
    fn find_sd_with_extended_right<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        right_guid: Option<Guid>,
    ) -> std::result::Result<HashMap<i64, InterestingAcl>, anyhow::Error> {
        Ok(database.data_table().find_acls(move |ace| {
            if ace.header().mask().contains(AccessMaskFlag::CONTROL_ACCESS) {
                match right_guid {
                    None if ace.object_type().is_none() => Some(ace.sid()),
                    Some(guid) => {
                        if let Some(object_type) = ace.object_type() {
                            if guid == *object_type {
                                Some(ace.sid())
                            } else {
                                None
                            }
                        } else {
                            None
                        }
                    }
                    _ => None,
                }
            } else {
                None
            }
        }))
    }

    fn find_objects_with_sd<'info, 'db>(
        &self,
        database: &'info CDatabase<'info, 'db>,
        sd_id: i64,
    ) -> anyhow::Result<Vec<DataTableRecord<'info, 'db>>> {
        let mut objects = Vec::new();
        if let Some(record_ids) = database
            .data_table()
            .data_table()
            .metadata()
            .sd_objects()
            .get(&sd_id)
        {
            for record_id in record_ids {
                let record = database
                    .data_table()
                    .entry(crate::EntryId::Id(*record_id))?
                    .expect("illegal record pointer");
                objects.push(record);
            }
        }
        Ok(objects)
    }
}