lr-openwrt-tool/src-tauri/src/commands/router.rs

use ssh2::Session;
use std::process::Command;
use std::{
    fs::File,
    io::{Read, Write},
    net::{TcpStream, ToSocketAddrs},
    path::{Path, PathBuf},
    thread,
    time::Duration,
};
use tauri::{AppHandle, Manager};

fn delay() {
    thread::sleep(Duration::from_millis(350));
}

fn resource_path(app: &AppHandle, relative_path: &str) -> Result<PathBuf, String> {
    app.path()
        .resource_dir()
        .map_err(|error| format!("failed to resolve resource directory: {}", error))
        .map(|resource_dir| resource_dir.join("resources").join(relative_path))
}

#[tauri::command]
pub async fn detect_router(ip: String) -> Result<bool, String> {
    delay();

    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    Ok(true)
}

fn authenticate_router(session: &Session, password: &str) -> Result<(), String> {
    if !password.trim().is_empty() {
        if session.userauth_password("root", password).is_ok() && session.authenticated() {
            return Ok(());
        }
    }

    if session.userauth_agent("root").is_ok() && session.authenticated() {
        return Ok(());
    }

    Err("SSH authentication failed for root".into())
}

fn run_system_ssh(ip: &str, command: &str) -> Result<String, String> {
    let target = format!("root@{}", ip);

    let output = Command::new("ssh")
        .args([
            "-o",
            "BatchMode=yes",
            "-o",
            "ConnectTimeout=6",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=NUL",
            &target,
            command,
        ])
        .output()
        .map_err(|error| format!("failed to run system ssh: {}", error))?;

    let stdout = String::from_utf8_lossy(&output.stdout);
    let stderr = String::from_utf8_lossy(&output.stderr);

    if output.status.success() {
        return Ok(format!("{}{}", stdout, stderr));
    }

    Err(format!(
        "system ssh failed with status {:?}:\n{}\n{}",
        output.status.code(),
        stderr,
        stdout
    ))
}

fn open_router_session(ip: &str, password: &str) -> Result<Session, String> {
    let address = format!("{}:22", ip);

    let socket_address = address
        .to_socket_addrs()
        .map_err(|error| format!("failed to resolve router address: {}", error))?
        .next()
        .ok_or_else(|| format!("failed to resolve router address {}", address))?;

    let tcp = TcpStream::connect_timeout(&socket_address, Duration::from_secs(8))
        .map_err(|error| format!("failed to connect to SSH on {}: {}", ip, error))?;

    let _ = tcp.set_read_timeout(Some(Duration::from_secs(30)));
    let _ = tcp.set_write_timeout(Some(Duration::from_secs(30)));

    let mut session =
        Session::new().map_err(|error| format!("failed to create SSH session: {}", error))?;

    session.set_tcp_stream(tcp);

    session
        .handshake()
        .map_err(|error| format!("SSH handshake failed for {}: {}", ip, error))?;

    authenticate_router(&session, password).map_err(|error| format!("root@{}: {}", ip, error))?;

    Ok(session)
}

fn scp_file_from_disk(
    session: &Session,
    local_path: &str,
    remote_path: &str,
    mode: i32,
) -> Result<(), String> {
    let mut local_file =
        File::open(local_path).map_err(|error| format!("failed to open {}: {}", local_path, error))?;

    let file_size = local_file
        .metadata()
        .map_err(|error| format!("failed to read metadata for {}: {}", local_path, error))?
        .len();

    let mut buffer = Vec::new();

    local_file
        .read_to_end(&mut buffer)
        .map_err(|error| format!("failed to read {}: {}", local_path, error))?;

    let mut remote_file = session
        .scp_send(Path::new(remote_path), mode, file_size, None)
        .map_err(|error| format!("failed to start SCP upload to {}: {}", remote_path, error))?;

    remote_file
        .write_all(&buffer)
        .map_err(|error| format!("failed to write {}: {}", remote_path, error))?;

    remote_file.send_eof().ok();
    remote_file.wait_eof().ok();
    remote_file.close().ok();
    remote_file.wait_close().ok();

    Ok(())
}

fn scp_string(session: &Session, content: &str, remote_path: &str, mode: i32) -> Result<(), String> {
    let bytes = content.as_bytes();

    let mut remote_file = session
        .scp_send(Path::new(remote_path), mode, bytes.len() as u64, None)
        .map_err(|error| format!("failed to start SCP upload to {}: {}", remote_path, error))?;

    remote_file
        .write_all(bytes)
        .map_err(|error| format!("failed to write {}: {}", remote_path, error))?;

    remote_file.send_eof().ok();
    remote_file.wait_eof().ok();
    remote_file.close().ok();
    remote_file.wait_close().ok();

    Ok(())
}

fn run_ssh_command(session: &Session, command: &str) -> Result<String, String> {
    let mut channel = session
        .channel_session()
        .map_err(|error| format!("failed to open SSH channel: {}", error))?;

    channel
        .exec(command)
        .map_err(|error| format!("failed to run command `{}`: {}", command, error))?;

    let mut output = String::new();

    let _ = channel.read_to_string(&mut output);
    let _ = channel.wait_close();

    Ok(output)
}

#[tauri::command]
pub async fn upload_firmware(ip: String, firmware_path: String) -> Result<String, String> {
    delay();

    Ok(format!("uploaded {} to {}:/tmp/firmware.bin", firmware_path, ip))
}

#[tauri::command]
pub async fn flash_router(ip: String, remote_firmware_path: String) -> Result<String, String> {
    delay();

    Ok(format!("sysupgrade started on {} with {}", ip, remote_firmware_path))
}

#[tauri::command]
pub async fn wait_for_ssh(ip: String) -> Result<bool, String> {
    for _ in 0..3 {
        delay();
    }

    if ip == "198.51.100.1" || ip == "192.168.1.1" {
        Ok(true)
    } else {
        Err(format!("SSH timeout waiting for {}", ip))
    }
}

#[tauri::command]
pub async fn run_provisioning(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let command = "cd /tmp && ./provision.sh";

    if password.trim().is_empty() {
        return run_system_ssh(&ip, command);
    }

    let session = open_router_session(&ip, &password)?;

    run_ssh_command(&session, command)
}

#[tauri::command]
pub async fn capture_wireguard_public_key(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let command = "wg show wg0 public-key";

    let output = if password.trim().is_empty() {
        run_system_ssh(&ip, command)?
    } else {
        let session = open_router_session(&ip, &password)?;
        run_ssh_command(&session, command)?
    };

    let public_key = output.trim();

    if public_key.is_empty() {
        return Err("WireGuard public key was empty".into());
    }

    if !public_key.ends_with('=') || public_key.len() < 40 {
        return Err(format!("unexpected WireGuard public key output: {}", public_key));
    }

    Ok(public_key.to_string())
}

#[tauri::command]
pub async fn verify_router(ip: String) -> Result<bool, String> {
    delay();

    if ip == "198.51.100.1" {
        Ok(true)
    } else {
        Err("router verification failed".into())
    }
}

#[tauri::command]
pub async fn upload_firmware_to_router(
    app: AppHandle,
    ip: String,
    password: String,
) -> Result<String, String> {
    let local_firmware_path = resource_path(
        &app,
        "firmware/openwrt-23.05.5-zbt-we826-16m-litoral-golden-sysupgrade.bin",
    )?;

    let remote_firmware_path = "/tmp/firmware.bin";

    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    if password.trim().is_empty() {
        let target = format!("root@{}:{}", ip, remote_firmware_path);
        let local_firmware_path_string = local_firmware_path.to_string_lossy().to_string();

        let output = Command::new("scp")
            .args([
                "-O",
                "-o",
                "BatchMode=yes",
                "-o",
                "ConnectTimeout=10",
                "-o",
                "StrictHostKeyChecking=accept-new",
                &local_firmware_path_string,
                &target,
            ])
            .output()
            .map_err(|error| format!("failed to run system scp: {}", error))?;

        if output.status.success() {
            return Ok(format!(
                "Firmware uploaded to {}:{} using system scp",
                ip, remote_firmware_path
            ));
        }

        return Err(format!(
            "system scp failed:\n{}\n{}",
            String::from_utf8_lossy(&output.stderr),
            String::from_utf8_lossy(&output.stdout)
        ));
    }

    let session = open_router_session(&ip, &password)?;

    let mut local_file = File::open(&local_firmware_path).map_err(|error| {
        format!(
            "failed to open local firmware file {}: {}",
            local_firmware_path.display(),
            error
        )
    })?;

    let firmware_size = local_file
        .metadata()
        .map_err(|error| format!("failed to read firmware metadata: {}", error))?
        .len();

    let mut buffer = Vec::new();

    local_file
        .read_to_end(&mut buffer)
        .map_err(|error| format!("failed to read firmware file: {}", error))?;

    let mut remote_file = session
        .scp_send(Path::new(remote_firmware_path), 0o644, firmware_size, None)
        .map_err(|error| {
            format!(
                "failed to start SCP upload to {}: {}",
                remote_firmware_path, error
            )
        })?;

    remote_file
        .write_all(&buffer)
        .map_err(|error| format!("failed to upload firmware via SCP: {}", error))?;

    remote_file
        .send_eof()
        .map_err(|error| format!("failed to send SCP EOF: {}", error))?;

    remote_file
        .wait_eof()
        .map_err(|error| format!("failed waiting for SCP EOF: {}", error))?;

    remote_file
        .close()
        .map_err(|error| format!("failed closing SCP channel: {}", error))?;

    remote_file
        .wait_close()
        .map_err(|error| format!("failed waiting for SCP close: {}", error))?;

    Ok(format!(
        "Firmware uploaded to {}:{} ({} bytes)",
        ip, remote_firmware_path, firmware_size
    ))
}

#[tauri::command]
pub async fn flash_router_sysupgrade(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let flash_command = "test -f /tmp/firmware.bin && sysupgrade -n /tmp/firmware.bin";

    if password.trim().is_empty() {
        let output = run_system_ssh(&ip, flash_command);

        return match output {
            Ok(output) => Ok(format!(
                "sysupgrade command submitted on {} using system ssh. {}",
                ip, output
            )),
            Err(error) => {
                let lowered = error.to_lowercase();

                if lowered.contains("commencing upgrade")
                    || lowered.contains("closing all shell sessions")
                    || lowered.contains("connection failed")
                    || lowered.contains("connection reset")
                    || lowered.contains("broken pipe")
                    || lowered.contains("closed")
                    || lowered.contains("disconnect")
                    || lowered.contains("sysupgrade")
                {
                    Ok(format!(
                        "sysupgrade started on {}; SSH disconnected as expected.",
                        ip
                    ))
                } else {
                    Err(error)
                }
            }
        };
    }

    let session = open_router_session(&ip, &password)?;

    let mut channel = session
        .channel_session()
        .map_err(|error| format!("failed to open SSH channel: {}", error))?;

    channel
        .exec(flash_command)
        .map_err(|error| format!("failed to start sysupgrade: {}", error))?;

    let mut output = String::new();
    let _ = channel.read_to_string(&mut output);
    let _ = channel.wait_close();

    let lowered = output.to_lowercase();

    if lowered.contains("commencing upgrade")
        || lowered.contains("closing all shell sessions")
        || lowered.contains("sysupgrade")
    {
        return Ok(format!(
            "sysupgrade started on {}; router should reboot shortly. {}",
            ip, output
        ));
    }

    Ok(format!(
        "sysupgrade command submitted on {}. Router should reboot shortly. {}",
        ip, output
    ))
}

#[tauri::command]
pub async fn reconnect_router_after_flash(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let max_attempts = 30;
    let wait_between_attempts = Duration::from_secs(5);

    for attempt in 1..=max_attempts {
        if password.trim().is_empty() {
            match run_system_ssh(&ip, "ubus call system board") {
                Ok(output) => {
                    return Ok(format!(
                        "Router reconnected after flash on attempt {}/{} using system ssh.\n{}",
                        attempt, max_attempts, output
                    ));
                }
                Err(_) => {
                    thread::sleep(wait_between_attempts);
                    continue;
                }
            }
        }

        match open_router_session(&ip, &password) {
            Ok(session) => match run_ssh_command(&session, "ubus call system board") {
                Ok(output) => {
                    return Ok(format!(
                        "Router reconnected after flash on attempt {}/{}.\n{}",
                        attempt, max_attempts, output
                    ));
                }
                Err(_) => thread::sleep(wait_between_attempts),
            },
            Err(_) => thread::sleep(wait_between_attempts),
        }
    }

    Err(format!(
        "Router did not become reachable over SSH at {} after {} attempts. Replug Ethernet and retry.",
        ip, max_attempts
    ))
}

#[tauri::command]
pub async fn check_router_after_flash(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    if password.trim().is_empty() {
        return run_system_ssh(&ip, "ubus call system board");
    }

    let session = open_router_session(&ip, &password)?;
    run_ssh_command(&session, "ubus call system board")
}

#[tauri::command]
pub async fn upload_provisioning_bundle(
    app: AppHandle,
    ip: String,
    password: String,
    env_content: String,
) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let local_script_path = resource_path(&app, "provisioning/provision.sh")?;
    let remote_script_path = "/tmp/provision.sh";
    let remote_env_path = "/tmp/router.env";

    if password.trim().is_empty() {
        let script_target = format!("root@{}:{}", ip, remote_script_path);
        let local_script_path_string = local_script_path.to_string_lossy().to_string();

        let script_upload = Command::new("scp")
            .args([
                "-O",
                "-o",
                "BatchMode=yes",
                "-o",
                "ConnectTimeout=10",
                "-o",
                "StrictHostKeyChecking=no",
                "-o",
                "UserKnownHostsFile=NUL",
                &local_script_path_string,
                &script_target,
            ])
            .output()
            .map_err(|error| format!("failed to run scp for provision.sh: {}", error))?;

        if !script_upload.status.success() {
            return Err(format!(
                "failed to upload provision.sh:\n{}\n{}",
                String::from_utf8_lossy(&script_upload.stderr),
                String::from_utf8_lossy(&script_upload.stdout)
            ));
        }

        let env_command = format!(
            "cat > {} <<'EOF'\n{}\nEOF\nchmod +x {}",
            remote_env_path, env_content, remote_script_path
        );

        run_system_ssh(&ip, &env_command)?;

        return Ok(format!(
            "uploaded provision.sh and router.env to {} using system ssh/scp",
            ip
        ));
    }

    let session = open_router_session(&ip, &password)?;

    scp_file_from_disk(
        &session,
        local_script_path.to_string_lossy().as_ref(),
        remote_script_path,
        0o755,
    )?;

    scp_string(&session, &env_content, remote_env_path, 0o600)?;

    run_ssh_command(&session, "chmod +x /tmp/provision.sh")?;

    Ok(format!("uploaded provision.sh and router.env to {}", ip))
}

#[tauri::command]
pub async fn upload_udp2raw_setup_script(
    app: AppHandle,
    ip: String,
    password: String,
) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let local_script_path = resource_path(&app, "udp2raw/setup_udp2raw.sh")?;
    let remote_script_path = "/tmp/setup_udp2raw.sh";

    if password.trim().is_empty() {
        let target = format!("root@{}:{}", ip, remote_script_path);
        let local_script_path_string = local_script_path.to_string_lossy().to_string();

        let output = Command::new("scp")
            .args([
                "-O",
                "-o",
                "BatchMode=yes",
                "-o",
                "ConnectTimeout=10",
                "-o",
                "StrictHostKeyChecking=no",
                "-o",
                "UserKnownHostsFile=NUL",
                &local_script_path_string,
                &target,
            ])
            .output()
            .map_err(|error| format!("failed to run scp for setup_udp2raw.sh: {}", error))?;

        if !output.status.success() {
            return Err(format!(
                "failed to upload setup_udp2raw.sh:\n{}\n{}",
                String::from_utf8_lossy(&output.stderr),
                String::from_utf8_lossy(&output.stdout)
            ));
        }

        run_system_ssh(&ip, "chmod +x /tmp/setup_udp2raw.sh")?;

        return Ok(format!("uploaded setup_udp2raw.sh to {}", ip));
    }

    let session = open_router_session(&ip, &password)?;

    scp_file_from_disk(
        &session,
        local_script_path.to_string_lossy().as_ref(),
        remote_script_path,
        0o755,
    )?;

    run_ssh_command(&session, "chmod +x /tmp/setup_udp2raw.sh")?;

    Ok(format!("uploaded setup_udp2raw.sh to {}", ip))
}

#[tauri::command]
pub async fn run_udp2raw_setup(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let command = "sh /tmp/setup_udp2raw.sh";

    if password.trim().is_empty() {
        return run_system_ssh(&ip, command);
    }

    let session = open_router_session(&ip, &password)?;

    run_ssh_command(&session, command)
}

#[tauri::command]
pub async fn check_udp2raw_router_status(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let command = r#"
    echo "== udp2raw binary =="
    if command -v udp2raw >/dev/null 2>&1; then
    command -v udp2raw
    else
    echo "missing"
    fi

    echo ""
    echo "== init script =="
    if [ -x /etc/init.d/udp2raw-wg ]; then
    echo "present"
    else
    echo "missing"
    fi

    echo ""
    echo "== process =="
    if pgrep -af "^/usr/bin/udp2raw" >/dev/null 2>&1; then
    pgrep -af "^/usr/bin/udp2raw"
    else
    echo "not running"
    fi

    echo ""
    echo "== service status =="
    if [ -x /etc/init.d/udp2raw-wg ]; then
    /etc/init.d/udp2raw-wg status || true
    else
    echo "service unavailable"
    fi

    echo ""
    echo "== WireGuard configured endpoint =="
    uci get network.wgserver.endpoint_host 2>/dev/null || true
    uci get network.wgserver.endpoint_port 2>/dev/null || true

    echo ""
    echo "== local listener =="
    netstat -ln 2>/dev/null | grep -E '127.0.0.1:4999|:4999' || echo "listener not confirmed"

    echo ""
    echo "== WireGuard runtime endpoint =="
    wg show wg0 2>/dev/null | grep -A8 '^peer:' || echo "wg0 unavailable"
    "#;

    if password.trim().is_empty() {
        return run_system_ssh(&ip, command);
    }

    let session = open_router_session(&ip, &password)?;

    run_ssh_command(&session, command)
}

#[tauri::command]
pub async fn test_udp2raw_tunnel(ip: String, password: String) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let command = r#"
    echo "== udp2raw process =="
    if pgrep -af "^/usr/bin/udp2raw" >/dev/null 2>&1; then
    pgrep -af "^/usr/bin/udp2raw"
    else
    echo "ERROR: udp2raw is not running"
    exit 20
    fi

    echo ""
    echo "== WireGuard configured endpoint =="
    uci get network.wgserver.endpoint_host 2>/dev/null || true
    uci get network.wgserver.endpoint_port 2>/dev/null || true

    echo ""
    echo "== local listener =="
    netstat -ln 2>/dev/null | grep -E '127.0.0.1:4999|:4999' || echo "WARNING: listener not confirmed"

    echo ""
    echo "== ping VPS public IP =="
    ping -c 2 -W 2 146.59.230.190 || true

    echo ""
    echo "== WireGuard status =="
    wg show wg0 2>/dev/null || echo "wg0 not available"

    echo ""
    echo "== route check =="
    ip route || true

    echo ""
    echo "UDP2RAW tunnel test completed"
    "#;

    if password.trim().is_empty() {
        return run_system_ssh(&ip, command);
    }

    let mut last_error = String::new();

    for attempt in 1..=5 {
        match open_router_session(&ip, &password) {
            Ok(session) => {
                return run_ssh_command(&session, command);
            }
            Err(error) => {
                last_error = format!("SSH attempt {}/5 failed: {}", attempt, error);
                thread::sleep(Duration::from_secs(2));
            }
        }
    }

    Err(last_error)
}

#[tauri::command]
pub async fn upload_udp2raw_binary(
    app: AppHandle,
    ip: String,
    password: String,
) -> Result<String, String> {
    if ip.trim().is_empty() {
        return Err("router IP is required".into());
    }

    let local_binary_path = resource_path(&app, "udp2raw/udp2raw")?;
    let remote_binary_path = "/usr/bin/udp2raw";

    if password.trim().is_empty() {
        let target = format!("root@{}:{}", ip, remote_binary_path);
        let local_binary_path_string = local_binary_path.to_string_lossy().to_string();

        let output = Command::new("scp")
            .args([
                "-O",
                "-o",
                "BatchMode=yes",
                "-o",
                "ConnectTimeout=10",
                "-o",
                "StrictHostKeyChecking=no",
                "-o",
                "UserKnownHostsFile=NUL",
                &local_binary_path_string,
                &target,
            ])
            .output()
            .map_err(|error| format!("failed to run scp for udp2raw binary: {}", error))?;

        if !output.status.success() {
            return Err(format!(
                "failed to upload udp2raw binary:\n{}\n{}",
                String::from_utf8_lossy(&output.stderr),
                String::from_utf8_lossy(&output.stdout)
            ));
        }

        run_system_ssh(
            &ip,
            "chmod +x /usr/bin/udp2raw && /usr/bin/udp2raw --help >/dev/null 2>&1 || true",
        )?;

        return Ok("uploaded udp2raw binary to /usr/bin/udp2raw".into());
    }

    let session = open_router_session(&ip, &password)?;

    scp_file_from_disk(
        &session,
        local_binary_path.to_string_lossy().as_ref(),
        remote_binary_path,
        0o755,
    )?;

    run_ssh_command(
        &session,
        "chmod +x /usr/bin/udp2raw && ls -l /usr/bin/udp2raw",
    )?;

    Ok("uploaded udp2raw binary to /usr/bin/udp2raw".into())
}