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

use std::io::{BufRead, BufReader, Write};

use std::process::{Command, Stdio};

/// Represents a shell with root privileges
/// The shell is automatically closed when the struct is dropped
/// The shell is opened with the pkexec command
/// It uses the sh shell
pub struct RootShell {
    shell_process: std::process::Child,
}

/// String that is appended to the end of each command, to indicate that the command has finished
const END_OF_COMMAND: &str = "~end-of-command~";

/// Implementation of RootShell
impl RootShell {
    /// Creates a new root shell
    /// Returns None if the root shell could not be created
    /// or if the user did not enter the password
    pub fn new() -> Option<Self> {
        let shell_process = Command::new("pkexec")
            .arg("sh")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()
            .unwrap();

        let mut root_shell = Self { shell_process };

        // Make sure we are root now
        let user = root_shell.execute("whoami");
        if !user.trim().eq("root") {
            return None;
        }

        Some(root_shell)
    }

    /// Executes a command in the root shell and returns the output trimmed
    /// Blocks the current thread until the command is finished
    pub fn execute(&mut self, command: impl AsRef<str>) -> String {
        // Append end of command string to the command
        let command = command.as_ref().to_string();

        // Write the actual command to stdin of the root shell
        let mut shell_stdin = self.shell_process.stdin.as_mut().unwrap();
        writeln!(&mut shell_stdin, "{}", command).unwrap();
        shell_stdin.flush().unwrap();

        // Write "end of command" string to stdin of the root shell
        writeln!(&mut shell_stdin, "echo {}", END_OF_COMMAND).unwrap();
        shell_stdin.flush().unwrap();

        // Read piped stdout from the root shell
        let stdout = self.shell_process.stdout.as_mut().unwrap();
        let mut stdout_reader = BufReader::new(stdout);

        // Read until the "end of command" string is found
        let mut string_data = String::new();
        loop {
            let mut line = String::new();
            stdout_reader.read_line(&mut line).unwrap();
            string_data.push_str(&line);

            if line.contains(END_OF_COMMAND) {
                break;
            }
        }

        // Clean up the string
        let cmd_response = string_data.replace(END_OF_COMMAND, "");
        let cmd_response = cmd_response.trim().to_string();

        cmd_response
    }

    /// Exits the root shell and waits for the process to finish
    pub fn exit(&mut self) {
        let mut shell_stdin = self.shell_process.stdin.as_mut().unwrap();
        writeln!(&mut shell_stdin, "exit").unwrap();
        shell_stdin.flush().unwrap();
        self.shell_process.wait().expect("failed to wait on child");
    }
}

/// Drop implementation for RootShell
impl Drop for RootShell {
    /// Exits the root shell and waits for the process to finish
    fn drop(&mut self) {
        self.exit();
        self.shell_process.wait().expect("failed to wait on child");
    }
}