#!/usr/bin/env bash

# This script is Bootstrapped with bash to open a screen session, to keep it running on SSH session exit

""""echo
echo Automated EVTEC Image Update Script
echo ===================================
sleep 1

SESSION_NAME="doImageUpdate"

SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
SCRIPT_FILE="${SCRIPT_DIR}/$(basename "$0")"

# Check if the screen session exists
screen -ls | grep -q "\.$SESSION_NAME\s"

if [ $? -eq 0 ]; then
    # The screen session exists, so attach to it
    screen -rx $SESSION_NAME
    exit
fi
# The screen session doesn't exist, so create a new one

scriptArgs="--noScreen $@"

# Setup logging of the session
timestamp=`date +%Y-%m-%d_%H-%M-%S`
logfile="${SCRIPT_DIR}/imageUpdate_${timestamp}.log"

screen -L -Logfile ${logfile} -m -S $SESSION_NAME bash -c "
    echo 'Screen session opened (log: ${logfile}), launching python...'
    python3 ${SCRIPT_FILE} ${scriptArgs}
    read -t 0.001 < /dev/null
    echo Logfile: ${logfile}
    echo 'Press any key to exit...'
    read -s -n 1
"

exit
# The `exit` causes the bash interpreter to stop and not continue into python code below.
"""


# TODO: more retry mechanism
# TODO: get bundle from neighbor charger

import argparse
import asyncio
import contextlib
import fcntl
import functools
import hashlib
import json
import os
import platform
import re
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from getpass import getpass, getuser
from json import JSONDecodeError
from pathlib import Path
from socket import AI_NUMERICHOST, AF_INET
from typing import List, Dict, Any, Optional, BinaryIO, Tuple, Callable, Union
from zipfile import ZipFile

import aiohttp
from dbus_next import BusType
from dbus_next.aio import MessageBus
from file_read_backwards import FileReadBackwards

try:
    # Primarily for better highlighting in the IDE
    import asyncssh
except ImportError:
    pass

SCREEN_NAME = "doImageUpdate"

DEFAULT_UPDATE_SERVER_URL_TEMPLATE = (
    "https://{}software.evtec.ch/images/productive/{}/evtec-image-update-bundle-evtec-msi-icore-mx6dl.raucb"
)

RAUC_LOGIN_ENV = "RAUC_LOGIN"
SUDO_PASS_ENV = "SUDO_PASS"

ECP_DATA_DIR = "/mnt/persistent-data/evtec-ecp-data"
CONFIG_FILE_LOCATION = "config/config.json"
LOG_DIR = "log"
WORKDIR_DEFAULT = "/mnt/persistent-data/rauc-work"
HW_WD_GPIO_PATH = Path("/sys/class/gpio/gpio0/value")

DOWNLOAD_KB_TIMEOUT_S = 30
HASH_TIMEOUT_S = 60 * 3
PROGRESS_LOG_INTERVAL_S = 10
INSTALL_TIMEOUT_MIN = 10
INSTALL_TIMEOUT_DISPLAY_S = 15 * 60
T_WAIT_BEFORE_SHUTDOWN = 180

WATCHDOG_DEVICE = "/dev/watchdog"
WDIOC_KEEPALIVE = 0x80045705
MAGIC_CLOSE_CHARACTER = b"V"

FULL_DOC = f"""
Provide required update server credentials through the environment variable RAUC_LOGIN
and sudo password with SUDO_PASS, example:

    export RAUC_LOGIN="user:passwd"
    export SUDO_PASS="secret123"

The script can be called from an arbitrary path, but update files will be stored in:
    "{WORKDIR_DEFAULT}"

Note: per default the external watchdog/reset IC is armed and will force reset the board on any unexpected behaviour.

Example:
    SUDO_PASS=<PASS> RAUC_LOGIN=<USER>:<PASS> ./doImageUpdate.py --bundleHash=57795a4ee5c6adf353ad6af3b2655355 -d 1.0.22

"""


class Rauc:
    """Interfaces with the system rauc deamon via d-bus"""

    def __init__(self):
        self._bus = None
        self._installer = None
        self._progressTask: Optional[asyncio.Task] = None
        self._installCompleted = False
        self._completedEvent = asyncio.Event()

    async def connect(self) -> None:
        """Connect to d-bus"""
        self._bus = await MessageBus(bus_type=BusType.SYSTEM).connect()

        introspection = await self._bus.introspect("de.pengutronix.rauc", "/")
        proxy = self._bus.get_proxy_object("de.pengutronix.rauc", "/", introspection)
        self._installer = proxy.get_interface("de.pengutronix.rauc.Installer")

        # attach signal handler
        self._installer.on_completed(self.handelCompleted)  # noqa

    async def getSlotStatus(self) -> Dict[str, Dict[str, Any]]:
        """Get the status of all slots

        Example:
        {
        "rootfs.0": {
            "installed.timestamp": "2023-06-20T14:10:03Z",
            "size": 357381084,
            "sha256": "c69f8b6cca06fa39614904d7ea93d240cd5ca7145a7c3e2d09c29ec6afec41aa",
            "bundle.build": "20230620062026",
            "bundle.compatible": "Evtec Rauc Update",
            "installed.count": 3,
            "bundle.version": "1.0",
            "type": "ext4",
            "bootname": "A",
            "class": "rootfs",
            "activated.count": 3,
            "bundle.description": "evtec-image-update-bundle version 1.0-r0",
            "boot-status": "good",
            "activated.timestamp": "2023-06-20T14:10:03Z",
            "mountpoint": "/",
            "state": "booted",
            "status": "ok",
            "device": "/dev/mmcblk0p2"
        }

        """
        result: List = await self._installer.call_get_slot_status()
        # transform [["slotName", {"paramName": Variant, ...}] ...] to plain dict
        parsed = dict(result)
        for slotName, slotData in parsed.items():
            parsed[slotName] = {paramName: param.value for paramName, param in slotData.items()}

        return parsed

    async def getPrimarySlot(self) -> str:
        """Get the currently selected boot slot

        :return: e.g. 'rootfs.1'
        """
        result: str = await self._installer.call_get_primary()
        return result

    async def installBundle(self, bundleFile: Path) -> None:
        """Install the specified bundle in not booted slot

        RAUC daemon will start the process in the background

        :param bundleFile: path to bundle file
        """
        self._installCompleted = False
        self._completedEvent.clear()
        await self._installer.call_install_bundle(bundleFile.as_posix(), {})
        self.startProgressMonitor()

    def handelCompleted(self, status: int) -> None:
        """Callback for completed installation of bundle

        :param status: '0' if successful'
        """
        print(f"--> Bundle install completed: {status}")
        self._installCompleted = True
        self._completedEvent.set()

    def startProgressMonitor(self) -> None:
        """Start a task which monitors the update process"""

        async def _progressMonitor():
            while not self._installCompleted:
                percentage, message, nestingDepth = await self._installer.get_progress()
                print(f"    install progress:  {percentage:3d}% ({message})")
                await asyncio.sleep(PROGRESS_LOG_INTERVAL_S)

        self._progressTask = asyncio.create_task(_progressMonitor())

    async def stopProgressMonitor(self) -> None:
        """Stop the process monitor task, if running"""
        if self._progressTask:
            with contextlib.suppress(asyncio.CancelledError):
                self._progressTask.cancel()
                await self._progressTask
            self._progressTask = None

    async def waitForCompletion(self) -> None:
        """Wait until update has completed"""
        await self._completedEvent.wait()

    @staticmethod
    def prettyFormatSlotData(slotData: Dict) -> str:
        """Format the full slot data dict into a human-readable form

        :param slotData: slot data dict
        :return: pretty formatted data
        """
        infos = []
        for name, details in slotData.items():
            infos.append(
                f"--> {name} - Slot {details['bootname']}"
                + "\n    ".join(
                    [""] + [f"{info}={details.get(info, 'unknown')}" for info in ("boot-status", "state", "status")]
                )
            )
        return "\n".join(infos)


class HttpError(Exception):
    """Custom exception for HTTP download"""

    def __init__(self, status: int):
        self.status = status


class Downloader:
    """HTTP downloader, capable of resuming"""

    def __init__(self, file: Path, url: str, noResume: bool = False, maxFreeSpace: int = None):
        """Initialize downloader with source and destination

        :param file: destination file
        :param url: source URL
        :param noResume: do not try to resume if `True` is passed
        """
        self._file = file
        self._complete = False
        self._noResume = noResume
        self._progress = 0
        self._url = url
        self._maxFreeSpace = maxFreeSpace  # if set, check if there is enough available space
        self._session = aiohttp.ClientSession()

    async def getSize(self) -> int:
        """Retrieve remote file size from HTTP header

        :return: the remote file size in bytes
        :raises: aiohttp.ClientError, HttpError
        """
        async with self._session.head(self._url) as response:
            if response.status == 200:
                return response.content_length
            else:
                raise HttpError(response.status)

    async def download(self) -> None:
        """Start download process.

        Raises 'asyncio.TimeoutError' if the download stops unexpectedly."""
        # Check remote size (in bytes)
        remoteSize = await self.getSize()
        print(f"    file size on remote is: {remoteSize} bytes\n")

        if self._maxFreeSpace and remoteSize > self._maxFreeSpace:
            print("    Not enough free space for download!\n     Need to free up some space before downloading...")
            self._complete = False
            return

        # Detect if we need to continue
        sizeDownloaded = self._file.stat().st_size if self._file.is_file() else 0

        if self._noResume:
            print("    force full download of image bundle")
            self._file.unlink(missing_ok=True)
            sizeDownloaded = 0

        if sizeDownloaded == remoteSize:
            print("    file is already completely downloaded, skip...")
            self._complete = True
            return
        elif sizeDownloaded > remoteSize:
            print("    local file is larger then remote, must be corrupted or changed, full download...")
            sizeDownloaded = 0

        resumeHeader = {"Range": f"bytes={sizeDownloaded}-"} if sizeDownloaded else None
        mode = "ab" if sizeDownloaded else "wb"
        progressTask = None
        try:
            async with self._session.get(self._url, headers=resumeHeader) as response:
                totalSize = sizeDownloaded or 0
                with self._file.open(mode) as f:
                    f: BinaryIO
                    # start progress bar
                    progressTask = asyncio.create_task(self.logProgress())
                    while True:
                        chunk = await asyncio.wait_for(response.content.read(1024), DOWNLOAD_KB_TIMEOUT_S)
                        if not chunk:
                            break
                        f.write(chunk)
                        totalSize += len(chunk)
                        self._progress = float(totalSize) / float(remoteSize) * 100.0
                print("    Download completed!\n")
                self._complete = True

        except aiohttp.ClientError as e:
            print(f"    Error occurred during download: {e}")
            return False
        finally:
            if progressTask:
                with contextlib.suppress(asyncio.CancelledError):
                    progressTask.cancel()
                    await progressTask

    async def logProgress(self):
        """Callback to log the progress during download"""
        while True:
            print(f"    download progress: {self._progress:5.1f}%")
            await asyncio.sleep(PROGRESS_LOG_INTERVAL_S)

    async def getHash(self) -> Optional[str]:
        """Calculate MD5 hash of downloaded file

        :return: hash or `None` on error
        """

        def hasher(file: Path):
            md5Hash = hashlib.md5()
            with file.open("rb") as f:
                while True:
                    if chunk := f.read(1024):
                        md5Hash.update(chunk)
                    else:
                        break
            return md5Hash.hexdigest()

        if not self._file.is_file():
            return

        print("    calculate MD5 hash of image bundle...", end="", flush=True)
        loop = asyncio.get_running_loop()
        try:
            md5Hash = await asyncio.wait_for(loop.run_in_executor(None, hasher, self._file), timeout=HASH_TIMEOUT_S)
        except asyncio.TimeoutError:
            print("TIMEOUT ERROR!")
            md5Hash = None
        else:
            print(f" {md5Hash or 'unknown or error'}")
        return md5Hash

    async def close(self):
        if self._session:
            await self._session.close()
            self._session = None


class SysInfo:
    """Collect various information about the underlying system and provide system level actions"""

    OS_RELEASE_FILE = Path("/etc/os-release")
    OPKG_PERSISTENT_DIR = Path("/mnt/persistent-data/opkg-persistent")
    ACTIVE_CHARGING_STATES = ("contractAuth", "initConnection", "currentDemand")

    def __init__(self, noAsk=False):
        """
        :param noAsk: pass `True` if interactive user input should be suppressed
        """
        self._serial = None
        self._role = None
        self._imageVersion = None
        self._hwWdTask = None
        self._linuxWdTask = None
        self._onTarget = True
        self._nrCores = os.cpu_count()
        self._ramMb = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / (1024.0**2)
        self._configuredPlugs = []
        self._sudoPassword = os.getenv(SUDO_PASS_ENV)
        self._noAsk = noAsk

        if "arm" not in platform.machine().lower():
            print("    !!! not on target device !!!")
            self._onTarget = False
            return

        configFile = Path(ECP_DATA_DIR) / CONFIG_FILE_LOCATION
        if not configFile.is_file():
            print("ERROR: No ECP config found")
        else:
            config = json.loads(configFile.read_bytes())
            with contextlib.suppress(KeyError):
                self._serial = config["Device"]["Serial"]["Serial"]
                self._role = config["Device"]["Role"]
                self._configuredPlugs = config["Plugs"]["ConfiguredPlugs"]

        osRelease = self.OS_RELEASE_FILE.read_text()
        if match := re.search(r'"?([\d.]+)', osRelease):
            self._imageVersion = match[1]

    @property
    def serial(self) -> str:
        """Get device serial number from config

        :return: e.g. '1234'
        """
        return self._serial

    @property
    def isDisplay(self) -> bool:
        """Check if running on display board

        :return: `True` on display
        """
        DISPLAY_ROLES = ("RemoteDisplay1", "RemoteHMI1", "RemoteHMI2")
        return self._role in DISPLAY_ROLES

    @property
    def configuredPlugs(self) -> str:
        """Read configured plugs from config

        :return: e.g. `['Plug0', 'Plug2']
        """
        return ", ".join(self._configuredPlugs) or "none"

    @property
    def imageVersion(self) -> str:
        """Currently installed image version

        :return: e.g. '1.0.20'
        """
        return self._imageVersion

    @property
    def cpuCores(self) -> int:
        """Number of CPU cores

        :return: number of cores
        """
        return self._nrCores

    @property
    def systemRam(self) -> int:
        """Total system RAM

        :return: system RAM in MB
        """
        return int(round(self._ramMb))

    def askForSudoPassword(self) -> str:
        """Interactively (if not cashed or passed in env) ask for sudo password

        :return: the sudo password
        """
        if self._sudoPassword is None:
            assert not self._noAsk, f"Provide sudo password via env {SUDO_PASS_ENV}"
            self._sudoPassword = getpass("Provide sudo password: ")
        return self._sudoPassword

    def displayIps(self) -> Tuple[str, str]:
        """Calculate possible IPs of displayboards from serial number

        :return: e.g. `("10.2.101.23", "10.2.201.23")`
        """
        highbyte = int((self._serial - 1) / 100)
        highbyte %= 100  # Prevent overflow if Serial higher than 9999
        lowbyte = (self._serial - 1) % 100 + 1

        # 10.2.1ss.ss -> ssss -> Serial No (except 100 * x Serial No)
        addressDisplay1 = f"10.2.1{str(highbyte).zfill(2)}.{str(lowbyte)}"

        # 10.2.2ss.ss -> ssss -> Serial No (except 100 * x Serial No)
        addressDisplay2 = f"10.2.2{str(highbyte).zfill(2)}.{str(lowbyte)}"

        return addressDisplay1, addressDisplay2

    def baseboardIp(self) -> str:
        """Calculate baseboard IP from serial number

        :return: e.g. "10.2.1.23"
        """
        highbyte = int((self._serial - 1) / 100)
        highbyte %= 100  # Prevent overflow if Serial higher than 9999
        lowbyte = (self._serial - 1) % 100 + 1

        # 10.2.ss.ss -> ssss -> Serial No (except 100 * x Serial No)
        return f"10.2.{highbyte}.{lowbyte}"

    def getPlugStates(self) -> Dict[str, str]:
        """Get current plug state of configured plugs

        :return: e.g. ``{"Plug0": "ready"}`
        """
        statePattern = re.compile(r"^([0-9-:.]+;){3}([^;]+)")
        # Check plug logs if we see an active session
        plugStates = {}
        for plug in self._configuredPlugs:
            plugStates[plug] = "unknown"
            logFile = Path(ECP_DATA_DIR) / LOG_DIR / f"{plug}.log"
            if not logFile.is_file():
                continue
            with FileReadBackwards(logFile) as plugLog:
                count = 3
                for line in plugLog:
                    if match := re.search(statePattern, line):
                        plugStates[plug] = match[2].strip()
                        break
                    elif not count:
                        break
                    count -= 1
        return plugStates

    @staticmethod
    def checkActiveSessionInPlugStates(plugStates: Dict[str, str]) -> List[str]:
        """Check if any plug is in charging state

        :param plugStates: the plug states dict
        :return: list of plugs which are in charging states
        """
        return [
            plug
            for plug, plugState in plugStates.items()
            if any(s.startswith(plugState) for s in SysInfo.ACTIVE_CHARGING_STATES)
        ]

    async def pingDisplays(self) -> Tuple[bool, bool]:
        """Ping displays

        :return: tuple of ping states
        """
        displays = self.displayIps()
        d1 = await self.pingHost(displays[0])
        d2 = await self.pingHost(displays[1])
        return d1, d2

    @staticmethod
    async def pingHost(host: str) -> bool:
        """Ping the given host

        :param host: host to ping
        :return: `True` if ping was successful
        """

        async def _runPingCommand():
            cmd = ["ping", "-c", "1", host]  # Ping once
            proc = await asyncio.create_subprocess_exec(*cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            await proc.communicate()  # Wait for the process to complete
            return proc.returncode

        try:
            # Run the ping command with a timeout of 5 seconds
            result = await asyncio.wait_for(_runPingCommand(), timeout=5)
            return result == 0
        except asyncio.TimeoutError:
            return False

    def enableHardwareWatchdogFeed(self) -> None:
        """Enable periodic feeding of external watchdog IC

        The corresponding GPIO is set to low most of the time, a periodic high will delay external reset
        """
        if not self._onTarget:
            return

        async def togglePin():
            try:
                while True:
                    HW_WD_GPIO_PATH.write_text("1")
                    HW_WD_GPIO_PATH.write_text("0")
                    await asyncio.sleep(10)
            except asyncio.CancelledError:
                HW_WD_GPIO_PATH.write_text("0")
                raise

        if self._hwWdTask is None:
            self._hwWdTask = asyncio.create_task(togglePin(), name="HardwareWD-toggle")

    async def disableHardwareWatchdog(self, letStarve: bool = False) -> None:
        """Disable external watchdog IC

        :param letStarve: if set, the external reset will happen soon
        """
        if self._hwWdTask:
            with contextlib.suppress(asyncio.CancelledError):
                self._hwWdTask.cancel()
                await self._hwWdTask
            HW_WD_GPIO_PATH.write_text("0" if letStarve else "1")

    def enableLinuxWatchdogFeed(self) -> None:
        """Enable periodic Linux watchdog feed"""
        if not self._onTarget:
            return

        async def feedWd():
            fd = None
            try:
                while True:
                    try:
                        if fd is None:
                            fd = os.open(WATCHDOG_DEVICE, os.O_WRONLY)
                    except OSError:
                        fd = None
                    else:
                        fcntl.ioctl(fd, WDIOC_KEEPALIVE, 0)
                    await asyncio.sleep(10)
            except asyncio.CancelledError:
                if fd:
                    os.close(fd)
                raise

        if self._linuxWdTask is None:
            self._linuxWdTask = asyncio.create_task(feedWd(), name="LinuxWD-feed")

    async def disableLinuxWatchdogFeed(self, letStarve: bool = False):
        """Disable Linux watchdog feeding

        :param letStarve: if set, the reset will happen soon
        """
        if self._linuxWdTask:
            with contextlib.suppress(asyncio.CancelledError):
                self._linuxWdTask.cancel()
                await self._linuxWdTask

        if not letStarve:
            fd = None
            try:
                fd = os.open(WATCHDOG_DEVICE, os.O_WRONLY)
                os.write(fd, MAGIC_CLOSE_CHARACTER)
            except OSError:
                print("WARNING: can't disable LinuxWatchdog")
            finally:
                if fd:
                    os.close(fd)

    async def setTime(self, epoch: float) -> None:
        """Set system time

        :param epoch: unix epoch timestamp
        """
        sudoPass = self.askForSudoPassword()
        cmd = f'echo "{sudoPass}" | sudo -S date -s "@{int(round(epoch))}"'
        proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        await proc.communicate()

    @staticmethod
    async def checkServiceRunning(service: str = None) -> bool:
        """Check if given service is running

        :param service: name of service, e.g. "ECP"
        :return: `True` if service is running
        """
        service = service or "ECP"
        cmd = ["systemctl", "is-active", "--quiet", f"{service}.service"]
        proc = await asyncio.create_subprocess_exec(*cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        await proc.communicate()
        return proc.returncode == 0  # will return 0 for active else inactive.

    async def stopService(self, service: str = None) -> bool:
        """Stop given service

        :param service: name of service, e.g. "ECP"
        :return: `True` on success
        """
        sudoPass = self.askForSudoPassword()
        service = service or "ECP"

        cmd = f'echo "{sudoPass}" | sudo -S systemctl stop {service}.service'
        proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        await proc.communicate()
        return proc.returncode == 0  # 0 on success

    async def deleteOpkgPersistent(self) -> None:
        """Delete the opkg persistent folder"""
        sudoPass = self.askForSudoPassword()

        if not self.OPKG_PERSISTENT_DIR.is_dir():
            return

        cmd = f'echo "{sudoPass}" | sudo -S rm -rf {self.OPKG_PERSISTENT_DIR.as_posix()}'
        proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        await proc.communicate()

    async def askForReboot(self) -> bool:
        """Interactively ask user to reboot

        Skip if in non-interactive mode.

        :return: True if reboot is requested
        """
        if self._noAsk:
            print("    reboot skipped! (non-interactive)")
            return False
        loop = asyncio.get_running_loop()
        userInput = await loop.run_in_executor(None, input, "\n===> press [y] if you want to reboot... ")
        if userInput.lower() == "y":
            print("    rebooting...")
            return True

        print("    reboot skipped!")
        return False

    @staticmethod
    async def parsePartitionSpace() -> int:
        """

        Returns: free space in bytes

        If parsing for the '/mnt/persistent-data' partition should fail, returns 0

        """
        cmd = ["df", "--block-size=1"]
        proc = await asyncio.create_subprocess_exec(*cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdoutStream: bytes = await proc.stdout.read()

        lines: List[str] = stdoutStream.decode().split("\n")

        await proc.communicate()

        parsedValues = [None for _ in range(6)]
        for l in lines:
            if "/mnt/persistent-data" in l:
                singleSpaced = " ".join(l.split())
                parsedValues = singleSpaced.split(" ")

        info = {
            "device": parsedValues[0] or "device not found",
            "size": int(parsedValues[1]) or 0,
            "used": int(parsedValues[2]) or 0,
            "avail": int(parsedValues[3]) * 1024 or 0,
            "use%": parsedValues[4] or 0,
            "mounted on": parsedValues[5] or "Mount point not found",
        }

        print(f"    Available disk space: {info['avail']} ({info['used']} of {info['size']} bytes used.)")

        return info["avail"]


class SshAgent:
    """An agent class to interact with a remote system via SSH

    We use AsyncSSH under the hood. We have to get it from PyPi, because it is not part of the EVTEC image.
    """

    ASYNCSSH_WHEEL = (
        "https://files.pythonhosted.org/packages/b7/f2/"
        "83301b225123cc991659903a1a5fedd9fcdb7f3d2efe1461f25f1eab8c5b/asyncssh-2.8.1-py3-none-any.whl"
    )
    # older version which still works with cryptography v2.8

    def __init__(self, ip: str, sudoPass: str):
        """Initialize with desired host and sudo password

        :param ip: remote host
        :param sudoPass: password
        """
        self._ip = ip
        self._sudoPass = sudoPass
        self._proc = None
        self._connection: Optional[asyncssh.SSHClientConnection] = None

    @property
    def ip(self):
        """Remote hosts IP address"""
        return self._ip

    @staticmethod
    def asyncSshWheelFile() -> Path:
        """Local file path to AsynSSH wheel"""
        packageDir = Path(WORKDIR_DEFAULT) / "site-packages"
        return packageDir / SshAgent.ASYNCSSH_WHEEL.split("/")[-1]

    @staticmethod
    async def preparePackage() -> None:
        """Get and install AsyncSSH locally"""
        try:
            import asyncssh
        except ImportError:
            print("    AsyncSSH unavailable, install...")
        else:
            return

        wheelFile = SshAgent.asyncSshWheelFile()
        packageDir = wheelFile.parent
        packageDir.mkdir(parents=True, exist_ok=True)
        if wheelFile.is_file():
            print("    AsyncSSH already available, skip download")
        else:
            async with aiohttp.ClientSession() as session:
                async with session.get(SshAgent.ASYNCSSH_WHEEL) as resp:
                    if resp.status == 200:
                        with wheelFile.open("wb") as f:
                            f.write(await resp.read())
            print("    downloaded AsyncSSH")

        # extract zip
        with ZipFile(wheelFile, "r") as f:
            f.extractall(packageDir)
        print("    wheel extracted")

        # try to load the package
        sys.path.append(packageDir.as_posix())
        import asyncssh  # noqa

    async def connect(self) -> None:
        """Connect to remote host"""
        import asyncssh  # noqa

        flags = AI_NUMERICHOST
        print(f"    Try SSH connection to host {self._ip}")
        self._connection = await asyncssh.connect(
            self._ip, password=self._sudoPass, known_hosts=None, flags=flags, family=AF_INET
        )
        print(f"    SSH connection to host {self._ip} established")

    async def transferBundle(self, bundle: Path) -> None:
        """Transfer remote bundle to local system

        :param bundle: path to bundle on remote host
        """
        import asyncssh  # noqa

        # make sure bundle file has user permissions
        cmds = (
            f'echo "{self._sudoPass}" | sudo -S chown {getuser()}:{getuser()} {bundle.as_posix()}',
            f'echo "{self._sudoPass}" | sudo -S chmod 664 {bundle.as_posix()}',
        )
        for cmd in cmds:
            proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            await proc.communicate()

        lastProgress = [0]
        progressHandler = functools.partial(self._scpProgressHandler, lastProgress)

        print(f"    transfer {bundle.name}")
        await asyncssh.scp((self._connection, bundle), bundle, progress_handler=progressHandler)

    async def uploadScript(self) -> None:
        """Upload the update script and AsyncSSH package to remote host"""
        import asyncssh  # noqa

        lastProgress = [0]
        progressHandler = functools.partial(self._scpProgressHandler, lastProgress)

        wheelFile = SshAgent.asyncSshWheelFile()
        print(f"    upload {wheelFile.name}")
        await self._connection.run(f"mkdir -p {wheelFile.parent.as_posix()}")
        lastProgress[0] = 0
        await asyncssh.scp(wheelFile, (self._connection, wheelFile), progress_handler=progressHandler)

        scriptFile = Path(__file__)
        print(f"    upload {scriptFile.name}")
        workDir = Path(WORKDIR_DEFAULT)
        lastProgress[0] = 0
        await asyncssh.scp(scriptFile, (self._connection, workDir / scriptFile.name), progress_handler=progressHandler)

    async def performUpdate(self, version: str, sudoPass: str, force: bool, noFix: bool) -> None:
        """Call the update script on the remote host

        :param version: image version
        :param sudoPass: sudo password
        :param force: force update if system image version is already the same
        :param noFix: do not apply configuration file fixes
        """
        import asyncssh  # noqa

        scriptFile = Path(__file__)
        workDir = Path(WORKDIR_DEFAULT)
        cmd = (
            # the 'u' is important to un-buffer stdout
            f"{SUDO_PASS_ENV}={sudoPass} python3 -u {(workDir / scriptFile.name).as_posix()} "
            f"--noScreen --bundleHash=skip --noAsk --noDownload "
            f" --noWatchdog "  # TODO: remove
            f"{'--forceInstall ' if force else ''}"
            f"{'--noFix ' if noFix else ''}"
            f"--setTime {int(round(datetime.utcnow().timestamp()))} "
            f"{version}"
        )
        async with self._connection.create_process(cmd) as process:
            while True:
                line = await process.stdout.readline()
                if not line:
                    break
                print(f"### Display {self._ip} ### {line}", end="", flush=True)
            print(f"### Display {self._ip} command done ###")
            error = await process.stderr.read()
            if error:
                print(f"### Display {self._ip} ### ERROR ###\n{error}\n####################")

    async def doReboot(self, sudoPass: str) -> None:
        """Reboot remote host

        Does also arm the external reset IC.
        :param sudoPass: password
        """
        import asyncssh  # noqa

        cmds = (
            f"echo 0 > {HW_WD_GPIO_PATH.as_posix()}",  # let external WD starve
            f"echo '{sudoPass}' | sudo -S systemctl reboot",  # perform reboot
        )

        for cmd in cmds:
            await self._connection.run(cmd)

    @staticmethod
    def _scpProgressHandler(lastProgress: List[int], srcPath: bytes, destPath: bytes, copied: int, total: int):
        del (srcPath, destPath)
        progress = int(copied / total * 100.0)
        if progress >= lastProgress[0] + 10:  # print in 10% steps
            print(f"    scp progress: {progress:3d}%")
            lastProgress[0] = progress


class FixAction(Enum):
    """Enumeration of possible configuration fixes"""

    REMOVE = "remove"
    # can be expanded if needed


@dataclass
class FixRule:
    """Definition of a config fix rule"""

    name: str
    """Name of the fix rule"""
    fullKey: str
    """Full configuration path to old key, forward slash separated"""
    actions: List[FixAction]
    """Kind of fix"""
    replacement: Optional[str] = None
    """Optional replacement value"""
    valueMatch: Optional[Union[str, Callable]] = None
    """Rule is only applied if old value matches exactly or the given lambda evaluates to true (old value is passed in)"""


class ConfigFixer:
    """Configuration fixer to change, remove rewrite configs after image update"""

    def __init__(self, configFile: Path):
        """Pass configuration file

        :param configFile: path to config.json
        """
        self._file = configFile
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        self._fileBackup = configFile.parent / f"{configFile.stem}_backup_imageUpdate_{timestamp}{configFile.suffix}"
        self._rules: List[FixRule] = []
        self._configContent: str = ""
        self._configParsed: Dict[Any, Any] = {}

    def addRule(self, rule: FixRule) -> None:
        """Add a fix rule

        :param rule: the rule
        """
        self._rules.append(rule)

    def execute(self):
        """Try to apply all fix rules

        Will create a backup of old configuration (e.g. `config_backup_imageUpdate_2023-08-21_13-46-10.json`).
        If current config file is not a valid JSON document, the process is skipped.
        """
        if not self._file.is_file():
            print(f"    ERROR: config file not found ({self._file.as_posix()})")
            return
        self._configContent = self._file.read_text()
        self._fileBackup.write_text(self._configContent)
        print(f"   created backup: {self._fileBackup.name}")
        try:
            self._configParsed = json.loads(self._configContent)
        except JSONDecodeError as ex:
            print(f"    ERROR: config contains invalid JSON ({self._file.as_posix()})\n{ex}")
            return

        configChanged = False
        for rule in self._rules:
            print(f"--> execute config rule '{rule.name}'")
            fullKey = rule.fullKey.split("/")
            currentValue = self._recursiveDictGet(self._configParsed, fullKey)
            if callable(rule.valueMatch):
                doSkip = not rule.valueMatch(currentValue)
            else:
                doSkip = rule.valueMatch != currentValue
            if doSkip:
                print("    rule not applicable -> skip")
                continue

            if FixAction.REMOVE in rule.actions:
                try:
                    self._recursiveDictDelete(self._configParsed, fullKey)
                    print(f"    removed '{rule.fullKey}'")
                except KeyError:
                    print(f"    key '{rule.fullKey}' to remove not found")
                else:
                    configChanged = True

        if configChanged:
            print("    write changed config")
            self._configContent = json.dumps(self._configParsed, indent=4)
            self._file.write_text(self._configContent)
        else:
            print("    config did not change")

    @staticmethod
    def _recursiveDictDelete(d: Dict, keys: List[str]) -> None:
        """Traverse a nexted dict to delete a nested key.

        :param d: nested dict
        :param keys: nested key, values in list
        """
        if not keys:
            return
        key = keys[0]
        if key in d:
            if len(keys) == 1:
                del d[key]
            else:
                ConfigFixer._recursiveDictDelete(d[key], keys[1:])
        else:
            raise KeyError()

    @staticmethod
    def _recursiveDictGet(d: Dict, keys: List[str]):
        """Traverse a nexted dict to get a nested key.

        :param d: nested dict
        :param keys: nested key, values in list
        """
        if not keys:
            return
        key = keys[0]
        if key in d:
            if len(keys) != 1:
                return ConfigFixer._recursiveDictGet(d.get(key), keys[1:])
            return d.get(key)


def _validateHashArg(value: str, pat=re.compile(r"^[a-f0-9A-F]{32}$")):
    """Arg parse validator for MD5 hashes"""
    if not pat.match(value) and value != "skip":
        raise argparse.ArgumentTypeError("value is not a MD5 hash")
    return value.lower()


async def runUpdateProcess(args, sysInfo: SysInfo) -> None:
    """Run the whole update process

    If the update process fails, an exception will be raised.
    Among possible exceptions are: HttpError, asyncio.TimeoutError, aiohttp.ClientConnectorError

    :param args: parsed arguments from ArgParse
    :param sysInfo: system information object
    """
    # Build download URL
    print("Remote bundle URL:")
    if args.noDownload:
        print("   download disabled!")
        url = None
    else:
        credentials = (os.getenv(RAUC_LOGIN_ENV) or "").strip()
        credentialsMasked = ""
        if credentials:
            credentialsMasked = f"{credentials.partition(':')[0]}:****@"
            credentials += "@"

        url = DEFAULT_UPDATE_SERVER_URL_TEMPLATE.format(credentials, args.version)
        print(f"    {DEFAULT_UPDATE_SERVER_URL_TEMPLATE.format(credentialsMasked, args.version)}")

    # Setup work dir and bundle file
    workdir = Path(WORKDIR_DEFAULT)
    workdir.mkdir(exist_ok=True, parents=True)
    bundleFile = workdir / f"update-bundle-{args.version.replace('.','_')}.raucb"
    print(f"    target bundle file: '{bundleFile.as_posix()}'")

    print("===> Gather SystemInfo")
    print(
        f"    Serial:           {sysInfo.serial}\n"
        f"    Is DisplayBoard:  {sysInfo.isDisplay}\n"
        f"    Image Version:    {sysInfo.imageVersion}\n"
        f"    CPU Cores:        {sysInfo.cpuCores}\n"
        f"    System RAM:       {sysInfo.systemRam}MB"
    )

    if sysInfo.cpuCores < 2 or sysInfo.systemRam < 300:
        print("    WARNING: slow hardware (SOM) detected!")

    if args.setTime:
        print("===> Set System Time")
        setTime = datetime.fromtimestamp(float(args.setTime))
        if abs(datetime.utcnow() - setTime) > timedelta(minutes=5):
            print(f"    system time is out of sync, setting to: {setTime.strftime('%d.%m.%Y %H:%M')}")
            await sysInfo.setTime(setTime.timestamp())
        else:
            print("    system time already in sync")

    displaysReachable = (False, False)
    if args.includeDisplays and not sysInfo.isDisplay:
        print("===> Ping Displays")
        displaysReachable = await sysInfo.pingDisplays()
        print(f"    reached displays (#1, #2): {displaysReachable}")

    if any(displaysReachable):
        await SshAgent.preparePackage()

    displayAgents: List[Optional[SshAgent]] = [
        SshAgent(ip=sysInfo.displayIps()[i], sudoPass=sysInfo.askForSudoPassword()) if displaysReachable[i] else None
        for i in range(1)
    ]
    for agent in displayAgents:
        if agent is None:
            continue
        await agent.connect()

    print("===> Check for active charging sessions")
    if sysInfo.isDisplay:
        print("    skip on display board")
    elif not await sysInfo.checkServiceRunning():
        print("    skip, ECP not running")
    else:
        print(f"    configured Plugs: {sysInfo.configuredPlugs}")
        plugStates = sysInfo.getPlugStates()
        plugCharging = sysInfo.checkActiveSessionInPlugStates(plugStates)
        print(f"    read plug states from logs: {plugStates}")
        print(f"    active charging sessions on: {', '.join(plugCharging) or 'none'}")

        startWaitingAt = time.monotonic()
        while plugCharging and args.waitTime >= 0:
            print(f"    waiting for charge end (charging: {', '.join(plugCharging)}")
            await asyncio.sleep(10)
            plugStates = sysInfo.getPlugStates()
            plugCharging = sysInfo.checkActiveSessionInPlugStates(plugStates)
            if (time.monotonic() - startWaitingAt) > args.waitTime != 0:
                print(f"    maximum wait time reached ({args.waitTime} seconds)")
                break

        print("    charge end reached, continue...")

    print("===> Stop Services")
    for service in ("ECP", "Display", "HmiControl"):
        if not await sysInfo.checkServiceRunning(service):
            print(f"    {service} not running")
            continue
        print(f"    stop {service}")
        await sysInfo.stopService(service)
    await asyncio.sleep(2)

    print("===> Start Linux Watchdog feeding")
    sysInfo.enableLinuxWatchdogFeed()

    if not args.noWatchdog:
        print("===> Arming external Watchdog")
        sysInfo.enableHardwareWatchdogFeed()

    if url and not sysInfo.isDisplay:
        print("===> Download update bundle")

        availableSpace: int = await sysInfo.parsePartitionSpace()

        try:
            down = Downloader(file=bundleFile, url=url, noResume=args.noResume, maxFreeSpace=availableSpace)
            await down.download()
            await down.close()

        except asyncio.TimeoutError:
            # download did timeout
            print("    ERROR: download did timeout!")
            raise

        if args.bundleHash != "skip":
            downloadedHash = await down.getHash()
            if downloadedHash == args.bundleHash:
                print("    downloaded bundle has expected MD5 hash")
            elif args.bundleHash:
                print("    ERROR: downloaded bundle does not match expected MD5 hash")
                return

    # Copy update files to displays
    for agent in displayAgents:
        if agent is None:
            continue
        print(f"===> Copy update to Display {agent.ip}")
        await agent.uploadScript()
        print(f"===> Perform update on Display {agent.ip}")
        # make sure bundle file has user permissions
        cmds = (
            f'echo "{sysInfo.askForSudoPassword()}" | sudo -S chown {getuser()}:{getuser()} {bundleFile.as_posix()}',
            f'echo "{sysInfo.askForSudoPassword()}" | sudo -S chmod 664 {bundleFile.as_posix()}',
        )
        for cmd in cmds:
            proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            await proc.communicate()
        try:
            await asyncio.wait_for(
                agent.performUpdate(
                    version=args.version,
                    sudoPass=sysInfo.askForSudoPassword(),
                    force=args.forceInstall,
                    noFix=args.noFix,
                ),
                INSTALL_TIMEOUT_DISPLAY_S,
            )
        except asyncio.TimeoutError:
            print("    ERROR: Timeout happened during update of display!")
            raise

    print("===> Start RAUC update process")
    doInstall = not args.noRauc
    if not doInstall:
        print("    skipped by user")
    elif sysInfo.imageVersion != args.version:
        print(f"    update version {sysInfo.imageVersion} -> {args.version}")
    elif args.forceInstall:
        print(f"    local image is already on version {args.version} but installation is force")
    else:
        print(f"    skipped, image is already on version {args.version}")
        doInstall = False

    # Get bundle from baseboard if we are on display
    if doInstall and sysInfo.isDisplay:
        print("    transfer update bundle from baseboard")
        agent = SshAgent(ip=sysInfo.baseboardIp(), sudoPass=sysInfo.askForSudoPassword())
        await agent.preparePackage()
        await agent.connect()
        await agent.transferBundle(bundleFile)

    if doInstall:
        rauc = Rauc()
        await rauc.connect()

        print("===> Pre-update slot data")
        primaryPre = await rauc.getPrimarySlot()
        print(f"    Primary Slot: {primaryPre}")
        slotDataPre = await rauc.getSlotStatus()
        print(rauc.prettyFormatSlotData(slotDataPre))

        print(f"===> Install Bundle (timeout: {INSTALL_TIMEOUT_MIN} minutes)")
        await asyncio.sleep(1)
        await rauc.installBundle(bundleFile)

        try:
            await asyncio.wait_for(rauc.waitForCompletion(), timeout=INSTALL_TIMEOUT_MIN * 60)
        except asyncio.TimeoutError:
            print("    ERROR: installation did timeout")
            raise

        print("===> Post-update slot data")
        primaryPost = await rauc.getPrimarySlot()
        print(f"    Primary Slot: {primaryPost}")
        slotDataPost = await rauc.getSlotStatus()
        print(rauc.prettyFormatSlotData(slotDataPost))
        print()
        if primaryPre != primaryPost:
            print("    Primary slot changed")
        else:
            print("    WARNING: primary slot did not change with update (ok if update was repeated)")

        primarySlotState: str = slotDataPost[primaryPost]["state"]
        if primarySlotState == "booted":
            print(f"    WARNING: active slot after update '{primaryPost}' is already '{primarySlotState}'!")

        print("===> Delete OPKG persistent")
        await sysInfo.deleteOpkgPersistent()
        print("    ... deleted")

    print("===> Fix config.json")
    if not args.noFix:
        fixer = ConfigFixer(Path(ECP_DATA_DIR) / CONFIG_FILE_LOCATION)
        fixer.addRule(
            FixRule(
                name="LAN DHCP, old key removal",
                fullKey="Interfaces/Network/DHCP",
                actions=[FixAction.REMOVE],
                valueMatch=lambda v: isinstance(v, str),
            )
        )
        fixer.execute()
    else:
        print("   skipped by user")

    # Do reboot (if desired)
    if await sysInfo.askForReboot():
        sudoPass = sysInfo.askForSudoPassword()
        for agent in displayAgents:
            if agent is None:
                continue
            print(f"===> Reboot display {agent.ip}")
            await agent.doReboot(sudoPass=sudoPass)

        # let watchdogs starve
        await sysInfo.disableHardwareWatchdog(True)
        await sysInfo.disableLinuxWatchdogFeed(True)

        print("     Do not forget to remove rauc-bundle manually.")

        # Perform reboot
        print("===> Reboot baseboard")
        cmd = f"echo '{sudoPass}' | sudo -S systemctl reboot"
        proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        await proc.communicate()
        sys.exit(os.EX_OK)
    else:
        print("===> Reboot skipped.")
        print("     Stop feed of LinuxWatchdog.")
        # Do not reboot
        await sysInfo.disableHardwareWatchdog()
        await sysInfo.disableLinuxWatchdogFeed()


class ArgFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
    """Custom formatter to show defaults in arguments and rich format the long documentation"""


async def main(argList: List[str]):
    """Main async entry point

    Runs the update process with correct arguments. If the update fails, it must raise an exception,
    which will be caught inside this function. In this case, an appropriate message will be prompted.
    """
    parser = argparse.ArgumentParser(
        description="Automated EVTEC Image Update Script", epilog=FULL_DOC, formatter_class=ArgFormatter
    )
    parser.add_argument("version", help="image version to update to, e.g. '1.0.8' or 'latest'")
    parser.add_argument("--bundleHash", type=_validateHashArg, help="expected MD5 hash of bundle")
    parser.add_argument(
        "-d", "--includeDisplays", action="store_true", help="update all available display boards as well"
    )
    parser.add_argument("--noResume", action="store_true", help="force full download of bundle, do not resume")
    parser.add_argument(
        "-f",
        "--forceInstall",
        action="store_true",
        help="force installation of bundle if installed image is on the same version",
    )
    parser.add_argument(
        "-i", "--noDownload", action="store_true", help="do not download a bundle, use locally available"
    )
    parser.add_argument("--noRauc", action="store_true", help="disable interaction with the RAUC service completely")
    parser.add_argument("--noWatchdog", action="store_true", help="disable arming the watchdog device")
    parser.add_argument(
        "--waitTime", type=int, help="secs to wait for charge end, -1 will skip waiting, 0 waits forever", default=6000
    )
    parser.add_argument("--noAsk", action="store_true", help="do not ask interactively for user input")
    parser.add_argument("--noFix", action="store_true", help="do not try to fix the config")
    parser.add_argument("--setTime", type=str, help="set system time (epoch format) prior bundle verification")

    args = parser.parse_args(argList)
    sysInfo = SysInfo(noAsk=args.noAsk)

    print("=== Start updating progress ===")
    # Run process in a dedicated coroutine, to have error handling here
    updateSuccess = False
    try:
        await runUpdateProcess(args, sysInfo)
    except HttpError as ex:
        if ex.status == 401:
            print(f"ERROR: update server unauthorized, have you provided `{RAUC_LOGIN_ENV}`")
        else:
            print(f"ERROR: can't get update bundle (status= {ex.status})")
    except aiohttp.ClientConnectorError:
        print("ERROR: could not connect to update server, is the default gateway correct?")
    except Exception as ex:
        print(f"ERROR: unexpected exception: {ex}")
    else:
        updateSuccess = True

    if not updateSuccess:
        print("===> ERROR OCCURRED DURING UPDATE. Wait 3 minutes before triggering reboot.")
        await asyncio.sleep(T_WAIT_BEFORE_SHUTDOWN)
        print("     Stopping feed of LinuxWatchdog. Disable HW watchdog. Reboot will happen soon.")
        await sysInfo.disableLinuxWatchdogFeed(True)
        await sysInfo.disableHardwareWatchdog()

    print("===> DONE!!!")


if __name__ == "__main__":
    argsList = sys.argv[1:]
    SCREEN_ARG = "--noScreen"

    # check that python is called on the scope of screen (hidden argument `--noScreen` is used to detect that)
    if SCREEN_ARG not in argsList:
        print("WARNING: Script is called outside of a screen session, make sure to call the script in bash")
        input("press any key to exit... ")
        sys.exit(os.EX_USAGE)
    else:
        argsList.pop(argsList.index(SCREEN_ARG))

    # check that the script us not run as root (we need the sudo password in plaintext)
    if getuser() == "root":
        print(f"WARNING: Script is called as root, use environment variable '{SUDO_PASS_ENV}' instead")
        input("press any key to exit... ")

    try:
        asyncio.run(main(argsList))
    except KeyboardInterrupt:
        print("\nScript cancelled by user...")
