use crate::error::Error; use crate::process::PtyProcess;
pub use crate::reader::ReadUntil;
use crate::reader::{NBReader, Regex};
use std::fs::File;
use std::io::prelude::*;
use std::io::LineWriter;
use std::ops::{Deref, DerefMut};
use std::process::Command;
use tempfile;
pub struct StreamSession<W: Write> {
pub writer: LineWriter<W>,
pub reader: NBReader,
}
impl<W: Write> StreamSession<W> {
pub fn new<R: Read + Send + 'static>(reader: R, writer: W, timeout_ms: Option<u64>) -> Self {
Self {
writer: LineWriter::new(writer),
reader: NBReader::new(reader, timeout_ms),
}
}
pub fn send_line(&mut self, line: &str) -> Result<usize, Error> {
let mut len = self.send(line)?;
len += self.writer.write(&[b'\n'])?;
Ok(len)
}
pub fn send(&mut self, s: &str) -> Result<usize, Error> {
self.writer.write(s.as_bytes()).map_err(Error::from)
}
pub fn send_control(&mut self, c: char) -> Result<(), Error> {
let code = match c {
'a'..='z' => c as u8 + 1 - b'a',
'A'..='Z' => c as u8 + 1 - b'A',
'[' => 27,
'\\' => 28,
']' => 29,
'^' => 30,
'_' => 31,
_ => return Err(Error::SendContolError(c)),
};
self.writer.write_all(&[code])?;
self.writer.flush()?;
Ok(())
}
pub fn flush(&mut self) -> Result<(), Error> {
self.writer.flush().map_err(Error::from)
}
pub fn read_line(&mut self) -> Result<String, Error> {
match self.exp(&ReadUntil::String('\n'.to_string())) {
Ok((mut line, _)) => {
if line.ends_with('\r') {
line.pop().expect("this never happens");
}
Ok(line)
}
Err(e) => Err(e),
}
}
pub fn try_read(&mut self) -> Option<char> {
self.reader.try_read()
}
fn exp(&mut self, needle: &ReadUntil) -> Result<(String, String), Error> {
self.reader.read_until(needle)
}
pub fn exp_eof(&mut self) -> Result<String, Error> {
self.exp(&ReadUntil::EOF).map(|(_, s)| s)
}
pub fn exp_regex(&mut self, regex: &str) -> Result<(String, String), Error> {
self.exp(&ReadUntil::Regex(Regex::new(regex)?))
}
pub fn exp_string(&mut self, needle: &str) -> Result<String, Error> {
self.exp(&ReadUntil::String(needle.to_string()))
.map(|(s, _)| s)
}
pub fn exp_char(&mut self, needle: char) -> Result<String, Error> {
self.exp(&ReadUntil::String(needle.to_string()))
.map(|(s, _)| s)
}
pub fn exp_any(&mut self, needles: Vec<ReadUntil>) -> Result<(String, String), Error> {
self.exp(&ReadUntil::Any(needles))
}
}
#[allow(dead_code)]
pub struct PtySession {
pub process: PtyProcess,
pub stream: StreamSession<File>,
}
impl Deref for PtySession {
type Target = StreamSession<File>;
fn deref(&self) -> &StreamSession<File> {
&self.stream
}
}
impl DerefMut for PtySession {
fn deref_mut(&mut self) -> &mut StreamSession<File> {
&mut self.stream
}
}
impl PtySession {
fn new(process: PtyProcess, timeout_ms: Option<u64>) -> Result<Self, Error> {
let f = process.get_file_handle();
let reader = f.try_clone()?;
let stream = StreamSession::new(reader, f, timeout_ms);
Ok(Self { process, stream })
}
}
fn tokenize_command(program: &str) -> Result<Vec<String>, Error> {
comma::parse_command(program).ok_or(Error::BadProgramArguments)
}
pub fn spawn(program: &str, timeout_ms: Option<u64>) -> Result<PtySession, Error> {
if program.is_empty() {
return Err(Error::EmptyProgramName);
}
let mut parts = tokenize_command(program)?;
let prog = parts.remove(0);
let mut command = Command::new(prog);
command.args(parts);
spawn_command(command, timeout_ms)
}
pub fn spawn_command(command: Command, timeout_ms: Option<u64>) -> Result<PtySession, Error> {
let mut process = PtyProcess::new(command)?;
process.set_kill_timeout(timeout_ms);
PtySession::new(process, timeout_ms)
}
pub struct PtyReplSession {
pub prompt: String,
pub pty_session: PtySession,
pub quit_command: Option<String>,
pub echo_on: bool,
}
impl PtyReplSession {
pub fn wait_for_prompt(&mut self) -> Result<String, Error> {
self.pty_session.exp_string(&self.prompt)
}
pub fn execute(&mut self, cmd: &str, ready_regex: &str) -> Result<(), Error> {
self.send_line(cmd)?;
if self.echo_on {
self.exp_string(cmd)?;
}
self.exp_regex(ready_regex)?;
Ok(())
}
pub fn send_line(&mut self, line: &str) -> Result<usize, Error> {
let bytes_written = self.pty_session.send_line(line)?;
if self.echo_on {
self.exp_string(line)?;
}
Ok(bytes_written)
}
}
impl Deref for PtyReplSession {
type Target = PtySession;
fn deref(&self) -> &PtySession {
&self.pty_session
}
}
impl DerefMut for PtyReplSession {
fn deref_mut(&mut self) -> &mut PtySession {
&mut self.pty_session
}
}
impl Drop for PtyReplSession {
fn drop(&mut self) {
if let Some(ref cmd) = self.quit_command {
self.pty_session
.send_line(cmd)
.expect("could not run `exit` on bash process");
}
}
}
pub fn spawn_bash(timeout: Option<u64>) -> Result<PtyReplSession, Error> {
let mut rcfile = tempfile::NamedTempFile::new().unwrap();
let _ = rcfile
.write(
b"include () { [[ -f \"$1\" ]] && source \"$1\"; }\n\
include /etc/bash.bashrc\n\
include ~/.bashrc\n\
PS1=\"~~~~\"\n\
unset PROMPT_COMMAND\n",
)
.expect("cannot write to tmpfile");
let mut c = Command::new("bash");
c.args(&[
"--rcfile",
rcfile.path().to_str().unwrap_or("temp file does not exist"),
]);
spawn_command(c, timeout).and_then(|p| {
let new_prompt = "[REXPECT_PROMPT>";
let mut pb = PtyReplSession {
prompt: new_prompt.to_string(),
pty_session: p,
quit_command: Some("quit".to_string()),
echo_on: false,
};
pb.exp_string("~~~~")?;
rcfile.close()?;
pb.send_line(&("PS1='".to_string() + new_prompt + "'"))?;
pb.wait_for_prompt()?;
Ok(pb)
})
}
pub fn spawn_python(timeout: Option<u64>) -> Result<PtyReplSession, Error> {
spawn_command(Command::new("python"), timeout).map(|p| PtyReplSession {
prompt: ">>> ".to_string(),
pty_session: p,
quit_command: Some("exit()".to_string()),
echo_on: true,
})
}
pub fn spawn_stream<R: Read + Send + 'static, W: Write>(
reader: R,
writer: W,
timeout_ms: Option<u64>,
) -> StreamSession<W> {
StreamSession::new(reader, writer, timeout_ms)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_read_line() -> Result<(), Error> {
let mut s = spawn("cat", Some(100000))?;
s.send_line("hans")?;
assert_eq!("hans", s.read_line()?);
let should = crate::process::wait::WaitStatus::Signaled(
s.process.child_pid,
crate::process::signal::Signal::SIGTERM,
false,
);
assert_eq!(should, s.process.exit()?);
Ok(())
}
#[test]
fn test_expect_eof_timeout() -> Result<(), Error> {
let mut p = spawn("sleep 3", Some(1000)).expect("cannot run sleep 3");
match p.exp_eof() {
Ok(_) => panic!("should raise Timeout"),
Err(Error::Timeout { .. }) => {}
Err(_) => panic!("should raise TimeOut"),
}
Ok(())
}
#[test]
fn test_expect_eof_timeout2() {
let mut p = spawn("sleep 1", Some(1100)).expect("cannot run sleep 1");
assert!(p.exp_eof().is_ok(), "expected eof");
}
#[test]
fn test_expect_string() -> Result<(), Error> {
let mut p = spawn("cat", Some(1000)).expect("cannot run cat");
p.send_line("hello world!")?;
p.exp_string("hello world!")?;
p.send_line("hello heaven!")?;
p.exp_string("hello heaven!")?;
Ok(())
}
#[test]
fn test_read_string_before() -> Result<(), Error> {
let mut p = spawn("cat", Some(1000)).expect("cannot run cat");
p.send_line("lorem ipsum dolor sit amet")?;
assert_eq!("lorem ipsum dolor sit ", p.exp_string("amet")?);
Ok(())
}
#[test]
fn test_expect_any() -> Result<(), Error> {
let mut p = spawn("cat", Some(1000)).expect("cannot run cat");
p.send_line("Hi")?;
match p.exp_any(vec![
ReadUntil::NBytes(3),
ReadUntil::String("Hi".to_string()),
]) {
Ok(s) => assert_eq!(("".to_string(), "Hi".to_string()), s),
Err(e) => panic!("got error: {}", e),
}
Ok(())
}
#[test]
fn test_expect_empty_command_error() {
let p = spawn("", Some(1000));
match p {
Ok(_) => panic!("should raise an error"),
Err(Error::EmptyProgramName) => {}
Err(_) => panic!("should raise EmptyProgramName"),
}
}
#[test]
fn test_kill_timeout() -> Result<(), Error> {
let mut p = spawn_bash(Some(1000))?;
p.execute("cat <(echo ready) -", "ready")?;
Ok(())
}
#[test]
fn test_bash() -> Result<(), Error> {
let mut p = spawn_bash(Some(1000))?;
p.send_line("cd /tmp/")?;
p.wait_for_prompt()?;
p.send_line("pwd")?;
assert_eq!("/tmp\r\n", p.wait_for_prompt()?);
Ok(())
}
#[test]
fn test_bash_control_chars() -> Result<(), Error> {
let mut p = spawn_bash(Some(1000))?;
p.execute("cat <(echo ready) -", "ready")?;
p.send_control('c')?; p.wait_for_prompt()?;
p.execute("cat <(echo ready) -", "ready")?;
p.send_control('z')?; p.exp_regex(r"(Stopped|suspended)\s+cat .*")?;
p.send_line("fg")?;
p.execute("cat <(echo ready) -", "ready")?;
p.send_control('c')?;
Ok(())
}
#[test]
fn test_tokenize_command() {
let res = tokenize_command("prog arg1 arg2");
assert_eq!(vec!["prog", "arg1", "arg2"], res.unwrap());
let res = tokenize_command("prog -k=v");
assert_eq!(vec!["prog", "-k=v"], res.unwrap());
let res = tokenize_command("prog 'my text'");
assert_eq!(vec!["prog", "my text"], res.unwrap());
let res = tokenize_command(r#"prog "my text""#);
assert_eq!(vec!["prog", r#"my text"#], res.unwrap());
let res = tokenize_command(r#"prog "my text with quote'""#);
assert_eq!(vec!["prog", r#"my text with quote'"#], res.unwrap());
let res = tokenize_command(r#"prog "my broken text"#);
assert!(res.is_err());
}
}