use crate::process::PtyProcess;
use crate::reader::{NBReader, Regex};
pub use crate::reader::ReadUntil;
use std::fs::File;
use std::io::LineWriter;
use std::process::Command;
use std::io::prelude::*;
use std::ops::{Deref, DerefMut};
use crate::errors::*;
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> {
let mut len = self.send(line)?;
len += self.writer
.write(&['\n' as u8])
.chain_err(|| "cannot write newline")?;
Ok(len)
}
pub fn send(&mut self, s: &str) -> Result<usize> {
self.writer
.write(s.as_bytes())
.chain_err(|| "cannot write line to process")
}
pub fn send_control(&mut self, c: char) -> Result<()> {
let code = match c {
'a'..='z' => c as u8 + 1 - 'a' as u8,
'A'..='Z' => c as u8 + 1 - 'A' as u8,
'[' => 27,
'\\' => 28,
']' => 29,
'^' => 30,
'_' => 31,
_ => return Err(format!("I don't understand Ctrl-{}", c).into()),
};
self.writer
.write_all(&[code])
.chain_err(|| "cannot send control")?;
self.writer
.flush()
.chain_err(|| "cannot flush after sending ctrl keycode")?;
Ok(())
}
pub fn flush(&mut self) -> Result<()> {
self.writer.flush().chain_err(|| "could not flush")
}
pub fn read_line(&mut self) -> Result<String> {
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)> {
self.reader.read_until(needle)
}
pub fn exp_eof(&mut self) -> Result<String> {
self.exp(&ReadUntil::EOF).and_then(|(_, s)| Ok(s))
}
pub fn exp_regex(&mut self, regex: &str) -> Result<(String, String)> {
let res = self.exp(&ReadUntil::Regex(Regex::new(regex).chain_err(|| "invalid regex")?))
.and_then(|s| Ok(s));
res
}
pub fn exp_string(&mut self, needle: &str) -> Result<String> {
self.exp(&ReadUntil::String(needle.to_string()))
.and_then(|(s, _)| Ok(s))
}
pub fn exp_char(&mut self, needle: char) -> Result<String> {
self.exp(&ReadUntil::String(needle.to_string()))
.and_then(|(s, _)| Ok(s))
}
pub fn exp_any(&mut self, needles: Vec<ReadUntil>) -> Result<(String, String)> {
self.exp(&ReadUntil::Any(needles))
}
}
#[allow(dead_code)]
pub struct PtySession {
pub process: PtyProcess,
pub stream: StreamSession<File>,
pub commandname: String,
}
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>, commandname: String) -> Result<Self> {
let f = process.get_file_handle();
let reader = f.try_clone().chain_err(|| "couldn't open write stream")?;
let stream = StreamSession::new(reader, f, timeout_ms);
Ok(Self {
process,
stream,
commandname: commandname,
})
}
}
fn tokenize_command(program: &str) -> Vec<String> {
let re = Regex::new(r#""[^"]+"|'[^']+'|[^'" ]+"#).unwrap();
let mut res = vec![];
for cap in re.captures_iter(program) {
res.push(cap[0].to_string());
}
res
}
pub fn spawn(program: &str, timeout_ms: Option<u64>) -> Result<PtySession> {
if program.is_empty() {
return Err(ErrorKind::EmptyProgramName.into());
}
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> {
let commandname = format!("{:?}", &command);
let mut process = PtyProcess::new(command)
.chain_err(|| "couldn't start process")?;
process.set_kill_timeout(timeout_ms);
PtySession::new(process, timeout_ms, commandname)
}
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> {
self.pty_session.exp_string(&self.prompt)
}
pub fn execute(&mut self, cmd: &str, ready_regex: &str) -> Result<()> {
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> {
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> {
let mut rcfile = tempfile::NamedTempFile::new().unwrap();
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_else(|| return "temp file does not exist".into())]);
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().chain_err(|| "cannot delete temporary rcfile")?;
pb.send_line(&("PS1='".to_string() + new_prompt + "'"))?;
pb.wait_for_prompt()?;
Ok(pb)
})
}
pub fn spawn_python(timeout: Option<u64>) -> Result<PtyReplSession> {
spawn_command(Command::new("python"), timeout).and_then(|p| {
Ok(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<()> {
let mut s = spawn("cat", Some(1000))?;
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(())
}()
.unwrap_or_else(|e| panic!("test_read_line failed: {}", e));
}
#[test]
fn test_expect_eof_timeout() {
|| -> Result<()> {
let mut p = spawn("sleep 3", Some(1000)).expect("cannot run sleep 3");
match p.exp_eof() {
Ok(_) => assert!(false, "should raise Timeout"),
Err(Error(ErrorKind::Timeout(_, _, _), _)) => {}
Err(_) => assert!(false, "should raise TimeOut"),
}
Ok(())
}()
.unwrap_or_else(|e| panic!("test_timeout failed: {}", e));
}
#[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<()> {
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(())
}()
.unwrap_or_else(|e| panic!("test_expect_string failed: {}", e));
}
#[test]
fn test_read_string_before() {
|| -> Result<()> {
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(())
}()
.unwrap_or_else(|e| panic!("test_read_string_before failed: {}", e));
}
#[test]
fn test_expect_any() {
|| -> Result<()> {
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\r".to_string()), s),
Err(e) => assert!(false, format!("got error: {}", e)),
}
Ok(())
}()
.unwrap_or_else(|e| panic!("test_expect_any failed: {}", e));
}
#[test]
fn test_expect_empty_command_error() {
let p = spawn("", Some(1000));
match p {
Ok(_) => assert!(false, "should raise an error"),
Err(Error(ErrorKind::EmptyProgramName, _)) => {}
Err(_) => assert!(false, "should raise EmptyProgramName"),
}
}
#[test]
fn test_kill_timeout() {
|| -> Result<()> {
let mut p = spawn_bash(Some(1000))?;
p.execute("cat <(echo ready) -", "ready")?;
Ok(())
}().unwrap_or_else(|e| panic!("test_kill_timeout failed: {}", e));
}
#[test]
fn test_bash() {
|| -> Result<()> {
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(())
}()
.unwrap_or_else(|e| panic!("test_bash failed: {}", e));
}
#[test]
fn test_bash_control_chars() {
|| -> Result<()> {
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(())
}()
.unwrap_or_else(|e| panic!("test_bash_control_chars failed: {}", e));
}
#[test]
fn test_tokenize_command() {
let res = tokenize_command("prog arg1 arg2");
assert_eq!(vec!["prog", "arg1", "arg2"], res);
let res = tokenize_command("prog -k=v");
assert_eq!(vec!["prog", "-k=v"], res);
let res = tokenize_command("prog 'my text'");
assert_eq!(vec!["prog", "'my text'"], res);
let res = tokenize_command(r#"prog "my text""#);
assert_eq!(vec!["prog", r#""my text""#], res);
}
}