breakmancer 0.5.0

//! Controller side of the protocol (run on the developer's laptop).

use std::{
    io::{self, Write},
    net::{IpAddr, Ipv6Addr, SocketAddr, ToSocketAddrs},
    sync::LazyLock,
    time::Duration,
};

use clap::{Parser, ValueEnum};
use indoc::indoc;
use regex::{self, Regex};
use rustyline_async::{Readline, ReadlineEvent};
use tokio::{net::TcpListener, signal};

use crate::{
    protocol::{
        self, AuthVersion, Command, Connection, ExitBreak, OutputType, SecureMsg, TerminateCmd,
        XWingKeypair,
    },
    sigint, util,
};

/// Start a new session.
///
/// When the controller starts, it will give you a command to run on
/// the breakpoint host. The breakpoint end will try to open a connection
/// to the callback address, at which point a secure handshake
/// will occur and a debugging session can start.
#[derive(Parser, Debug)]
pub struct Cli {
    /// Authentication protocol.
    ///
    /// Only one supported value: `dual-pub`
    ///
    /// If not provided, you will be prompted with some questions to
    /// determine whether the available protocol will work.
    #[arg(long, value_enum)]
    auth: Option<AuthVersion>,

    /// Callback address and port
    ///
    /// IP address or domain name that the developer's machine can be reached at,
    /// and the public-facing port to use. This must not be
    /// obstructed by NAT or firewall.
    ///
    /// IPV6 addresses must be enclosed in `[square brackets]`.
    #[arg(required = true, value_name = "HOST:PORT")]
    callback: String,

    /// Local port to listen on for callback.
    ///
    /// If not specified, defaults to the port from the callback
    /// option. Set this value if you use port-forwarding and have
    /// differing external and internal ports.
    #[arg(long)]
    local_port: Option<u16>,
}

struct Params {
    pub callback: String,
    pub listen: SocketAddr,
}

// Should work for listening on both IPv4 and IPv6.
const UNSPECIFIED_IP: IpAddr = IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0));

fn parse_args(args: &Cli) -> Result<Params, String> {
    // We can't really validate the callback address beyond checking
    // if it's well-formed; it's possible that the breakpoint needs to
    // be given a domain name that won't resolve from the controller's
    // vantage point.
    let (_, callback_port) = util::parse_address(&args.callback)
        .map_err(|err| format!("Invalid callback address: {err}"))?;
    // We'll still use it if it doesn't resolve, but at least warn the user.
    match &args.callback.to_socket_addrs() {
        Ok(addrs) => {
            if addrs.len() == 0 {
                eprintln!("[WARNING] Callback host did not resolve to any IP addresses.");
            }
        }
        Err(err) => eprintln!("[WARNING] Could not resolve callback host to an IP address: {err}"),
    }

    let listen_port = args.local_port.unwrap_or(callback_port);

    Ok(Params {
        callback: args.callback.to_string(),
        listen: SocketAddr::new(UNSPECIFIED_IP, listen_port),
    })
}

/// Print bytes as string (lossy Unicode conversion) and add a newline
/// on the end if there wasn't one already.
fn print_with_trailing_newline(data: &[u8]) {
    print!("{}", String::from_utf8_lossy(data));
    if !data.ends_with("\n".as_bytes()) {
        println!();
    }
}

const CMD_EXIT: &str = "exit";
static CMD_EXIT_PATTERN: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(&format!("(?i)\\s*{CMD_EXIT}\\s*")).unwrap());

/// What the user did when we asked them for a command at the REPL.
enum ReplInput {
    /// A command the user wants to have run.
    Command(String),
    /// An interrupt of some sort, possibly signaling a desire to
    /// leave the REPL.
    Interrupted,
}

/// Prompt the user for a command at the REPL.
async fn ask_user_next_command(
    cmd_history: &mut Vec<String>,
    sigint_manager: &mut sigint::Manager,
) -> Result<ReplInput, String> {
    // Set up a new readline each time we need one, because
    // otherwise we can't use regular printlns everywhere. (It has
    // to be dropped in order to get the terminal out of raw mode
    // again.) Maybe
    // change to a persistent readline once we switch to a unified
    // logging system or something. The big downside is having to
    // maintain an external history and re-import it repeatedly.
    let (mut readline, rl_writer) = Readline::new(">> ".to_string())
        .map_err(|err| format!("Could not set up input prompt: {err}"))?;
    for cmd in &mut *cmd_history {
        readline.add_history_entry((*cmd).to_string());
    }

    // Suppress the default ^C handler we have set up.
    let suppress_sigint = sigint_manager.suppress();
    let readline_result = readline.readline().await;
    suppress_sigint.release();

    // Prevent dangling prompt
    readline.flush().unwrap();
    // Keep the writer open this long. Otherwise we get `Couldn't read
    // input: line writers closed` and the session keeps restarting.
    drop(rl_writer);

    match readline_result {
        Ok(ReadlineEvent::Line(command)) => Ok(ReplInput::Command(command)),
        Ok(ReadlineEvent::Eof | ReadlineEvent::Interrupted) => Ok(ReplInput::Interrupted),
        Err(err) => Err(format!("Couldn't read input: {err}")),
    }
}

/// Tell the breakpoint to exit (continue);
///
/// Consumes the connection, as it will stop working immediately anyhow.
async fn hangup(conn: Connection) -> Result<(), String> {
    conn.send(SecureMsg::ExitBreak(ExitBreak {}))
        .await
        .map_err(|err| format!("Could not tell breakpoint to exit: {err}"))?;
    Ok(())
}

/// What the user wants to have happen after the session ends.
enum AfterSession {
    /// Breakpoint should exit, but controller should wait for a new connection.
    Iterate,
    /// Both breakpoint and controller should exit.
    Exit,
}

/// Listen for a connection and perform one shell session.
///
/// The `cmd_history` will be updated with new commands run in this
/// session.
///
/// An Ok return means the session ended normally, without an error
/// (that is, the user asked to end the session). The Ok arm indicates
/// what should happen next.
///
/// The connection is consumed; when the function ends, the connection
/// is dropped, which cancels its tasks.
async fn run_session(
    mut conn: Connection,
    cmd_history: &mut Vec<String>,
    sigint_manager: &mut sigint::Manager,
) -> Result<AfterSession, String> {
    println!("Session ready. You can enter single-line commands. Use `{CMD_EXIT}` to exit.");

    // Each iteration of the loop is a command/response interaction on the REPL.
    let mut cmd_seq = 0;
    loop {
        // Try to get a command to run, or get an interrupt/exit
        // request. Convert these to Some vs. None.
        let command: Option<String> =
            match ask_user_next_command(cmd_history, sigint_manager).await? {
                ReplInput::Command(command) => {
                    if command.is_empty() {
                        continue; // just re-show the prompt
                    }

                    if CMD_EXIT_PATTERN.is_match(&command) {
                        // Another way of exiting the session, besides
                        // ^C/^D. We'll treat it the same way as an
                        // interrupt.
                        None
                    } else {
                        Some(command)
                    }
                }
                ReplInput::Interrupted => None,
            };

        // If there was an interrupt/exit, figure out if they meant it
        // (and return) or whether they actually want to stay in the
        // REPL.
        let Some(command) = command else {
            match ask_user_exit_choice().await {
                ExitChoice::Resume => {
                    // User maybe interrupted by accident, and just
                    // wants to go back to entering commands.
                    continue;
                }
                ExitChoice::Iterate => {
                    hangup(conn).await?;
                    return Ok(AfterSession::Iterate);
                }
                ExitChoice::Exit => {
                    hangup(conn).await?;
                    return Ok(AfterSession::Exit);
                }
            }
        };

        // We actually have a command to run! Save it in the history
        // if it's new and not space-prefixed. (ignorespace,
        // ignoredups)
        if !command.starts_with(' ') && Some(&command) != cmd_history.last() {
            cmd_history.push(command.clone());
        }
        conn.send(SecureMsg::Command(Command {
            command: command.clone(),
            seq: cmd_seq,
        }))
        .await
        .map_err(|err| format!("Could not send command to breakpoint: {err}"))?;

        // Once we send the message, we expect zero or more output
        // lines and then an outcome.

        // Whether we're currently trying to kill the remote command
        // (e.g. because it's hanging).
        let mut killing = false;
        loop {
            tokio::select! {
                // Always check the ^C first, to ensure it is handled
                // on a fast stream.
                biased;

                // We turn ^C into a termination command
                _ = signal::ctrl_c() => {
                    if !killing {
                        println!(); // get onto a new line after the ^C
                        println!("Attempting to terminate command...");
                        conn.send(
                            SecureMsg::TerminateCmd(TerminateCmd { seq: cmd_seq })
                        ).await.map_err(|err| format!("Could not send termination command to breakpoint: {err}"))?;
                        killing = true;
                    }
                },

                // But otherwise just stream the output
                msg = conn.recv() => match msg.ok_or("Error reading from breakpoint.")? {
                    SecureMsg::OutputLine(output) => {
                        if output.seq != cmd_seq {
                            return Err(format!(
                                "Received output for command #{} instead of #{}",
                                output.seq, cmd_seq
                            ));
                        }
                        let log_tag = match output.gender {
                            OutputType::Stdout => "out",
                            OutputType::Stderr => "err",
                        };
                        print!("[{log_tag}] ");
                        print_with_trailing_newline(&output.bytes);
                    }
                    SecureMsg::Finished(outcome) => {
                        if outcome.seq != cmd_seq {
                            return Err(format!(
                                "Received completion notification for command #{} instead of #{}",
                                outcome.seq, cmd_seq
                            ));
                        }
                        match outcome.exit_code {
                            Some(status) => println!("[exit: {status}]"),
                            None => println!("[exit: terminated by signal]"),
                        }
                        break;
                    }
                    wrong_variant => Err(format!(
                        "Received unexpected message type: {}",
                        Into::<&str>::into(wrong_variant)
                    ))?,
                },
            };
        }

        cmd_seq += 1;
    }
}

/// User's confirmation of how they want to exit the REPL.
enum ExitChoice {
    /// They don't actually want to exit, and want to stay in the REPL.
    Resume,
    /// See [`AfterSession::Iterate`].
    Iterate,
    /// See [`AfterSession::Exit`].
    Exit,
}

/// The user has used an interrupt while the REPL was idle -- what do they want to do next?
async fn ask_user_exit_choice() -> ExitChoice {
    println!(indoc! {
        "You've left the command shell. Do you want to resume, or exit?

         [1] Resume: Go back to issuing commands to the breakpoint.
         [2] Iterate: Exit from breakpoint on remote end, running script normally, but leave controller running with same keys.
         [3] Exit: Exit from breakpoint and controller. For additional debugging you'll need to provision the breakpoint with a new key and secret."
    });
    println!();

    let (mut readline, mut rl_writer) = Readline::new("Choice [1/2/3]: ".to_string()).unwrap();
    let choice = loop {
        match readline.readline().await {
            Ok(ReadlineEvent::Line(choice)) => match choice.as_str() {
                "1" => break ExitChoice::Resume,
                "2" => break ExitChoice::Iterate,
                "3" => break ExitChoice::Exit,
                _ => (),
            },
            Ok(ReadlineEvent::Eof) => (),
            Ok(ReadlineEvent::Interrupted) => break ExitChoice::Exit,
            Err(err) => {
                rl_writer
                    .write_all(format!("Failed to read input, exiting: {err}\n").as_bytes())
                    .unwrap();
                break ExitChoice::Exit;
            }
        }
    };
    readline.flush().unwrap(); // prevent dangling prompt
    choice
}

/// Print connection instructions for the user.
fn print_connection_instructions(controller_keypair: &XWingKeypair, params: &Params) {
    let controller_digest = protocol::controller_digest(&controller_keypair.public_key);
    println!(
        indoc! {r#"Run the following command from the environment you want to inspect:

          breakmancer break --auth dual-pub -c {} -i {}

        You can add --which="some location" to distinguish between breakpoints.

        (Now listening on {} for callback...)"#},
        &params.callback, &controller_digest, &params.listen,
    );
}

/// Ask the user if the protocol will work for them.
///
/// Return None if no, Err if can't understand answer.
fn protocol_wizard() -> Result<Option<AuthVersion>, String> {
    println!(
        "Will you be able to view stdout from your process in realtime? \
         This is required for finding the verification string during \
         connection setup."
    );
    print!("Answer [y/n]: ");
    io::stdout().flush().unwrap();
    let mut answer = String::new();
    io::stdin().read_line(&mut answer).unwrap();

    let answer = answer.trim().to_lowercase();
    if answer == "y" || answer == "yes" {
        Ok(Some(AuthVersion::DualPub))
    } else if answer == "n" || answer == "no" {
        Ok(None)
    } else {
        Err("Unknown answer, not yes or no".to_string())
    }
}

pub async fn run(args: &Cli) -> Result<(), String> {
    let params = parse_args(args)?;

    if args.auth.is_none() {
        let chosen = match protocol_wizard()? {
            None => {
                println!(
                    "This program currently cannot meet your needs. \
                 If after reviewing your answers you find they are unchanged, \
                 consider submitting a bug report describing your use-case \
                 and its constraints."
                );
                return Ok(());
            }
            Some(proto) => proto,
        };
        let arg_val = chosen
            .to_possible_value()
            .expect("None of the selectable AuthVersion variants should be skipped.");
        println!(
            "Tip: You can skip this question in the future by passing `--auth {}` \
             (as long as your answers would be the same).",
            arg_val.get_name()
        );
    }

    // Allow override by debug keys
    let (controller_keypair, debug_breakpoint_digest) = match crate::debug_utils::get_debug_keys() {
        Some(keys) => (
            keys.controller_keypair,
            Some(protocol::breakpoint_digest(
                &keys.breakpoint_keypair.public_key,
            )),
        ),
        None => (XWingKeypair::new(), None),
    };

    // We need this signal manager for suppressing ^C during certain
    // operations.
    //
    // Except... that's not quite true. We actually *already* suppress
    // ^C at all times, because rustyline-async and tokio's ctrl_c
    // both register a signal handler, and that removes the default
    // action of ^C. So what this manager *actually* does is register
    // yet another handler that exits on ^C, and then provides a way
    // to *specifically* suppress that one handler for certain
    // operations.
    let mut sigint_manager = sigint::Manager::setup();

    let server_socket = TcpListener::bind(params.listen)
        .await
        .map_err(|err| format!("Could not listen on {}: {}", &params.listen, err))?;

    print_connection_instructions(&controller_keypair, &params);

    // Share the readline (and its history) across sessions.
    let mut command_history = Vec::<String>::new();

    loop {
        let backoff = Duration::from_secs(10);
        // New connection created (and destroyed) on every loop
        // iteration. This minimizes lingering state to reduce the
        // opportunity for vulnerabilities or other bugs (including
        // leftover in-flight messages).
        let conn_partial = match protocol::controller_open_connection(
            controller_keypair.clone(),
            &server_socket,
        )
        .await
        {
            Ok(conn) => conn,
            Err(err) => {
                println!(
                    "Failed to start controller connection; waiting {backoff:?} and trying again: {err}"
                );
                tokio::time::sleep(backoff).await;
                continue;
            }
        };

        let expected_breakpoint_digest = match debug_breakpoint_digest.as_ref() {
            None => {
                print!("Enter the verification string the breakpoint has printed to stdout: ");
                io::stdout().flush().unwrap();
                let mut expected_breakpoint_digest = String::new();
                io::stdin()
                    .read_line(&mut expected_breakpoint_digest)
                    .map_err(|err| format!("Error while reading input: {err}"))?;
                expected_breakpoint_digest
            }
            Some(debug_value) => {
                println!("[WARNING] Using debug breakpoint digest.");
                debug_value.clone()
            }
        };
        let (conn, breakpoint_intro) = match conn_partial
            .verify_breakpoint_and_complete_setup(&expected_breakpoint_digest)
            .await
        {
            Ok(conn_and_info) => conn_and_info,
            Err(err) => {
                println!("Unable to complete connection: {err}");
                continue;
            }
        };

        if let Some(which) = breakpoint_intro.which {
            println!("Breakpoint: {which}");
        }

        match run_session(conn, &mut command_history, &mut sigint_manager).await {
            Ok(AfterSession::Iterate) => {
                println!("Will wait for new connection...");
            }
            Ok(AfterSession::Exit) => {
                return Ok(());
            }
            Err(err) => {
                println!("Encountered error, restarting connection: {err}");
            }
        }
    }
}