234 lines · 7.0 KB · Rust
RawDownload
use std::collections::BTreeMap; #[derive(Debug, Clone, PartialEq, Eq)]pub struct ProcessIdentity {    pub uid: u32,    pub gid: u32,    pub euid: u32,    pub egid: u32,    pub suid: u32,    pub sgid: u32,    pub supplementary_gids: Vec<u32>,} impl Default for ProcessIdentity {    fn default() -> Self {        Self {            uid: 1000,            gid: 1000,            euid: 1000,            egid: 1000,            suid: 1000,            sgid: 1000,            supplementary_gids: vec![1000],        }    }} #[derive(Debug, Clone, PartialEq, Eq)]pub struct UserAccount {    pub uid: u32,    pub gid: u32,    pub username: String,    pub homedir: String,    pub shell: String,    pub gecos: String,    pub supplementary_gids: Vec<u32>,} #[derive(Debug, Clone, PartialEq, Eq)]pub struct GroupRecord {    pub gid: u32,    pub name: String,    pub members: Vec<String>,} #[derive(Debug, Clone, Default, PartialEq, Eq)]pub struct UserConfig {    pub uid: Option<u32>,    pub gid: Option<u32>,    pub euid: Option<u32>,    pub egid: Option<u32>,    pub username: Option<String>,    pub homedir: Option<String>,    pub shell: Option<String>,    pub gecos: Option<String>,    pub group_name: Option<String>,    /// Supplementary groups are VM configuration, not guest-mutable state.    /// The primary gid is always injected and duplicate gids are dropped.    pub supplementary_gids: Vec<u32>,    pub accounts: Vec<UserAccount>,    pub groups: Vec<GroupRecord>,} #[derive(Debug, Clone, PartialEq, Eq)]pub struct UserManager {    pub uid: u32,    pub gid: u32,    pub euid: u32,    pub egid: u32,    pub username: String,    pub homedir: String,    pub shell: String,    pub gecos: String,    pub group_name: String,    pub supplementary_gids: Vec<u32>,    accounts_by_uid: BTreeMap<u32, UserAccount>,    account_uids_by_name: BTreeMap<String, u32>,    groups_by_gid: BTreeMap<u32, GroupRecord>,    group_gids_by_name: BTreeMap<String, u32>,} impl Default for UserManager {    fn default() -> Self {        Self::from_config(UserConfig::default())    }} impl UserManager {    pub fn new() -> Self {        Self::default()    }     pub fn from_config(config: UserConfig) -> Self {        let uid = config.uid.unwrap_or(1000);        let gid = config.gid.unwrap_or(1000);        let username = config.username.unwrap_or_else(|| String::from("agentos"));        let supplementary_gids = normalize_supplementary_gids(gid, config.supplementary_gids);         let primary_account = UserAccount {            uid,            gid,            username: username.clone(),            homedir: config                .homedir                .clone()                .unwrap_or_else(|| String::from("/home/agentos")),            shell: config                .shell                .clone()                .unwrap_or_else(|| String::from("/bin/sh")),            gecos: config.gecos.clone().unwrap_or_default(),            supplementary_gids: supplementary_gids.clone(),        };        let mut accounts_by_uid = BTreeMap::new();        let mut account_uids_by_name = BTreeMap::new();        for mut account in config.accounts {            account.supplementary_gids =                normalize_supplementary_gids(account.gid, account.supplementary_gids);            account_uids_by_name.insert(account.username.clone(), account.uid);            accounts_by_uid.insert(account.uid, account);        }        account_uids_by_name.insert(primary_account.username.clone(), primary_account.uid);        accounts_by_uid.insert(primary_account.uid, primary_account.clone());         let primary_group_name = config.group_name.unwrap_or_else(|| username.clone());        let mut groups_by_gid = BTreeMap::new();        let mut group_gids_by_name = BTreeMap::new();        for group in config.groups {            group_gids_by_name.insert(group.name.clone(), group.gid);            groups_by_gid.insert(group.gid, group);        }        groups_by_gid.entry(gid).or_insert_with(|| GroupRecord {            gid,            name: primary_group_name.clone(),            members: vec![username.clone()],        });        group_gids_by_name.insert(primary_group_name.clone(), gid);         Self {            uid,            gid,            euid: config.euid.unwrap_or(uid),            egid: config.egid.unwrap_or(gid),            username: username.clone(),            homedir: primary_account.homedir,            shell: primary_account.shell,            gecos: primary_account.gecos,            group_name: primary_group_name,            supplementary_gids,            accounts_by_uid,            account_uids_by_name,            groups_by_gid,            group_gids_by_name,        }    }     pub fn identity(&self) -> ProcessIdentity {        ProcessIdentity {            uid: self.uid,            gid: self.gid,            euid: self.euid,            egid: self.egid,            suid: self.euid,            sgid: self.egid,            supplementary_gids: self.supplementary_gids.clone(),        }    }     pub fn getgroups(&self) -> Vec<u32> {        self.supplementary_gids.clone()    }     pub fn getpwuid(&self, uid: u32) -> Option<String> {        self.accounts_by_uid.get(&uid).map(render_passwd)    }     pub fn getpwnam(&self, username: &str) -> Option<String> {        self.account_uids_by_name            .get(username)            .and_then(|uid| self.getpwuid(*uid))    }     pub fn getgrgid(&self, gid: u32) -> Option<String> {        self.groups_by_gid.get(&gid).map(render_group).or_else(|| {            let members = self                .accounts_by_uid                .values()                .filter(|account| account.supplementary_gids.contains(&gid))                .map(|account| account.username.as_str())                .collect::<Vec<_>>();            (!members.is_empty()).then(|| format!("group{gid}:x:{gid}:{}", members.join(",")))        })    }     pub fn getgrnam(&self, name: &str) -> Option<String> {        self.group_gids_by_name            .get(name)            .and_then(|gid| self.getgrgid(*gid))    }     pub fn passwd_entries(&self) -> Vec<String> {        self.accounts_by_uid.values().map(render_passwd).collect()    }     pub fn group_entries(&self) -> Vec<String> {        self.groups_by_gid.values().map(render_group).collect()    }     pub fn account(&self, uid: u32) -> Option<&UserAccount> {        self.accounts_by_uid.get(&uid)    }} fn render_passwd(account: &UserAccount) -> String {    format!(        "{}:x:{}:{}:{}:{}:{}",        account.username, account.uid, account.gid, account.gecos, account.homedir, account.shell    )} fn render_group(group: &GroupRecord) -> String {    format!("{}:x:{}:{}", group.name, group.gid, group.members.join(","))} fn normalize_supplementary_gids(primary_gid: u32, supplementary_gids: Vec<u32>) -> Vec<u32> {    let mut normalized = Vec::with_capacity(supplementary_gids.len() + 1);    normalized.push(primary_gid);    for gid in supplementary_gids {        if !normalized.contains(&gid) {            normalized.push(gid);        }    }    normalized}