device_handler.py

import sys
import json
import subprocess
import time
import os
import re

import database
import logging_setup
from device import Device

RAC_IP = "ras.torizon.io"
logger = logging_setup.setup_logging()


def process_devices(devices, cloud, env_vars, args):
    for index, device in enumerate(devices):
        uuid = device["deviceUuid"]
        hardware_id = cloud.get_hardware_id(uuid)
        logger.debug(f"hardware_id: {hardware_id}")

        if not database.device_exists(uuid):
            database.create_device(uuid)
            logger.info(f"Created new device entry for {uuid}")

        is_last_device = index == len(devices) - 1
        if is_last_device:
            logger.debug(
                f"Wasn't able to lock any of the n - 1 devices, trying again but busy-waiting device number n ({uuid}) until this time."
            )

        # Use fail_fast=False for the last device
        if database.try_until_locked(uuid, fail_fast=not is_last_device):
            logger.info(f"Lock acquired for device {uuid}")
            try:
                dut = Device(cloud, uuid, hardware_id, env_vars)
                dut.create_ssh_connnection()

                if not args.do_not_update:
                    if not dut.is_os_updated_to_latest(
                        env_vars["TARGET_BUILD_TYPE"]
                    ):
                        dut.update_to_latest(
                            env_vars["TARGET_BUILD_TYPE"],
                            args.ignore_different_secondaries_between_updates,
                        )
                        if not dut.is_os_updated_to_latest(
                            env_vars["TARGET_BUILD_TYPE"]
                        ):
                            logger.error(
                                f"Update unsuccessful for {uuid}: trying to get Aktualizr logs and raising an exception. This might take some time."
                            )
                            logger.info(
                                f"Trying an early SSH connection to {dut.remote_session_ip}:{dut.remote_session_port}. If the device didn't roll back successfully this might not work."
                            )

                            # Wait for the device to come back up
                            time.sleep(300)

                            try:
                                dut.connection.run(
                                    "journalctl -u aktualizr-torizon --no-pager"
                                )
                            except ConnectionError as e:
                                logger.error(
                                    f"Failed to connect to device {uuid} for log retrieval: {str(e)}"
                                )

                            raise Exception(f"Update unsuccessful for {uuid}.")

                logger.debug(dut.network_info)

                if dut.network_info:
                    with open("device_information.json", "w") as f:
                        json.dump(dut.network_info, f, ensure_ascii=False)

                if args.run_before_on_host:
                    logger.debug(f"Executing {args.run_before_on_host} on host")

                    if os.name == "nt":
                        subprocess.check_call(
                            [
                                "powershell",
                                "-Command",
                                args.run_before_on_host,
                            ],
                            stdout=sys.stdout,
                            stderr=subprocess.STDOUT,
                        )
                    elif os.name == "posix":
                        subprocess.check_call(
                            args.run_before_on_host,
                            shell=True,
                            stdout=sys.stdout,
                            stderr=subprocess.STDOUT,
                        )
                    else:
                        logger.error(f"Unsupported {os.name} OS")
                        raise Exception(f"Unsupported {os.name} OS")

                if args.hacking_session:
                    logger.info(
                        f"Opening interactive SSH terminal for device {uuid}"
                    )
                    try:
                        dut.connection.shell()
                        logger.info(
                            f"Interactive session finished for device {uuid}"
                        )
                    except Exception as e:
                        if re.search(r"Exit code:\s*1\b", str(e)):
                            logger.info(
                                f"Interactive shell closed normally (exit 1) for device {uuid}"
                            )
                        elif re.search(r"Exit code:\s*130\b", str(e)):
                            logger.info(
                                f"Interactive shell closed using SIGINT (exit 130) for device {uuid}"
                            )
                        else:
                            logger.error(
                                f"Failed to start interactive shell: {e}"
                            )

                if args.before:
                    dut.connection.run(args.before)

                if args.command:
                    dut.connection.run(args.command)
                    logger.info(
                        f"Command '{args.command}' executed for device {uuid} via connection at {dut.remote_session_ip}/{dut.remote_session_port}"
                    )

                if args.copy_artifact:
                    for i in range(0, len(args.copy_artifact), 2):
                        remote_path = args.copy_artifact[i]
                        local_output = args.copy_artifact[i + 1]
                        logger.info(
                            f"Copying artifact from {remote_path} to {local_output}"
                        )
                        dut.connection.get(remote_path, local_output)
                        logger.info(f"Artifact retrieved for device {uuid}")

            except Exception as e:
                logger.error(
                    f"An error occurred while processing device {uuid}: {e}"
                )
                sys.exit(1)
            finally:
                database.release_lock(uuid)
                logger.info(f"Lock released for device {uuid}")

            return True

    return False