super_shell/
lib.rs

1use std::io::{BufRead, BufReader, Write};
2
3use std::process::{Command, Stdio};
4
5/// Represents a shell with root privileges.
6pub struct RootShell {
7    shell_process: std::process::Child,
8}
9
10/// String that is appended to the end of each command, to indicate that the command has finished.
11const END_OF_COMMAND: &str = "~end-of-command~";
12
13/// Implementation of RootShell
14impl RootShell {
15    /// Creates a new root shell.
16    /// Defaults to pkexec as super user provider and sh as shell.
17    /// Returns None if the root shell could not be created.
18    /// # Example
19    /// ```
20    /// use super_shell::RootShell;
21    /// let mut root_shell = RootShell::new().expect("Failed to crate root shell");
22    /// println!("{}", root_shell.execute("echo Hello $USER"));
23    /// ```
24    #[cfg(target_os = "linux")]
25    pub fn new() -> Option<Self> {
26        Self::spawn_root_shell("pkexec", "sh")
27    }
28
29    #[cfg(target_os = "windows")]
30    pub fn new() -> Option<Self> {
31        None
32    }
33
34    #[cfg(target_os = "windows")]
35    pub fn new_custom(super_user_provider: &str, shell: &str) -> Option<Self> {
36        None
37    }
38
39    /// Creates a new root shell with the specified super user provider and shell.
40    /// Returns None if the root shell could not be created.
41    /// sudo as an interactive super user provider is currently not supported.
42    /// # Parameter
43    /// * `super_user_provider` - The command to use to get super user privileges
44    /// * `shell` - The shell to use
45    /// # Example
46    /// ```
47    /// use super_shell::RootShell;
48    /// let mut root_shell = RootShell::new_custom("gksu", "bash").expect("Failed to crate root shell");
49    /// println!("{}", root_shell.execute("echo Hello $USER"));
50    /// ```
51    #[cfg(target_os = "linux")]
52    pub fn new_custom(super_user_provider: &str, shell: &str) -> Option<Self> {
53        Self::spawn_root_shell(super_user_provider, shell)
54    }
55
56    /// Creates a new root shell with the specified super user provider and shell.
57    /// Returns None if the root shell could not be created.
58    /// # Parameter
59    /// * `super_user_provider` - The command to use to get super user privileges
60    /// * `shell` - The shell to use
61    fn spawn_root_shell(super_user_provider: &str, shell: &str) -> Option<RootShell> {
62        let shell_process = Command::new(super_user_provider)
63            .arg(shell)
64            .stdin(Stdio::piped())
65            .stdout(Stdio::piped())
66            .stderr(Stdio::inherit())
67            .spawn()
68            .unwrap();
69
70        let mut root_shell = Self { shell_process };
71
72        // Make sure we are root now
73        let user = root_shell.execute("whoami");
74        if !user.trim().eq("root") {
75            return None;
76        }
77
78        Some(root_shell)
79    }
80
81    /// Executes a command in the root shell and returns the output trimmed.
82    /// Blocks the current thread until the command is finished.
83    /// # Parameter
84    /// * `command` - The command to execute
85    /// # Example
86    /// ```
87    /// use super_shell::RootShell;
88    /// let mut root_shell = RootShell::new().expect("Failed to crate root shell");
89    /// assert!(root_shell.execute("echo Hello $USER").trim().eq("Hello root"));
90    /// ```
91    pub fn execute(&mut self, command: impl AsRef<str>) -> String {
92        // If not linux return empty string
93        #[cfg(not(target_os = "linux"))]
94        return "".to_string();
95
96        // Append end of command string to the command
97        let command = command.as_ref().to_string();
98
99        // Write the actual command to stdin of the root shell
100        let mut shell_stdin = self.shell_process.stdin.as_mut().unwrap();
101        writeln!(&mut shell_stdin, "{}", command).unwrap();
102        shell_stdin.flush().unwrap();
103
104        // Write "end of command" string to stdin of the root shell
105        writeln!(&mut shell_stdin, "echo {}", END_OF_COMMAND).unwrap();
106        shell_stdin.flush().unwrap();
107
108        // Read piped stdout from the root shell
109        let stdout = self.shell_process.stdout.as_mut().unwrap();
110        let mut stdout_reader = BufReader::new(stdout);
111
112        // Read until the "end of command" string is found
113        let mut string_data = String::new();
114        loop {
115            let mut line = String::new();
116            stdout_reader.read_line(&mut line).unwrap();
117            string_data.push_str(&line);
118
119            if line.contains(END_OF_COMMAND) {
120                break;
121            }
122        }
123
124        // Clean up the string
125        let cmd_response = string_data.replace(END_OF_COMMAND, "");
126        let cmd_response = cmd_response.trim().to_string();
127
128        cmd_response
129    }
130
131    /// Exits the root shell and waits for the process to finish.
132    pub fn exit(&mut self) {
133        let mut shell_stdin = self.shell_process.stdin.as_mut().unwrap();
134        writeln!(&mut shell_stdin, "exit").unwrap();
135        shell_stdin.flush().unwrap();
136        self.shell_process.wait().expect("failed to wait on child");
137    }
138}
139
140/// Drop implementation for RootShell.
141impl Drop for RootShell {
142    /// Exits the root shell and waits for the process to finish.
143    fn drop(&mut self) {
144        self.exit();
145        self.shell_process.wait().expect("failed to wait on child");
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use crate::RootShell;
152
153    #[test]
154    fn test_custom() {
155        let mut root_shell = RootShell::new_custom("sudo", "bash").unwrap();
156        assert!(root_shell
157            .execute("echo Hello $USER")
158            .trim()
159            .eq("Hello root"));
160    }
161}
super_shell/lib.rs

super_shell/
lib.rs
