Refactor: Chia nho file main thao thu muc va don dep theo yeu cau

core/ssh_flasher.py

@@ -0,0 +1,241 @@
+"""
+SSH-based Firmware Flasher for OpenWrt Devices
+
+Replicates the flash.ps1 logic using Python paramiko/scp:
+  1. Auto-accept SSH host key
+  2. Set device password via `passwd` command
+  3. Upload firmware via SCP to /tmp/
+  4. Verify, sync, and flash via sysupgrade
+"""
+
+import os
+import time
+import paramiko
+from scp import SCPClient
+
+
+def _create_ssh_client(ip, user, password, timeout=10):
+    """Create an SSH client with auto-accept host key policy."""
+    client = paramiko.SSHClient()
+    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    client.connect(ip, username=user, password=password, timeout=timeout,
+                   look_for_keys=False, allow_agent=False)
+    return client
+
+
+import telnetlib
+
+def set_device_password(ip, user="root", old_password="", new_password="admin123a",
+                        status_cb=None):
+    """
+    Set device password via Telnet (if raw/reset) or SSH.
+    """
+    if status_cb:
+        status_cb("Checking Telnet port for raw device...")
+
+    # 1. Thử Telnet trước (OpenWrt mặc định mở Telnet 23 và cấm SSH Root khi chưa có Pass)
+    try:
+        tn = telnetlib.Telnet(ip, timeout=5)
+        # Nếu vô được Telnet tức là thiết bị vừa Reset cứng chưa có pass
+        if status_cb:
+            status_cb("Telnet connected! Setting password...")
+        
+        # Đợi logo OpenWrt và prompt "root@OpenWrt:/# "
+        time.sleep(1)
+        tn.read_very_eager()
+        
+        # Gửi lệnh đổi pass
+        tn.write(b"passwd\n")
+        time.sleep(1)
+        tn.write(new_password.encode('ascii') + b"\n")
+        time.sleep(1)
+        tn.write(new_password.encode('ascii') + b"\n")
+        time.sleep(1)
+        
+        # Thoát telnet
+        tn.write(b"exit\n")
+        time.sleep(0.5)
+        tn.close()
+        
+        # Chờ 3 giây để OpenWrt kịp đóng Telnet và nổ tiến trình Dropbear (SSH Server)
+        if status_cb:
+            status_cb("Password set via Telnet. Waiting for SSH to start...")
+        time.sleep(3)
+        return "DONE"
+        
+    except ConnectionRefusedError:
+        # Port 23 đóng -> Tức là thiết bị đã có Pass và đã bật SSH, chuyển qua luồng mồi mật khẩu cũ
+        pass
+    except Exception as e:
+        # Các lỗi timeout khác có thể do ping không tới
+        return f"FAIL: Telnet check -> {e}"
+
+    # 2. Rơi xuống luồng SSH nếu thiết bị cũ (cổng Telnet đóng)
+    if status_cb:
+        status_cb("Connecting SSH for password update...")
+
+    try:
+        client = _create_ssh_client(ip, user, old_password, timeout=5)
+    except Exception as e:
+        return f"FAIL: Cannot connect SSH (Old pass incorrect?) — {e}"
+
+    try:
+        if status_cb:
+            status_cb("Setting password via SSH...")
+
+        shell = client.invoke_shell()
+        time.sleep(1)
+        if shell.recv_ready():
+            shell.recv(65535)
+
+        shell.send("passwd\n")
+        time.sleep(2)
+        shell.send(f"{new_password}\n")
+        time.sleep(1)
+        shell.send(f"{new_password}\n")
+        time.sleep(2)
+
+        if shell.recv_ready():
+            shell.recv(65535)
+
+        shell.send("exit\n")
+        time.sleep(0.5)
+        shell.close()
+
+        if status_cb:
+            status_cb("Password set ✓")
+        return "DONE"
+
+    except Exception as e:
+        return f"FAIL: {e}"
+    finally:
+        client.close()
+
+
+def flash_device_ssh(ip, firmware_path, user="root", password="admin123a",
+                     set_passwd=False, status_cb=None):
+    """
+    Flash firmware to an OpenWrt device via SSH/SCP.
+
+    Steps (mirroring flash.ps1):
+      1. (Optional) Set device password via passwd
+      2. Upload firmware via SCP to /tmp/
+      3. Verify uploaded file
+      4. Sync filesystem
+      5. Execute sysupgrade
+
+    Returns:
+        "DONE" on success, "FAIL: reason" on error
+    """
+
+    # ═══════════════════════════════════════════
+    #  STEP 0: Set password (optional)
+    # ═══════════════════════════════════════════
+    if set_passwd:
+        result = set_device_password(ip, user, "", password, status_cb)
+        if result.startswith("FAIL"):
+            # Try with current password as old password
+            result = set_device_password(ip, user, password, password, status_cb)
+            if result.startswith("FAIL"):
+                return result
+
+    # ═══════════════════════════════════════════
+    #  STEP 1: Connect SSH
+    # ═══════════════════════════════════════════
+    if status_cb:
+        status_cb("Connecting SSH...")
+
+    try:
+        client = _create_ssh_client(ip, user, password)
+    except paramiko.AuthenticationException:
+        return "FAIL: SSH authentication failed"
+    except paramiko.SSHException as e:
+        return f"FAIL: SSH error — {e}"
+    except Exception as e:
+        return f"FAIL: Cannot connect — {e}"
+
+    try:
+        # ═══════════════════════════════════════════
+        #  STEP 2: Upload firmware via SCP
+        # ═══════════════════════════════════════════
+        if status_cb:
+            status_cb("Uploading firmware via SCP...")
+
+        filename = os.path.basename(firmware_path)
+        remote_path = f"/tmp/{filename}"
+
+        try:
+            scp_client = SCPClient(client.get_transport(), socket_timeout=300)
+            scp_client.put(firmware_path, remote_path)
+            scp_client.close()
+        except Exception as e:
+            return f"FAIL: SCP upload failed — {e}"
+
+        time.sleep(2)
+
+        # ═══════════════════════════════════════════
+        #  STEP 3: Verify firmware uploaded
+        # ═══════════════════════════════════════════
+        if status_cb:
+            status_cb("Verifying firmware...")
+
+        stdin, stdout, stderr = client.exec_command(
+            f"test -f {remote_path} && ls -lh {remote_path}",
+            timeout=10
+        )
+        verify_output = stdout.read().decode("utf-8", errors="ignore").strip()
+        verify_err = stderr.read().decode("utf-8", errors="ignore").strip()
+
+        if not verify_output:
+            return f"FAIL: Firmware file not found on device after upload"
+
+        # ═══════════════════════════════════════════
+        #  STEP 4: Sync filesystem
+        # ═══════════════════════════════════════════
+        if status_cb:
+            status_cb("Syncing filesystem...")
+
+        client.exec_command("sync", timeout=10)
+        time.sleep(2)
+
+        # ═══════════════════════════════════════════
+        #  STEP 5: Flash firmware (sysupgrade)
+        # ═══════════════════════════════════════════
+        if status_cb:
+            status_cb("Flashing firmware (sysupgrade)...")
+
+        try:
+            # Capture output by redirecting to a file in /tmp first, or read from stdout
+            # Use -F to force upgrade and bypass "Image metadata not present" error for uImage files
+            stdin, stdout, stderr = client.exec_command(f"sysupgrade -F -v -n {remote_path} > /tmp/sysup.log 2>&1")
+            
+            # Wait up to 4 seconds to see if it immediately fails
+            # Sysupgrade takes time and normally drops the connection, so if it finishes in < 4s, it's an error.
+            time.sleep(4)
+            
+            if stdout.channel.exit_status_ready():
+                exit_code = stdout.channel.recv_exit_status()
+                
+                # Fetch the error log
+                _, log_out, _ = client.exec_command("cat /tmp/sysup.log")
+                err_msg = log_out.read().decode("utf-8", errors="ignore").strip()
+                
+                return f"FAIL: sysupgrade terminated early (Code {exit_code}). Details:\n{err_msg}"
+
+        except Exception:
+            # Connection drop during sysupgrade is exactly what we expect if it succeeds
+            pass
+
+        if status_cb:
+            status_cb("Rebooting...")
+        time.sleep(3)
+
+        return "DONE"
+
+    except Exception as e:
+        return f"FAIL: {e}"
+    finally:
+        try:
+            client.close()
+        except Exception:
+            pass