breakmancer 0.1.0

//! Listener/developer side of the protocol.

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

use alkali::asymmetric::kx::x25519blake2b as X25519;
use base64::prelude::{Engine as _, BASE64_STANDARD_NO_PAD};
use clap::Parser;
use indoc::indoc;
use once_cell::sync::Lazy;
use regex::{self, Regex};
use rustyline_async::{Readline, ReadlineEvent};
use tokio::{net::TcpListener, signal::ctrl_c, time::sleep};

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

#[derive(Parser, Debug)]
pub struct Cli {
    /// 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 listener'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: Lazy<Regex> =
    Lazy::new(|| Regex::new(&format!("(?i)\\s*{CMD_EXIT}\\s*")).unwrap());

/// 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).
async fn run_session(
    conn: &mut Connection,
    cmd_history: &mut Vec<String>,
    sigint_manager: &mut sigint::Manager,
) -> Result<(), String> {
    println!("Session ready. You can enter single-line commands. Use `{CMD_EXIT}` to exit.");

    let mut cmd_seq = 0;
    loop {
        // Set up a new readline each time we need one, because
        // otherwise we can't use regular printlns everywhere. 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 line = readline.readline();
        suppress_sigint.release();

        let command = match line.await {
            Ok(ReadlineEvent::Line(command)) => {
                // ignorespace, ignoredups
                if !command.starts_with(' ') && Some(&command) != cmd_history.last() {
                    readline.add_history_entry(command.clone());
                    cmd_history.push(command.clone());
                }
                command
            }
            // User has typed ^D or ^C to exit the session.
            Ok(ReadlineEvent::Eof | ReadlineEvent::Interrupted) => return Ok(()),
            Err(err) => return Err(format!("Couldn't read input: {err}")),
        };
        readline.flush().unwrap(); // prevent dangling prompt
        drop(readline); // get out of raw mode
        drop(rl_writer);

        // Another way of exiting the session, besides ^D.
        if CMD_EXIT_PATTERN.is_match(&command) {
            return Ok(());
        }

        if command.is_empty() {
            continue; // just re-show the prompt
        }

        conn.tx
            .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
                _ = ctrl_c() => {
                    if !killing {
                        println!(); // get onto a new line after the ^C
                        println!("Attempting to terminate command...");
                        conn.tx.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.rx.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;
                    }
                    // TODO: Info about what kind of message was received
                    _ => Err("Received unexpected message type from breakpoint")?,
                },
            };
        }

        cmd_seq += 1;
    }
}

enum AfterSession {
    Resume,
    Iterate,
    Exit,
}

async fn end_session_choose_next() -> AfterSession {
    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 listener running with same keys.
         [3] Exit: Exit from breakpoint and listener. 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 AfterSession::Resume,
                "2" => break AfterSession::Iterate,
                "3" => break AfterSession::Exit,
                _ => continue,
            },
            Ok(ReadlineEvent::Eof) => continue,
            Ok(ReadlineEvent::Interrupted) => break AfterSession::Exit,
            Err(err) => {
                rl_writer
                    .write_all(format!("Failed to read input, exiting: {err}\n").as_bytes())
                    .unwrap();
                break AfterSession::Exit;
            }
        };
    };
    readline.flush().unwrap(); // prevent dangling prompt
    choice
}

/// Print connection instructions for the user.
fn print_connection_instructions(
    oob_secret: &OobSecret,
    listener_pubkey: &X25519::PublicKey,
    params: &Params,
) {
    let oob_secret_b64 = &BASE64_STANDARD_NO_PAD.encode(oob_secret.as_slice());
    let listener_pubkey_b64 = &BASE64_STANDARD_NO_PAD.encode(listener_pubkey.as_slice());
    println!(
        indoc! {r#"On the host where you want to set a breakpoint, set the following environment variable in a secure way:

          BM_CALLBACK_SECRET={}

        (Warning: Anyone who can see the secret can connect to you and pretend to be the breakpoint system, so setting it in an encrypted vault or secret store is recommended. However, they will not be able to impersonate you to the breakpoint system.)

        Then, ensure the breakpoint end runs the following command from the environment you want to inspect:

          breakmancer break -c {} -i {}

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

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

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

    // Allow override by debug keys
    let (listener_keypair, oob_secret) = match crate::debug_utils::get_debug_keys() {
        Some(keys) => (keys.listener_keypair, keys.oob_secret),
        None => (
            X25519::Keypair::generate().unwrap(),
            OobSecret::new_random(),
        ),
    };

    // 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(&oob_secret, &listener_keypair.public_key, &params);

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

    loop {
        let backoff = Duration::from_secs(10);
        let (mut conn, breakpoint_intro) = match Connection::new_listener(
            &server_socket,
            &listener_keypair,
            &oob_secret,
        )
        .await
        {
            Ok(conn) => conn,
            Err(err) => {
                println!(
                    "Failed to start listener connection; waiting {backoff:?} and trying again: {err}"
                );
                sleep(backoff).await;
                continue;
            }
        };
        if let Some(which) = breakpoint_intro.which {
            println!("Breakpoint: {which}");
        }

        match run_session(&mut conn, &mut command_history, &mut sigint_manager).await {
            Ok(()) => match end_session_choose_next().await {
                AfterSession::Resume => {}
                AfterSession::Iterate => {
                    conn.tx
                        .send(SecureMsg::ExitBreak(ExitBreak {}))
                        .await
                        .map_err(|err| format!("Could not tell breakpoint to exit: {err}"))?;
                    println!("Will wait for new connection...");
                }
                AfterSession::Exit => {
                    conn.tx
                        .send(SecureMsg::ExitBreak(ExitBreak {}))
                        .await
                        .map_err(|err| format!("Could not tell breakpoint to exit: {err}"))?;
                    return Ok(());
                }
            },
            Err(err) => {
                println!("Encountered error, restarting connection: {err}");
            }
        }

        // After any session we want to close the connection in order
        // to clear state (and drop in-flight messages).
        conn.hangup();
    }
}