deltachat-remotecontrol-bot/src/handler.rs

use std::{
    collections::HashSet,
    fs::File,
    io::Read,
    path::PathBuf,
    sync::Arc,
    time::Duration,
};

use anyhow::{Context as _, Result as AnyhowResult};
use deltachat::{
    chat::{self, ChatId},
    contact::ContactId,
    context::Context,
    message::Message,
};
use eui48::MacAddress;
use russh::{
    client::AuthResult,
    keys::PrivateKeyWithHashAlg,
};
use tokio::sync::Mutex;

use crate::{AUTH_REQUIRED, config::BotConfig, data_path};

pub struct BotContext {
    authed_contacts: HashSet<ContactId>,
    config: BotConfig,
}

impl BotContext {
    pub fn new(config: BotConfig) -> BotContext {
        BotContext {
            authed_contacts: HashSet::new(),
            config,
        }
    }
}

pub async fn echo(dchat_ctx: Arc<Mutex<Context>>, msg: Message) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let dchat_ctx_lock = dchat_ctx.lock().await;
    chat::send_text_msg(&dchat_ctx_lock, chat_id, msg.get_text()).await?;
    Ok(())
}

pub async fn auth_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let mut ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !AUTH_REQUIRED {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            "Authentication is disabled".to_owned(),
        )
        .await?;
        return Ok(());
    }

    if ctx_lock.authed_contacts.contains(&contact_id) {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Already authenticated".to_owned()).await?;
        return Ok(());
    }

    let Some(password) = args.first() else {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            "Usage: /auth <password>".to_owned(),
        )
        .await?;
        return Ok(());
    };

    if *password == ctx_lock.config.auth.password {
        ctx_lock.authed_contacts.insert(contact_id);
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            "Authentication successful!".to_owned(),
        )
        .await?;
    } else {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Incorrect password".to_owned()).await?;
    }

    Ok(())
}

pub async fn wol_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !ensure_auth(&dchat_ctx_lock, &ctx_lock, chat_id, contact_id).await? {
        return Ok(());
    }

    let mac = match args.first() {
        None => match ctx_lock.config.machines.values().find(|m| m.is_default) {
            Some(machine) => machine.mac,
            None => {
                let usage = wol_usage(&ctx_lock.config);
                chat::send_text_msg(&dchat_ctx_lock, chat_id, usage).await?;
                return Ok(());
            }
        },
        Some(arg) => {
            if let Some(machine) = ctx_lock.config.machines.get(*arg) {
                machine.mac
            } else if let Ok(mac) = MacAddress::parse_str(arg) {
                mac
            } else {
                let usage = wol_usage(&ctx_lock.config);
                chat::send_text_msg(&dchat_ctx_lock, chat_id, usage).await?;
                return Ok(());
            }
        }
    };

    let reply = match send_magic_packet(mac).await {
        Ok(()) => format!("Magic packet sent to {}.", mac_to_string(mac)),
        Err(e) => format!("Failed to send magic packet: {e:#}"),
    };
    chat::send_text_msg(&dchat_ctx_lock, chat_id, reply).await?;

    Ok(())
}

fn wol_usage(config: &BotConfig) -> String {
    let mut forms: Vec<String> = Vec::new();

    if config.machines.values().any(|m| m.is_default) {
        forms.push("/wol".to_owned());
    }

    let mut names: Vec<&str> = config.machines.keys().map(String::as_str).collect();
    names.sort();
    if !names.is_empty() {
        forms.push(format!("/wol <{}>", names.join("|")));
    }

    forms.push("/wol <XX:XX:XX:XX:XX:XX>".to_owned());

    format!("Usage: {}", forms.join(" OR "))
}

fn mac_to_string(mac: MacAddress) -> String {
    let b = mac.as_bytes();
    format!(
        "{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}",
        b[0], b[1], b[2], b[3], b[4], b[5]
    )
}

async fn send_magic_packet(mac: MacAddress) -> AnyhowResult<()> {
    let mut packet = [0u8; 102];
    packet[..6].fill(0xFF);
    let mac_bytes = mac.as_bytes();
    for i in 0..16 {
        packet[6 + i * 6..6 + (i + 1) * 6].copy_from_slice(mac_bytes);
    }

    let socket = tokio::net::UdpSocket::bind("0.0.0.0:0")
        .await
        .context("Failed to bind UDP socket for WoL")?;
    socket
        .set_broadcast(true)
        .context("Failed to enable broadcast on UDP socket")?;
    socket
        .send_to(&packet, "255.255.255.255:9")
        .await
        .context("Failed to send magic packet")?;

    Ok(())
}

fn ssh_unlock_disk_usage(config: &BotConfig) -> String {
    if !config.machines.is_empty() {
        let list = config
            .machines
            .keys()
            .cloned()
            .collect::<Vec<_>>()
            .join("|");
        if config.default_machine().is_some() {
            format!("Usage: /ssh-unlock-disk [{list}] <password>")
        } else {
            format!("Usage: /ssh-unlock-disk {list} <password>")
        }
    } else {
        "Command /ssh-unlock-disk unavailable: no machines configured".to_string()
    }
}

pub async fn ssh_unlock_disk_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !ensure_auth(&dchat_ctx_lock, &ctx_lock, chat_id, contact_id).await? {
        return Ok(());
    }

    if args.is_empty() || args.len() > 2 {
        log::warn!("wrong number of args");
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            ssh_unlock_disk_usage(&ctx_lock.config),
        )
        .await?;
        return Ok(());
    }

    let Some((machine_name, machine)) = ({
        if args.len() > 1 {
            ctx_lock.config.machines.get(args[0]).map(|m| (args[0], m))
        } else if let Some((name, machine)) = ctx_lock.config.default_machine() {
            log::info!("default machine: {:?} {:?}", name, machine);
            Some((name, machine))
        } else {
            log::warn!("no machine found :/");
            None
        }
    }) else {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            ssh_unlock_disk_usage(&ctx_lock.config),
        )
        .await?;
        return Ok(());
    };

    let Some(remote_unlock_cfg) = machine.remote_unlock.as_ref() else {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            format!("Remote unlock is not configured for machine '{machine_name}"),
        )
        .await?;
        return Ok(());
    };

    let ssh_key_relative_path: PathBuf = remote_unlock_cfg.ssh_key_file.clone().into();
    let ssh_key_path = data_path().join(ssh_key_relative_path);
    log::debug!("Reading SSH key from {:?}", &ssh_key_path);
    let mut ssh_key_file = File::open(ssh_key_path)?;
    let mut ssh_key = String::new();
    ssh_key_file.read_to_string(&mut ssh_key)?;
    drop(ssh_key_file);

    let password = args[args.len() - 1].to_owned();
    let host = machine.static_ip.to_string();
    let ssh_key =
        russh::keys::PrivateKey::from_openssh(ssh_key).context("Invalid remote unlock SSH key")?;

    drop(ctx_lock);
    drop(dchat_ctx_lock);

    tokio::spawn(async move {
        let ssh_config = Arc::new(russh::client::Config {
            inactivity_timeout: None,
            ..Default::default()
        });
        // TODO implement real host key verification (needs config file changes)
        const SSH_TIMEOUT: Duration = Duration::from_secs(15);
        log::info!("Connecting to SSH {}:22", &host);
        let mut connection = tokio::time::timeout(
            SSH_TIMEOUT,
            russh::client::connect(ssh_config, (host, 22), crate::ssh::ClientHandler {}),
        )
        .await
        .context("Connection to host timed out")?
        .context("Cannot connect to host")?;

        log::info!("Trying to authenticate as root");
        let auth_result = tokio::time::timeout(
            SSH_TIMEOUT,
            connection.authenticate_publickey(
                "root",
                PrivateKeyWithHashAlg::new(
                    Arc::new(ssh_key),
                    connection.best_supported_rsa_hash().await?.flatten(),
                ),
            ),
        )
        .await
        .context("SSH authentication timed out")?
        .context("SSH auth against remote unlock target failed")?;

        // ignoring both fields because we only have the key
        if let AuthResult::Failure {
            remaining_methods: _,
            partial_success: _,
        } = auth_result
        {
            anyhow::bail!("SSH authentication failed");
        }

        log::info!("Opening session");
        let mut channel = tokio::time::timeout(SSH_TIMEOUT, connection.channel_open_session())
            .await
            .context("Timed out")?
            .context("Cannot open SSH session")?;

        log::info!("Requesting shell");
        tokio::time::timeout(SSH_TIMEOUT, channel.request_shell(true))
            .await
            .context("Timed out")?
            .context("Cannot request disk password prompt")?;

        log::info!("Sending password");
        let input = password + "\n";
        tokio::time::timeout(SSH_TIMEOUT, channel.data(input.as_bytes()))
            .await
            .context("Timed out")?
            .context("Cannot send password via SSH")?;

        loop {
            let Some(msg) = tokio::time::timeout(SSH_TIMEOUT, channel.wait())
                .await
                .context("Timed out")?
            else {
                break;
            };
            log::info!("Received message from SSH server: {:?}", msg);
        }

        Ok(())
    })
    .await?
    .context("SSH unlock error")?;

    let dchat_ctx_lock = dchat_ctx.lock().await;
    chat::send_text_msg(&dchat_ctx_lock, chat_id, "Unlock password sent".to_owned()).await?;

    Ok(())
}

async fn ensure_auth(
    dchat_ctx: &Context,
    ctx: &BotContext,
    chat_id: ChatId,
    contact_id: ContactId,
) -> AnyhowResult<bool> {
    if !AUTH_REQUIRED {
        return Ok(true);
    }
    if !ctx.authed_contacts.contains(&contact_id) {
        chat::send_text_msg(
            dchat_ctx,
            chat_id,
            "Authenticate yourself first with '/auth <password>'.".to_owned(),
        )
        .await?;
        Ok(false)
    } else {
        Ok(true)
    }
}