sima-cli 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

sima_cli/update/netboot.py

@@ -0,0 +1,408 @@
+from sima_cli.update.updater import download_image
+from sima_cli.utils.net import get_local_ip_candidates
+from sima_cli.update.remote import get_remote_board_info, copy_file_to_remote_board, DEFAULT_PASSWORD, run_remote_command, init_ssh_session
+import os
+import platform
+import threading
+import socket
+import select
+import time
+import logging
+import click
+from errno import EINTR
+from tftpy import TftpServer, TftpException, TftpTimeout, TftpTimeoutExpectACK, DEF_TFTP_PORT, DEF_TIMEOUT_RETRIES
+from tftpy.TftpContexts import TftpContextServer
+from tftpy.TftpPacketFactory import TftpPacketFactory
+
+# Configuration constants
+MAX_BLKSIZE = 1468  # Block size for MTU compatibility
+SOCK_TIMEOUT = 2    # Timeout for faster retransmits
+
+log = logging.getLogger("tftpy.InteractiveTftpServer")
+emmc_image_paths = []
+
+def flash_emmc(client_manager, emmc_image_paths):
+    """Flash eMMC on a selected client device."""
+    clients = [
+        (ip, info) for ip, info in client_manager.get_client_info()
+        if info.get("state") == "Connected"
+    ]
+
+    if not clients:
+        click.echo("📭 No connected clients available to flash.")
+        return
+
+    if len(clients) == 1:
+        selected_ip = clients[0][0]
+    else:
+        click.echo("👥 Multiple connected clients found. Select one to flash:")
+        for idx, (ip, info) in enumerate(clients, 1):
+            board_info = info.get("board_info") or "Unknown"
+            click.echo(f"   {idx}. {ip} - {board_info}")
+        while True:
+            choice = input("Enter the number of the client to flash: ").strip()
+            if choice.isdigit() and 1 <= int(choice) <= len(clients):
+                selected_ip = clients[int(choice) - 1][0]
+                break
+            click.echo("❌ Invalid choice. Try again.")
+
+    click.echo(f"📡 Selected client: {selected_ip}")
+    remote_dir = "/tmp"
+
+    for path in emmc_image_paths:
+        click.echo(f"📤 Copying {path} to {selected_ip}:{remote_dir}")
+        success = copy_file_to_remote_board(selected_ip, path, remote_dir, passwd=DEFAULT_PASSWORD)
+        if not success:
+            click.echo(f"❌ Failed to copy {path} to {selected_ip}. Aborting.")
+            return
+
+    try:
+        ssh = init_ssh_session(selected_ip, password=DEFAULT_PASSWORD)
+
+        # Step a: Check if eMMC exists
+        check_cmd = "[ -e /dev/mmcblk0p1 ] || (echo '❌ /dev/mmcblk0p1 not found'; exit 1)"
+        run_remote_command(ssh, check_cmd)
+
+        # Step b: Flash with bmaptool
+        wic_path = next((p for p in emmc_image_paths if p.endswith(".wic.gz")), None)
+        if not wic_path:
+            click.echo("❌ No .wic.gz image found in emmc_image_paths.")
+            return
+
+        # Step c: umount the emmc
+        pre_unmount_cmd = (
+            "sudo mount | grep mmcblk0 | awk '{print $3}' | while read mnt; do "
+            "sudo umount \"$mnt\"; done"
+        )
+        run_remote_command(ssh, pre_unmount_cmd)
+
+        # Step d: bmaptool copy the image
+        filename = os.path.basename(wic_path)
+        remote_path = f"/tmp/{filename}"        
+        flash_cmd = f"sudo bmaptool copy {remote_path} /dev/mmcblk0"
+        run_remote_command(ssh, flash_cmd)
+
+        # Step d: Fix GPT
+        fix_cmd = 'sudo printf "fix\n" | sudo parted ---pretend-input-tty /dev/mmcblk0 print'
+        run_remote_command(ssh, fix_cmd)
+
+        click.echo("✅ Flash completed. Please reboot the board to boot from eMMC.")
+    except Exception as e:
+        click.echo(f"❌ Flashing failed: {e}")
+
+class ClientManager:
+    """Manages TFTP client state and monitoring."""
+    def __init__(self):
+        self.clients = {}
+        self.lock = threading.Lock()
+        self.shutdown_event = threading.Event()
+
+    def add_client(self, ip, filename):
+        """Add a new client with initial state."""
+        with self.lock:
+            if ip not in self.clients:
+                start_time = time.time()
+                self.clients[ip] = {
+                    'state': 'Booting',
+                    'filename': filename,
+                    'timestamp': start_time,
+                    'board_info': None
+                }
+                click.echo(f"📥 New client connected: {ip}")
+                if filename:
+                    click.echo(f"📄 Client {ip} requested file: {filename}")
+                # Start monitoring thread
+                threading.Thread(
+                    target=self.monitor_client,
+                    args=(ip, start_time),
+                    daemon=True
+                ).start()
+
+    def monitor_client(self, ip, start_time):
+        """Monitor client by calling get_remote_board_info after 1 minutes, retry every 10 seconds."""
+        try:
+            # Wait ~1 minute before first attempt
+            if self.shutdown_event.wait(timeout=max(0, 60 - (time.time() - start_time))):
+                return
+
+            while not self.shutdown_event.is_set():
+                try:
+                    board_type, build_version = get_remote_board_info(ip)
+
+                    if board_type and build_version:
+                        with self.lock:
+                            self.clients[ip]['state'] = 'Connected'
+                            self.clients[ip]['board_info'] = f"{board_type}|{build_version}"
+                        click.echo(f"✅ Board info retrieved for {ip}: {board_type} with {build_version}")
+                        break 
+                    else:
+                        log.debug(f"Incomplete board info for {ip}: {board_type}, {build_version}")
+
+                except Exception as e:
+                    log.info(f"get_remote_board_info failed for {ip}, retrying in 10s: {e}")
+
+                if self.shutdown_event.wait(timeout=10):
+                    break 
+
+        except Exception as e:
+            log.error(f"Unexpected error while monitoring {ip}: {e}")
+
+    def get_client_info(self):
+        """Return sorted client information for display."""
+        with self.lock:
+            return sorted(self.clients.items(), key=lambda x: x[0])
+
+    def shutdown(self):
+        """Signal monitoring threads to exit."""
+        self.shutdown_event.set()
+
+class InteractiveTftpServer(TftpServer):
+    """Custom TFTP server with client logging and monitoring."""
+    def __init__(self, tftproot, client_manager):
+        super().__init__(tftproot)
+        self.client_manager = client_manager
+
+    def listen(self, listenip="", listenport=DEF_TFTP_PORT, timeout=SOCK_TIMEOUT, retries=DEF_TIMEOUT_RETRIES):
+        """Override listen to log client IPs and filenames."""
+        tftp_factory = TftpPacketFactory()
+        if not listenip:
+            listenip = "0.0.0.0"
+        log.info(f"Server requested on ip {listenip}, port {listenport}")
+        try:
+            self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            self.sock.bind((listenip, listenport))
+            self.sock.setblocking(0)
+            _, self.listenport = self.sock.getsockname()
+        except OSError as err:
+            raise err
+
+        self.is_running.set()
+        log.info("Starting receive loop...")
+        while True:
+            log.debug("shutdown_immediately is %s" % self.shutdown_immediately)
+            log.debug("shutdown_gracefully is %s" % self.shutdown_gracefully)
+            if self.shutdown_immediately:
+                log.info("Shutting down now. Session count: %d" % len(self.sessions))
+                self.sock.close()
+                for key in self.sessions:
+                    log.warning("Forcefully closed session with %s" % self.sessions[key].host)
+                    self.sessions[key].end()
+                self.sessions = []
+                self.is_running.clear()
+                self.shutdown_gracefully = self.shutdown_immediately = False
+                self.client_manager.shutdown()
+                break
+            elif self.shutdown_gracefully:
+                if not self.sessions:
+                    log.info("In graceful shutdown mode and all sessions complete.")
+                    self.sock.close()
+                    self.is_running.clear()
+                    self.shutdown_gracefully = self.shutdown_immediately = False
+                    self.client_manager.shutdown()
+                    break
+
+            inputlist = [self.sock]
+            for key in self.sessions:
+                inputlist.append(self.sessions[key].sock)
+
+            try:
+                readyinput, _, _ = select.select(inputlist, [], [], timeout)
+            except OSError as err:
+                if err.errno == EINTR:
+                    log.debug("Interrupted syscall, retrying")
+                    continue
+                else:
+                    raise
+
+            deletion_list = []
+            for readysock in readyinput:
+                if readysock == self.sock:
+                    log.debug("Data ready on our main socket")
+                    buffer, (raddress, rport) = self.sock.recvfrom(MAX_BLKSIZE)
+                    log.debug("Read %d bytes", len(buffer))
+
+                    if self.shutdown_gracefully:
+                        log.warning("Discarding data on main port, in graceful shutdown mode")
+                        continue
+
+                    key = f"{raddress}:{rport}"
+                    if key not in self.sessions:
+                        log.debug("Creating new server context for session key = %s" % key)
+                        filename = None
+                        if buffer[:2] == b'\x00\x01':  # RRQ packet
+                            filename = buffer[2:].split(b'\x00')[0].decode()
+                        self.client_manager.add_client(raddress, filename)
+                        self.sessions[key] = TftpContextServer(
+                            raddress,
+                            rport,
+                            timeout,
+                            self.root,
+                            self.dyn_file_func,
+                            self.upload_open,
+                            retries=retries
+                        )
+                        try:
+                            self.sessions[key].start(buffer)
+                        except TftpTimeoutExpectACK:
+                            self.sessions[key].timeout_expectACK = True
+                        except TftpException as err:
+                            deletion_list.append(key)
+                            log.error("Fatal exception thrown from session %s: %s" % (key, str(err)))
+                    else:
+                        log.warning("received traffic on main socket for existing session??")
+                    log.info("Currently handling these sessions:")
+                    for session_key, session in list(self.sessions.items()):
+                        log.info("    %s" % session)
+                else:
+                    for key in self.sessions:
+                        if readysock == self.sessions[key].sock:
+                            log.debug("Matched input to session key %s" % key)
+                            self.sessions[key].timeout_expectACK = False
+                            try:
+                                self.sessions[key].cycle()
+                                if self.sessions[key].state is None:
+                                    log.info("Successful transfer.")
+                                    deletion_list.append(key)
+                            except TftpTimeoutExpectACK:
+                                self.sessions[key].timeout_expectACK = True
+                            except TftpException as err:
+                                deletion_list.append(key)
+                                log.error("Fatal exception thrown from session %s: %s" % (key, str(err)))
+                            break
+                    else:
+                        log.error("Can't find the owner for this packet. Discarding.")
+
+            now = time.time()
+            for key in self.sessions:
+                try:
+                    self.sessions[key].checkTimeout(now)
+                except TftpTimeout as err:
+                    log.error(str(err))
+                    self.sessions[key].retry_count += 1
+                    if self.sessions[key].retry_count >= self.sessions[key].retries:
+                        log.debug("hit max retries on %s, giving up" % self.sessions[key])
+                        deletion_list.append(key)
+                    else:
+                        log.debug("resending on session %s" % self.sessions[key])
+                        self.sessions[key].state.resendLast()
+
+            for key in deletion_list:
+                log.info("Session %s complete" % key)
+                if key in self.sessions:
+                    log.debug("Gathering up metrics from session before deleting")
+                    self.sessions[key].end()
+                    metrics = self.sessions[key].metrics
+                    if metrics.duration == 0:
+                        log.info("Duration too short, rate undetermined")
+                    else:
+                        log.info("Transferred %d bytes in %.2f seconds" % (metrics.bytes, metrics.duration))
+                        log.info("Average rate: %.2f kbps" % metrics.kbps)
+                        click.echo(f"✅ Transfer to {self.sessions[key].host} complete: "
+                                   f"{metrics.bytes} bytes in {metrics.duration:.2f} s ({metrics.kbps:.2f} kbps)")
+                    log.info("%.2f bytes in resent data" % metrics.resent_bytes)
+                    log.info("%d duplicate packets" % metrics.dupcount)
+                    log.debug("Deleting session %s" % key)
+                    del self.sessions[key]
+                    log.debug("Session list is now %s" % self.sessions)
+                else:
+                    log.warning("Strange, session %s is not on the deletion list" % key)
+
+        self.is_running.clear()
+        self.shutdown_gracefully = self.shutdown_immediately = False
+        self.client_manager.shutdown()
+
+def run_cli(client_manager):
+    """Run the interactive CLI for netboot commands."""
+    click.echo("\n🛠  Type 'c' to see connected IPs and board info, 'f' to flash eMMC, or 'q' to quit.\n")
+    while True:
+        try:
+            user_input = input("netboot> ").strip().lower()
+            if user_input in {"q", "quit", "exit"}:
+                click.echo("🛑 Shutting down TFTP server.")
+                return True
+            elif user_input == "c":
+                client_info = client_manager.get_client_info()
+                if client_info:
+                    click.echo("🧾 TFTP client IPs and status:")
+                    for ip, info in client_info:
+                        state = info['state']
+                        filename = info['filename'] or "Unknown"
+                        timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(info['timestamp']))
+                        board_info = info['board_info']
+                        click.echo(f"   • {ip}: {state}, Initial File: {filename}, First seen: {timestamp}")
+                        if board_info:
+                            click.echo(f"     Board Info: {board_info}")
+                else:
+                    click.echo("📭 No TFTP client requests received yet.")
+            elif user_input == "f":
+                click.echo("🔧 Initiating eMMC flash (implementation pending).")
+                flash_emmc(client_manager, emmc_image_paths)
+            elif user_input == "":
+                continue
+            else:
+                click.echo("❓ Unknown command. Try 'c' to print client list, 'f' to flash emmc, or 'q'.")
+        except (KeyboardInterrupt, EOFError):
+            click.echo("\n🛑 Exiting netboot session.")
+            return True
+
+def setup_netboot(version: str, board: str, internal: bool = False, autoflash: bool = False):
+    """
+    Download and serve a bootable image for network boot over TFTP with client monitoring.
+
+    Parameters:
+        version (str): Firmware version to download (e.g., "1.6.0").
+        board (str): Target board type, e.g., "modalix" or "davinci".
+        internal (bool): Whether to use internal download sources. Defaults to False.
+        autoflash (bool): Whether to automatically flash the devkit when networked booted. Defaults to False.
+
+    Raises:
+        RuntimeError: If the download or TFTP setup fails.
+    """
+    global emmc_image_paths
+
+    if platform.system() == "Windows":
+        click.echo("❌ Netboot with built-in TFTP is not supported on Windows. Use macOS or Linux.")
+        exit(1)
+
+    try:
+        click.echo(f"⬇️ Downloading netboot image for version: {version}, board: {board}")
+        file_list = download_image(version, board, swtype="yocto", internal=internal, update_type='netboot')
+        if not isinstance(file_list, list):
+            raise ValueError("Expected list of extracted files, got something else.")
+        extract_dir = os.path.dirname(file_list[0])
+        click.echo(f"📁 Image extracted to: {extract_dir}")
+        # Extract specific image paths
+        wic_gz_file = next((f for f in file_list if f.endswith(".wic.gz")), None)
+        bmap_file = next((f for f in file_list if f.endswith(".wic.bmap")), None)
+        emmc_image_paths = [p for p in [wic_gz_file, bmap_file] if p]
+        click.echo(f"📁 eMMC image paths are: {emmc_image_paths}")
+
+    except Exception as e:
+        raise RuntimeError(f"❌ Failed to download and extract netboot image: {e}")
+
+    try:
+        click.echo(f"🚀 Starting TFTP server in: {extract_dir}")
+        ip_candidates = get_local_ip_candidates()
+        if not ip_candidates:
+            click.echo("❌ No suitable local IP addresses found.")
+            exit(1)
+
+        click.echo("🌐 TFTP server is listening on these interfaces (UDP port 69):")
+        for iface, ip in ip_candidates:
+            click.echo(f"   🔹 {iface}: {ip}")
+
+        client_manager = ClientManager()
+
+        server = InteractiveTftpServer(tftproot=extract_dir, client_manager=client_manager)
+        server_thread = threading.Thread(target=server.listen, args=('0.0.0.0', 69), daemon=True)
+        server_thread.start()
+
+        if run_cli(client_manager):
+            server.stop(now=True)
+            client_manager.shutdown()
+
+    except PermissionError:
+        raise RuntimeError("❌ Permission denied. You must run this command with sudo to bind to port 69.")
+    except OSError as e:
+        raise RuntimeError(f"❌ Failed to start TFTP server: {e}")

sima_cli/update/remote.py

@@ -83,7 +83,8 @@ def get_remote_board_info(ip: str, passwd: str = DEFAULT_PASSWORD) -> Tuple[str,
 
         return board_type, build_version
 
-    except Exception:
+    except Exception as e:
+        click.echo(f"Unable to retrieve board info with error: {e}, board may be still booting.")
         return "", ""
 
 
@@ -94,7 +95,7 @@ def _scp_file(sftp, local_path: str, remote_path: str):
     sftp.put(local_path, remote_path)
     click.echo("✅ Upload complete")
 
-def _run_remote_command(ssh, command: str, password: str = DEFAULT_PASSWORD):
+def run_remote_command(ssh, command: str, password: str = DEFAULT_PASSWORD):
     """
     Run a remote command over SSH and stream its output live to the console.
     If the command starts with 'sudo', pipe in the password.
@@ -139,6 +140,13 @@ def _run_remote_command(ssh, command: str, password: str = DEFAULT_PASSWORD):
     for line in remaining_err.splitlines():
         click.echo(f"⚠️ {line}")
 
+def init_ssh_session(ip: str, password: str = DEFAULT_PASSWORD):
+    ssh = paramiko.SSHClient()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+    ssh.connect(ip, username=DEFAULT_USER, password=password, timeout=10)
+    return ssh
+
 def reboot_remote_board(ip: str, passwd: str):
     """
     Reboot remote board by sending SSH command
@@ -149,12 +157,36 @@ def reboot_remote_board(ip: str, passwd: str):
 
         ssh.connect(ip, username=DEFAULT_USER, password=passwd, timeout=10)
 
-        _run_remote_command(ssh, "sudo systemctl stop watchdog", password=passwd)
-        _run_remote_command(ssh, "sudo bash -c 'echo b > /proc/sysrq-trigger'", password=passwd)
+        run_remote_command(ssh, "sudo systemctl stop watchdog", password=passwd)
+        run_remote_command(ssh, "sudo bash -c 'echo b > /proc/sysrq-trigger'", password=passwd)
 
     except Exception as reboot_err:
         click.echo(f"⚠️  Unable to connect to the remote board")
 
+
+def copy_file_to_remote_board(ip: str, file_path: str, remote_dir: str, passwd: str):
+    """
+    Copy a file to the remote board over SSH.
+    Assumes default credentials: sima / edgeai.
+    """
+    ssh = paramiko.SSHClient()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+    try:
+        ssh.connect(ip, username=DEFAULT_USER, password=passwd, timeout=10)
+        sftp = ssh.open_sftp()
+
+        # Upload the file
+        base_file_path = os.path.basename(file_path)
+        click.echo(f"📤 Uploading {file_path} → {remote_dir}")
+        sftp.put(file_path, os.path.join(remote_dir, base_file_path))
+        return True
+
+    except Exception as e:
+        click.echo(f"❌ Remote file copy failed: {e}")
+
+    return False
+
 def push_and_update_remote_board(ip: str, troot_path: str, palette_path: str, passwd: str, reboot_and_wait: bool):
     """
     Upload and install firmware images to remote board over SSH.
@@ -176,7 +208,7 @@ def push_and_update_remote_board(ip: str, troot_path: str, palette_path: str, pa
         click.echo("🚀 Uploaded tRoot image.")
 
         # Run tRoot update
-        _run_remote_command(
+        run_remote_command(
             ssh,
             f"sudo swupdate -H simaai-image-troot:1.0 -i /tmp/{troot_name}", password=passwd
         )
@@ -191,7 +223,7 @@ def push_and_update_remote_board(ip: str, troot_path: str, palette_path: str, pa
         click.echo("🚀 Uploaded system image.")
 
         # Run Palette update
-        _run_remote_command(
+        run_remote_command(
             ssh,
             f"sudo swupdate -H simaai-image-palette:1.0 -i /tmp/{palette_name}",
             password=passwd
@@ -205,11 +237,12 @@ def push_and_update_remote_board(ip: str, troot_path: str, palette_path: str, pa
             click.echo("🔁 Rebooting board after update. Waiting for reconnection...")
 
             try:
-                _run_remote_command(ssh, "sudo systemctl stop watchdog", password=passwd)
-                _run_remote_command(ssh, "sudo bash -c 'echo b > /proc/sysrq-trigger'", password=passwd)
+                run_remote_command(ssh, "sudo systemctl stop watchdog", password=passwd)
+                run_remote_command(ssh, "sudo bash -c 'echo b > /proc/sysrq-trigger'", password=passwd)
 
             except Exception as reboot_err:
                 click.echo(f"⚠️  SSH connection lost due to reboot (expected): {reboot_err}, please powercycle the board...")
+                click.confirm("⚠️  Have you powercycled the board?", default=True, abort=True)
 
             try:
                 ssh.close()
@@ -228,7 +261,7 @@ def push_and_update_remote_board(ip: str, troot_path: str, palette_path: str, pa
                 ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                 ssh.connect(ip, username=DEFAULT_USER, password=passwd, timeout=10)
 
-                _run_remote_command(ssh, "grep SIMA_BUILD_VERSION /etc/build 2>/dev/null || grep SIMA_BUILD_VERSION /etc/buildinfo 2>/dev/null", password=passwd)
+                run_remote_command(ssh, "grep SIMA_BUILD_VERSION /etc/build 2>/dev/null || grep SIMA_BUILD_VERSION /etc/buildinfo 2>/dev/null", password=passwd)
                 ssh.close()
             except Exception as e:
                 click.echo(f"❌ Unable to validate the version: {e}")