sparse/sparse-05/sparse-05-client/src/commands/connect/shell.rs

use std::{
    io::{self, Read, Write},
    path::PathBuf,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread,
};

use sparse_05_common::messages::{Capabilities, Command, Response};
use structopt::StructOpt;
use tokio::{
    io::{stderr, stdout, AsyncWriteExt},
    runtime::Handle,
    sync::mpsc::{channel, Receiver},
};

use super::{commands, Connection};

#[derive(StructOpt)]
#[structopt()]
pub enum SparseCommands {
    #[structopt(name = "#sysinfo")]
    SysInfo,
    #[structopt(name = "#exit")]
    Exit,
    Cd {
        folder: PathBuf,
    },
    #[structopt(name = "#upload")]
    UploadFile {
        local_file: PathBuf,
        remote_path: PathBuf,
    },
    #[structopt(name = "#download")]
    DownloadFile {
        remote_file: PathBuf,
        local_path: PathBuf,
    },
    #[structopt(name = "#edit")]
    EditFile {
        remote_path: PathBuf,
    },
}

/*macro_rules! libc_try {
    ($expr:expr) => {
        if unsafe { $expr } == -1 {
            return Err(std::io::Error::last_os_error())?;
        }
    };
}

#[cfg(unix)]
fn get_term_attrs<F: AsRawFd>(fd: &F) -> anyhow::Result<libc::termios> {
    let mut termios = unsafe { std::mem::zeroed() };
    libc_try!(libc::tcgetattr(fd.as_raw_fd(), &mut termios));
    Ok(termios)
}

#[cfg(unix)]
fn set_term_attrs<F: AsRawFd>(fd: &mut F, attrs: &libc::termios) -> anyhow::Result<()> {
    libc_try!(libc::tcsetattr(fd.as_raw_fd(), 0, attrs));
    Ok(())
}

#[cfg(unix)]
fn convert_termios_raw(attrs: &mut libc::termios) -> anyhow::Result<()> {
    unsafe { libc::cfmakeraw(attrs) };
    Ok(())
}*/

async fn run_command(
    stdin: &mut Receiver<Vec<u8>>,
    command: String,
    connection: Arc<Connection>,
) -> anyhow::Result<()> {
    connection
        .send_command(Command::RunCommand(command))
        .await?;

    let id = loop {
        let resp = connection.get_response().await?;
        if let Response::AckRunCommand(id) = resp {
            break id;
        }
    };

    loop {
        enum Event {
            Stdin(Vec<u8>),
            Remote(Response),
        }

        let Some(event) = tokio::select! (
            v = connection.get_response() => v.ok().map(Event::Remote),
            v = stdin.recv() => v.map(Event::Stdin)
        ) else {
            continue;
        };

        match event {
            Event::Remote(Response::SendStdout(bytes, cid)) if cid == id => {
                stdout().write(&bytes).await?;
            }
            Event::Remote(Response::SendStderr(bytes, cid)) if cid == id => {
                stderr().write(&bytes).await?;
            }
            Event::Remote(Response::CommandDone(cid, _)) if cid == id => break,
            Event::Stdin(stdin) => {
                let _ = connection
                    .send_command(Command::SendStdin(stdin, id))
                    .await?;
            }
            _ => {}
        }
    }

    Ok(())
}

async fn edit_file(
    cwd: &str,
    connection: Arc<Connection>,
    remote_path: PathBuf,
) -> anyhow::Result<()> {
    let local_path = tempfile::NamedTempFile::new()?.into_temp_path();

    let mut remote_path_used = PathBuf::from(cwd);
    remote_path_used.extend(&remote_path);

    commands::download::download_file(
        Arc::clone(&connection),
        remote_path_used.clone(),
        local_path.to_owned(),
    )
    .await?;

    let cmd = format!(
        "{} {local_path:?}",
        std::env::var("VISUAL")
            .or(std::env::var("EDITOR"))
            .unwrap_or("vim".to_owned())
    );

    tokio::process::Command::new("sh")
        .arg("-c")
        .arg(cmd)
        .spawn()?
        .wait()
        .await?;

    commands::upload::upload_file(
        Arc::clone(&connection),
        local_path.to_owned(),
        remote_path_used,
    )
    .await?;

    Ok(())
}

pub(super) async fn shell(
    connection: Arc<Connection>,
    capabilities: Capabilities,
) -> anyhow::Result<()> {
    println!("Source code available at https://github.com/r-a303931/sparse (feel free to give it a star!)");
    println!("Type #help to view a list of sparse commands\n");

    let mut stdout = io::stdout();
    /* let backup_term_attrs = get_term_attrs(&stdin)?;
    let mut raw_term_attrs = get_term_attrs(&stdin)?;
    convert_termios_raw(&mut raw_term_attrs)?;*/

    let mut cwd = "/".to_string();

    let (stdin_sender, mut stdin_receiver) = channel(64);

    let pause = Arc::new(AtomicBool::new(false));
    let pause_2 = Arc::clone(&pause);

    let handle = Handle::current();
    thread::spawn(move || {
        let mut stdin_buf = [0u8; 1024];
        loop {
            if pause_2.load(Ordering::Relaxed) {
                std::thread::sleep(std::time::Duration::from_millis(250));
                continue;
            }
            let mut stdin = std::io::stdin();
            let Ok(amount) = stdin.read(&mut stdin_buf) else {
                continue;
            };
            if pause_2.load(Ordering::Relaxed) {
                continue;
            }
            let stdin_buf = stdin_buf[..amount].to_vec();
            let stdin_sender = stdin_sender.clone();
            handle.spawn(async move {
                _ = stdin_sender.send(stdin_buf).await;
            });
        }
    });

    'outer: loop {
        print!(
            "{}@{}:{} {} ",
            capabilities.userent.as_deref().unwrap_or("unknown user!"),
            match &capabilities.hostname {
                Some(n) => n.clone(),
                None => format!("{}", connection.ip.ip()),
            },
            cwd,
            if capabilities.root { "#" } else { "$" }
        );
        stdout.flush().unwrap();

        let Some(cmd) = stdin_receiver.recv().await else {
            break;
        };

        if cmd.is_empty() {
            break;
        }

        let Ok(input) = std::str::from_utf8(&cmd) else {
            continue;
        };
        let input = input.trim();

        let (args, help) = if input.starts_with("#help") {
            (input.split(" ").collect::<Vec<_>>(), true)
        } else {
            (
                [&[""][..], &input.split(" ").collect::<Vec<_>>()].concat(),
                false,
            )
        };

        let parsed = SparseCommands::from_iter_safe(args.iter());

        match (parsed, help) {
            (help_info, true) => {
                commands::help::print_help(help_info.ok());
            }
            (Ok(SparseCommands::SysInfo), _) => {
                commands::sysinfo::print_capabilities(&capabilities, &connection.ip.ip())
            }
            (Ok(SparseCommands::Exit), _) => {
                break;
            }
            (
                Ok(SparseCommands::UploadFile {
                    local_file,
                    remote_path,
                }),
                _,
            ) => {
                let mut remote_path_used = PathBuf::from(&cwd);
                remote_path_used.extend(&remote_path);
                if let Err(e) = commands::upload::upload_file(
                    Arc::clone(&connection),
                    local_file,
                    remote_path_used,
                )
                .await
                {
                    eprintln!("{e:?}")
                }
            }
            (
                Ok(SparseCommands::DownloadFile {
                    local_path,
                    remote_file,
                }),
                _,
            ) => {
                let mut remote_path_used = PathBuf::from(&cwd);
                remote_path_used.extend(&remote_file);
                if let Err(e) = commands::download::download_file(
                    Arc::clone(&connection),
                    remote_path_used,
                    local_path,
                )
                .await
                {
                    eprintln!("{e:?}")
                }
            }
            (Ok(SparseCommands::EditFile { remote_path }), _) => {
                pause.store(true, Ordering::SeqCst);
                if let Err(e) = edit_file(&cwd, Arc::clone(&connection), remote_path).await {
                    eprintln!("{e:?}");
                }
                pause.store(false, Ordering::Relaxed);
            }
            (Ok(SparseCommands::Cd { folder }), _) => {
                let Ok(_) = connection.send_command(Command::Cd(folder.clone())).await else {
                    continue;
                };
                let resp = loop {
                    let Ok(resp) = connection.get_response().await else {
                        continue 'outer;
                    };
                    if let Response::CdDone(res) = resp {
                        break res;
                    }
                };
                match resp {
                    Ok(rcwd) => cwd = rcwd.to_string_lossy().to_string(),
                    Err(e) => eprintln!("{e}"),
                }
            }
            _ => {
                if !input.to_string().trim().is_empty() && !pause.load(Ordering::Relaxed) {
                    if let Err(e) = run_command(
                        &mut stdin_receiver,
                        input.to_string(),
                        Arc::clone(&connection),
                    )
                    .await
                    {
                        eprintln!("{e:?}");
                    };
                }
            }
        }
    }

    Ok(())
}