From f4f8a148d7f1bad3f36eb90504ca6514229d56ef Mon Sep 17 00:00:00 2001
From: remittor <remittor@gmail.com>
Date: Wed, 23 Apr 2025 10:39:06 +0300
Subject: [PATCH] gateway: Fix SSH-error after run "reboot" command

---
 gateway.py    | 6 +++++-
 install_fw.py | 8 ++------
 reboot.py     | 4 +++-
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/gateway.py b/gateway.py
index 09a1407..8fe8d21 100644
--- a/gateway.py
+++ b/gateway.py
@@ -901,7 +901,7 @@ class Gateway():
           return False
     return True
 
-  def run_cmd(self, command, msg = None, timeout = None, die_on_error = True):
+  def run_cmd(self, command, msg = None, timeout = None, die_on_error = True, reboot = False):
     error = 0
     reslist = [ ]
     self.errcode = -1
@@ -938,6 +938,10 @@ class Gateway():
             except Exception as e:
                 error = -10
             finally:
+                if reboot:
+                    self.shutdown()
+                    self.errcode = 0
+                    return ''
                 channel.close()
                 channel.wait_closed()
             if error != 0 and die_on_error:
diff --git a/install_fw.py b/install_fw.py
index c3226f3..dca269b 100644
--- a/install_fw.py
+++ b/install_fw.py
@@ -1116,12 +1116,8 @@ class XqFlash():
         else:
             import ssh2
             print('Send command "reboot" via SSH/Telnet ...')
-            try:
-                gw.run_cmd("reboot -f", die_on_error = False)
-                print("Forced REBOOT activated!")
-            except ssh2.exceptions.SocketRecvError as e:
-                print("Forced REBOOT Activated!")
-                pass 
+            gw.run_cmd("reboot -f", die_on_error = False, reboot = True)
+            print("Forced REBOOT activated!")
 
     def flash_data_to_mtd(self, img_name, img: Image, timeout, check = True):
         print(f'Writing {img_name} image to addr 0x{img.addr:08X} ...')
diff --git a/reboot.py b/reboot.py
index aa4002e..1ee8b93 100644
--- a/reboot.py
+++ b/reboot.py
@@ -3,6 +3,7 @@
 
 import os
 import sys
+import time
 import platform
 
 import xmir_base
@@ -15,7 +16,8 @@ gw = gateway.Gateway(detect_ssh = False)
 ssh = gw.detect_ssh(verbose = 1, interactive = True)
 if ssh > 0:
   print('Send command "reboot" via SSH ...')
-  gw.run_cmd("reboot")
+  gw.run_cmd("reboot", reboot = True)
+  time.sleep(1)
 else:
   if not gw.stok:
     gw.web_login()