roguewave 0.1.0

An imperative remote server automation tool
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::collections::BTreeMap;

use anyhow::Context;

use crate::Session;

impl Session {
    /// Read environment variables for the specified user. If no user if specified, the current user
    /// is used.
    pub async fn env(&mut self, user: Option<&str>) -> anyhow::Result<&BTreeMap<String, String>> {
        let cache_has_user = self
            .cache()
            .get::<EnvCache>()
            .map_or(false, |c| c.has_user(user));

        if !cache_has_user {
            let output = self
                .command(["env"])
                .user(user)
                .hide_command()
                .hide_stdout()
                .run()
                .await?;
            let mut env = BTreeMap::new();
            for line in output.stdout.split('\n') {
                if line.is_empty() {
                    continue;
                }
                let line = line.trim_end_matches('\n');
                let mut iter = line.splitn(2, '=');
                let name = iter.next().unwrap();
                let value = iter.next().context("missing '=' in env output")?;
                env.insert(name.to_string(), value.to_string());
            }
            let cache = self
                .cache()
                .entry::<EnvCache>()
                .or_insert_with(EnvCache::default);
            if let Some(user) = user {
                cache.other_users.insert(user.into(), env);
            } else {
                cache.current_user = Some(env);
            }
        }

        let cache = self.cache().get::<EnvCache>().unwrap();
        if let Some(user) = user {
            Ok(cache.other_users.get(user).unwrap())
        } else {
            Ok(cache.current_user.as_ref().unwrap())
        }
    }

    /// Fetch home directory location for the specified user. If no user if specified, the current user
    /// is used.
    pub async fn home_dir(&mut self, user: Option<&str>) -> anyhow::Result<String> {
        let env = self.env(user).await?;
        env.get("HOME")
            .context("missing remote env var \"HOME\"")
            .cloned()
    }

    /// Fetch the current user name.
    pub async fn current_user(&mut self) -> anyhow::Result<String> {
        let env = self.env(None).await?;
        env.get("USER")
            .context("missing remote env var \"USER\"")
            .cloned()
    }

    /// Fetch shell command for the specified user. If no user if specified, the current user
    /// is used.
    pub async fn shell(&mut self, user: Option<&str>) -> anyhow::Result<String> {
        let env = self.env(user).await?;
        env.get("SHELL")
            .context("missing remote env var \"SHELL\"")
            .cloned()
    }

    /// Set shell command for the specified user. If no user if specified, the current user
    /// is used.
    pub async fn set_shell(&mut self, shell: &str, user: Option<&str>) -> anyhow::Result<()> {
        if self.shell(user).await? != shell {
            let mut command = self.command(["chsh", "-s", shell]);
            if let Some(user) = user {
                command = command.arg(user);
            }
            command.run().await?;
        }
        Ok(())
    }
}

#[derive(Default)]
struct EnvCache {
    current_user: Option<BTreeMap<String, String>>,
    other_users: BTreeMap<String, BTreeMap<String, String>>,
}

impl EnvCache {
    fn has_user(&self, user: Option<&str>) -> bool {
        if let Some(user) = user {
            self.other_users.contains_key(user)
        } else {
            self.current_user.is_some()
        }
    }
}