shell_client 0.1.0

Rust进程调试工具,封装了与Unix域套接字通信、命令行自动补全、以及内置命令处理等功能。Client 可以连接到一个进程,发送自定义命令,并处理来自该进程的输出。实现了命令行读取、命令解析、进程查找、自动补全、命令执行等功能,并提供了一个循环来不断读取用户输入并执行相应的命令
Documentation
use crate::{autocomplete_reader::AutoCompleteReader, sys::get_process_list};
use shell_core::*;
use std::{
    os::unix::net::UnixStream,
    sync::{Arc, Mutex},
    thread::{sleep, spawn, JoinHandle},
    time::Duration,
};


pub struct Client {
    cmd_channel: Option<UnixStream>,
    output_channel: Option<UnixStream>,
    copy_stdout: Option<JoinHandle<()>>,
    reader: Arc<Mutex<Box<AutoCompleteReader>>>,
}

static DEFAULT_PS1: &str = "\x1B[33m>> \x1B[0m";

impl Default for Client {
    fn default() -> Self {
        Self::new()
    }
}

impl Client {
    pub fn new() -> Client {
        Client {
            cmd_channel: None,
            output_channel: None,
            copy_stdout: None,
            reader: AutoCompleteReader::new().unwrap(),
        }
    }

    fn find_process(&self, arg: &Argument) -> Vec<(String, u64)> {
        let result: Vec<(String, u64)> = get_process_list()
            .into_iter()
            .map(|(pid, name)| (name, pid.parse().unwrap()))
            .collect();

        match arg {
            Argument::Str(name) => result
                .into_iter()
                .filter(|(name_, _)| name_.starts_with(name))
                .collect::<Vec<(String, u64)>>(),
            Argument::Int(pid) => result
                .into_iter()
                .filter(|(_, pid_)| pid.to_owned() as u64 == *pid_)
                .collect(),
        }
    }

    fn make_uds_path(pid: &u64) -> (String, String) {
        (
            format!("/tmp/rust_shell_cmd_{}", pid),
            format!("/tmp/rust_shell_output_{}", pid),
        )
    }

    fn parse_auto_complete(line: &str) -> Vec<String> {
        line.split_whitespace().map(|s| s.to_owned()).collect()
    }

    fn attach_process(&mut self, args: &Vec<Argument>) -> Result<(), String> {
        if args.len() != 1 {
            return Err("argument number error".to_string());
        }

        let pids = self.find_process(&args[0]);
        if pids.is_empty() {
            return Err("process not found".to_string());
        }
        if pids.len() != 1 {
            return Err(format!("multiple process found: \n{}", {
                pids.iter()
                    .map(|(comm, pid)| format!("{}: {}", comm, pid))
                    .collect::<Vec<String>>()
                    .join("\n")
            }));
        }

        let (cmd_path, output_path) = Self::make_uds_path(&pids[0].1);
        self.cmd_channel = Some(UnixStream::connect(&cmd_path).map_err(|err| err.to_string())?);
        self.output_channel =
            Some(UnixStream::connect(&output_path).map_err(|err| err.to_string())?);

        if let Some(c) = &self.cmd_channel {
            self.reader
                .lock()
                .map_err(|err| err.to_string())?
                .append_debug_command_complete_data(
                    Client::parse_auto_complete(&read_line(
                        &mut c.try_clone().map_err(|err| err.to_string())?,
                    )?)
                    .into_iter()
                    .map(|x| (x.clone(), x.clone()))
                    .collect(),
                )
        }

        let mut output_channel_copy = self
            .output_channel
            .as_ref()
            .ok_or("not attach to process")?
            .try_clone()
            .map_err(|err| err.to_string())?;

        self.copy_stdout = Some(spawn(move || {
            let _ = std::io::copy(&mut output_channel_copy, &mut std::io::stdout());
        }));

        self.reader
            .lock()
            .map_err(|err| err.to_string())?
            .set_prompt(format!("\x1B[32m{} >> \x1B[0m", pids[0].0).as_str());

        Ok(())
    }

    fn exit() -> Result<(), String> {
        Err("exit".to_owned())
    }



    pub fn run_custom_command(&mut self, line: &String) -> Result<(), String> {
        match self.cmd_channel {
            None => Err("not attach to process".to_owned()),
            Some(ref mut cmd_channel) => write_line(cmd_channel, line),
        }
    }

    pub fn detach_process(&mut self) {
        self.cmd_channel = None;
        self.output_channel = None;
        self.copy_stdout = None;
        self.reader
            .lock()
            .expect("lock reader failed")
            .set_prompt(DEFAULT_PS1);
    }


    fn init_reader(&mut self) -> Result<(), String> {
        let mut r = self.reader.lock().map_err(|err| err.to_string())?;
        r.set_prompt(DEFAULT_PS1);
        r.set_debug_command_complete_data(vec![
            ("exit".to_owned(), "exit".to_owned()),
            ("attach".to_owned(), "attach".to_owned()),
            ("detach".to_owned(), "detach".to_owned()),
        ]);

        Ok(())
    }


    fn run_builtin_command(&mut self, cmd: &String, args: &Vec<Argument>) -> Result<(), String> {
        match cmd.as_str() {
            "attach" => self.attach_process(args),
            "detach" => {
                self.detach_process();
                Ok(())
            }
            "exit" => Self::exit(),
            _ => Err("custom".to_owned()),
        }
    }

    pub fn run(&mut self) -> Result<(), String> {
        self.init_reader()?;
        loop {
            let line: String;
            {
                line = self
                    .reader
                    .lock()
                    .expect("lock reader failed")
                    .read()
                    .expect("read failed");
            }
            if line.is_empty() {
                continue;
            }
            if let Some((cmd, args)) = split_command(line.trim()) {
                if let Err(err) = self.run_builtin_command(&cmd, &parse_arguments(&args)) {
                    match err.as_str() {
                        "exit" => break,
                        "custom" => {
                            if let Err(err) = self.run_custom_command(&line) {
                                println!("Error: {}", err);
                                self.detach_process()
                            }
                            sleep(Duration::from_millis(10));
                        }
                        _ => {
                            println!("Error: {}", err);
                        }
                    }
                }
            }
        }

        Ok(())
    }
}