Improve router provisioning workflow

src-tauri/src/commands/router.rs

@@ -1,4 +1,10 @@
+use ssh2::Session;
+use std::process::Command;
 use std::{
+    fs::File,
+    io::{ Read, Write },
+    net::{ TcpStream, ToSocketAddrs },
+    path::Path,
     thread,
     time::Duration,
 };
@@ -8,118 +14,703 @@ fn delay() {
 }
 
 #[tauri::command]
-pub async fn detect_router(
-    ip: String,
-) -> Result<bool, String> {
+pub async fn detect_router(ip: String) -> Result<bool, String> {
     delay();
 
     if ip.trim().is_empty() {
-        return Err(
-            "router IP is required".into(),
-        );
+        return Err("router IP is required".into());
     }
 
     Ok(true)
 }
 
-#[tauri::command]
-pub async fn upload_firmware(
-    ip: String,
-    firmware_path: String,
-) -> Result<String, String> {
-    delay();
+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(());
+        }
+    }
 
-    Ok(format!(
-        "uploaded {} to {}:/tmp/firmware.bin",
-        firmware_path, ip,
-    ))
+    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 flash_router(
-    ip: String,
-    remote_firmware_path: String,
-) -> Result<String, String> {
+pub async fn upload_firmware(ip: String, firmware_path: String) -> Result<String, String> {
     delay();
 
-    Ok(format!(
-        "sysupgrade started on {} with {}",
-        ip, remote_firmware_path,
-    ))
+    Ok(format!("uploaded {} to {}:/tmp/firmware.bin", firmware_path, ip))
 }
 
 #[tauri::command]
-pub async fn wait_for_ssh(
-    ip: String,
-) -> Result<bool, String> {
+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"
-    {
+    if ip == "198.51.100.1" || ip == "192.168.1.1" {
         Ok(true)
     } else {
-        Err(format!(
-            "SSH timeout waiting for {}",
-            ip,
-        ))
+        Err(format!("SSH timeout waiting for {}", ip))
     }
 }
 
 #[tauri::command]
-pub async fn upload_provisioning_bundle(
-    ip: String,
-    env_content: String,
-    script_content: String,
-) -> Result<String, String> {
-    delay();
+pub async fn run_provisioning(ip: String, password: String) -> Result<String, String> {
+    if ip.trim().is_empty() {
+        return Err("router IP is required".into());
+    }
 
-    Ok(format!(
-        "uploaded router.env ({} bytes) and provision.sh ({} bytes) to {}",
-        env_content.len(),
-        script_content.len(),
-        ip,
-    ))
+    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 run_provisioning(
-    ip: String,
-) -> Result<String, String> {
-    delay();
+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());
+    }
 
-    Ok(format!(
-        "provision.sh completed on {}; fw4/nftables, wg0, DNAT, LuCI over WireGuard configured",
-        ip,
-    ))
+    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 capture_wireguard_public_key(
-    _ip: String,
-) -> Result<String, String> {
-    delay();
-
-    Ok(
-        "MOCK_ROUTER_WIREGUARD_PUBLIC_KEY_BASE64="
-            .into(),
-    )
-}
-
-#[tauri::command]
-pub async fn verify_router(
-    ip: String,
-) -> Result<bool, String> {
+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(),
-        )
+        Err("router verification failed".into())
     }
-}
+}
+
+#[tauri::command]
+pub async fn upload_firmware_to_router(ip: String, password: String) -> Result<String, String> {
+    let local_firmware_path =
+        "resources/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());
+    }
+
+    /*
+     * FACTORY / FRESH OPENWRT PATH
+     * Uses system scp when no password is set.
+     */
+
+    if password.trim().is_empty() {
+        let target = format!("root@{}:{}", ip, remote_firmware_path);
+
+        let output = std::process::Command
+            ::new("scp")
+            .args([
+                "-O",
+                "-o",
+                "BatchMode=yes",
+                "-o",
+                "ConnectTimeout=10",
+                "-o",
+                "StrictHostKeyChecking=accept-new",
+                local_firmware_path,
+                &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)
+            )
+        );
+    }
+
+    /*
+     * PASSWORD SSH PATH
+     */
+
+    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)
+    })?;
+
+    tcp
+        .set_read_timeout(Some(Duration::from_secs(30)))
+        .map_err(|error| { format!("failed to set read timeout: {}", error) })?;
+
+    tcp
+        .set_write_timeout(Some(Duration::from_secs(30)))
+        .map_err(|error| { format!("failed to set write timeout: {}", error) })?;
+
+    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) })?;
+
+    session
+        .userauth_password("root", &password)
+        .map_err(|error| { format!("SSH authentication failed for root@{}: {}", ip, error) })?;
+
+    if !session.authenticated() {
+        return Err(format!("SSH authentication failed for root@{}", ip));
+    }
+
+    let mut local_file = File::open(local_firmware_path).map_err(|error| {
+        format!("failed to open local firmware file {}: {}", local_firmware_path, 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 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)
+    })?;
+
+    tcp.set_read_timeout(Some(Duration::from_secs(8))).ok();
+    tcp.set_write_timeout(Some(Duration::from_secs(8))).ok();
+
+    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)
+    })?;
+
+    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;
+                }
+            }
+        }
+
+        let address = format!("{}:22", ip);
+
+        let socket_address = match address.to_socket_addrs() {
+            Ok(mut addresses) =>
+                match addresses.next() {
+                    Some(socket_address) => socket_address,
+                    None => {
+                        thread::sleep(wait_between_attempts);
+                        continue;
+                    }
+                }
+            Err(_) => {
+                thread::sleep(wait_between_attempts);
+                continue;
+            }
+        };
+
+        let tcp = match TcpStream::connect_timeout(&socket_address, Duration::from_secs(4)) {
+            Ok(tcp) => tcp,
+            Err(_) => {
+                thread::sleep(wait_between_attempts);
+                continue;
+            }
+        };
+
+        let _ = tcp.set_read_timeout(Some(Duration::from_secs(8)));
+        let _ = tcp.set_write_timeout(Some(Duration::from_secs(8)));
+
+        let mut session = match Session::new() {
+            Ok(session) => session,
+            Err(_) => {
+                thread::sleep(wait_between_attempts);
+                continue;
+            }
+        };
+
+        session.set_tcp_stream(tcp);
+
+        if session.handshake().is_err() {
+            thread::sleep(wait_between_attempts);
+            continue;
+        }
+
+        if authenticate_router(&session, &password).is_err() {
+            thread::sleep(wait_between_attempts);
+            continue;
+        }
+
+        let mut channel = match session.channel_session() {
+            Ok(channel) => channel,
+            Err(_) => {
+                thread::sleep(wait_between_attempts);
+                continue;
+            }
+        };
+
+        if channel.exec("ubus call system board").is_err() {
+            thread::sleep(wait_between_attempts);
+            continue;
+        }
+
+        let mut output = String::new();
+
+        if channel.read_to_string(&mut output).is_err() {
+            thread::sleep(wait_between_attempts);
+            continue;
+        }
+
+        let _ = channel.wait_close();
+
+        return Ok(
+            format!(
+                "Router reconnected after flash on attempt {}/{}.\n{}",
+                attempt,
+                max_attempts,
+                output
+            )
+        );
+    }
+
+    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 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(4)).map_err(|error|
+        format!("SSH not ready at {}: {}", ip, error)
+    )?;
+
+    let _ = tcp.set_read_timeout(Some(Duration::from_secs(8)));
+    let _ = tcp.set_write_timeout(Some(Duration::from_secs(8)));
+
+    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: {}", error))?;
+
+    authenticate_router(&session, &password).map_err(|error| format!("root@{}: {}", ip, error))?;
+
+    let mut channel = session
+        .channel_session()
+        .map_err(|error| format!("failed to open SSH channel: {}", error))?;
+
+    channel
+        .exec("ubus call system board")
+        .map_err(|error| format!("failed to run ubus: {}", error))?;
+
+    let mut output = String::new();
+
+    channel
+        .read_to_string(&mut output)
+        .map_err(|error| format!("failed to read router info: {}", error))?;
+
+    let _ = channel.wait_close();
+
+    Ok(output)
+}
+
+#[tauri::command]
+pub async fn upload_provisioning_bundle(
+    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 = "resources/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 script_upload = Command::new("scp")
+            .args([
+                "-O",
+                "-o",
+                "BatchMode=yes",
+                "-o",
+                "ConnectTimeout=10",
+                "-o",
+                "StrictHostKeyChecking=no",
+                "-o",
+                "UserKnownHostsFile=NUL",
+                local_script_path,
+                &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, 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))
+}

src-tauri/src/lib.rs

@@ -7,11 +7,15 @@ use commands::{
         capture_wireguard_public_key,
         detect_router,
         flash_router,
+        flash_router_sysupgrade,
         run_provisioning,
         upload_firmware,
+        upload_firmware_to_router,
         upload_provisioning_bundle,
+        reconnect_router_after_flash,
         verify_router,
         wait_for_ssh,
+        check_router_after_flash
     },
     ssh::{
         inspect_router_with_password,
@@ -33,12 +37,16 @@ pub fn run() {
             inspect_router_with_password,
             detect_router,
             upload_firmware,
+            upload_firmware_to_router,
             flash_router,
+            flash_router_sysupgrade,
+            reconnect_router_after_flash,
             wait_for_ssh,
             upload_provisioning_bundle,
             run_provisioning,
             capture_wireguard_public_key,
             verify_router,
+            check_router_after_flash
         ])
         .run(tauri::generate_context!())
         .expect("error while running tauri application");