ANTA Device API

@abstractmethod
async def _collect(self, command: AntaCommand, *, collection_id: str | None = None) -> None:
    """Collect device command output.

    This abstract coroutine can be used to implement any command collection method
    for a device in ANTA.

    The `_collect()` implementation needs to populate the `output` attribute
    of the `AntaCommand` object passed as argument.

    If a failure occurs, the `_collect()` implementation is expected to catch the
    exception and implement proper logging, the `output` attribute of the
    `AntaCommand` object passed as argument would be `None` in this case.

    Parameters
    ----------
    command
        The command to collect.
    collection_id
        An identifier used to build the eAPI request ID.
    """

async def collect(self, command: AntaCommand, *, collection_id: str | None = None) -> None:
    """Collect the output for a specified command.

    When caching is activated on both the device and the command,
    this method prioritizes retrieving the output from the cache. In cases where the output isn't cached yet,
    it will be freshly collected and then stored in the cache for future access.
    The method employs asynchronous locks based on the command's UID to guarantee exclusive access to the cache.

    When caching is NOT enabled, either at the device or command level, the method directly collects the output
    via the private `_collect` method without interacting with the cache.

    Parameters
    ----------
    command
        The command to collect.
    collection_id
        An identifier used to build the eAPI request ID.
    """
    if self.cache is not None and command.use_cache:
        async with self.cache.locks[command.uid]:
            cached_output = await self.cache.get(command.uid)

            if cached_output is not None:
                logger.debug("Cache hit for %s on %s", command.command, self.name)
                command.output = cached_output
            else:
                await self._collect(command=command, collection_id=collection_id)
                await self.cache.set(command.uid, command.output)
    else:
        await self._collect(command=command, collection_id=collection_id)

async def collect_commands(self, commands: list[AntaCommand], *, collection_id: str | None = None) -> None:
    """Collect multiple commands.

    Parameters
    ----------
    commands
        The commands to collect.
    collection_id
        An identifier used to build the eAPI request ID.
    """
    await asyncio.gather(*(self.collect(command=command, collection_id=collection_id) for command in commands))

async def copy(self, sources: list[Path], destination: Path, direction: Literal["to", "from"] = "from") -> None:
    """Copy files to and from the device, usually through SCP.

    It is not mandatory to implement this for a valid AntaDevice subclass.

    Parameters
    ----------
    sources
        List of files to copy to or from the device.
    destination
        Local or remote destination when copying the files. Can be a folder.
    direction
        Defines if this coroutine copies files to or from the device.

    """
    _ = (sources, destination, direction)
    msg = f"copy() method has not been implemented in {self.__class__.__name__} definition"
    raise NotImplementedError(msg)

async def disconnect(self) -> None:
    """Disconnect the device and close any open connections.

    It is not mandatory to implement this for a valid AntaDevice subclass.
    If disconnection logic is not needed, implement this method as `pass`.
    Subclasses should document the concrete resources they close.

    NOTE:
        In ANTA 2.0, this method will be made abstract and must be implemented by all subclasses.
    """
    msg = f"disconnect() method has not been implemented in {self.__class__.__name__} definition"
    logger.warning(msg)

@abstractmethod
async def refresh(self) -> None:
    """Update attributes of an AntaDevice instance.

    This coroutine must update the following attributes of AntaDevice:

    - `is_online`: When the device IP is reachable and a port can be open.

    - `established`: When a command execution succeeds.

    - `hw_model`: The hardware model of the device.
    """

async def _collect(self, command: AntaCommand, *, collection_id: str | None = None) -> None:
    """Collect device command output from EOS using asynceapi.

    Supports outformat `json` and `text` as output structure.
    Gain privileged access using the `enable_password` attribute
    of the `AntaDevice` instance if populated.

    Parameters
    ----------
    command
        The command to collect.
    collection_id
        An identifier used to build the eAPI request ID.

    Raises
    ------
    RuntimeError
        If the eAPI client is closed. Call `refresh()` first to reconnect.
    """
    if self._client.is_closed:
        msg = f"Device {self.name}: httpx client is closed. Call refresh() to reconnect before collecting commands."
        raise RuntimeError(msg)
    async with self._command_semaphore:
        commands: list[EapiComplexCommand | EapiSimpleCommand] = []
        if self.enable and self._enable_password is not None:
            commands.append(
                {
                    "cmd": "enable",
                    "input": str(self._enable_password),
                },
            )
        elif self.enable:
            # No password
            commands.append(EapiComplexCommand(cmd="enable"))
        commands += [EapiComplexCommand(cmd=command.command, revision=command.revision)] if command.revision else [EapiComplexCommand(cmd=command.command)]
        try:
            response = await self._client.cli(
                commands=commands,
                ofmt=command.ofmt,
                version=command.version,
                req_id=f"ANTA-{collection_id}-{id(command)}" if collection_id else f"ANTA-{id(command)}",
            )
            # Do not keep response of 'enable' command
            command.output = response[-1]
        except asynceapi.EapiCommandError as e:
            # This block catches exceptions related to EOS issuing an error.
            self._handle_eapi_command_error(command, e)
        except EapiAuthenticationError as e:
            # This block catches authentication errors (HTTP 401) from eAPI when session auth is enabled.
            command.errors = [exc_to_str(e)]
            logger.error("Authentication failed while sending a command to %s: %s", self.name, e)
        except TimeoutException as e:
            # This block catches Timeout exceptions.
            command.errors = [exc_to_str(e)]
            timeouts = self._client.timeout.as_dict()
            logger.error(
                "%s occurred while sending a command to %s. Consider increasing the timeout.\nCurrent timeouts: Connect: %s | Read: %s | Write: %s | Pool: %s",
                exc_to_str(e),
                self.name,
                timeouts["connect"],
                timeouts["read"],
                timeouts["write"],
                timeouts["pool"],
            )
        except (ConnectError, OSError) as e:
            # This block catches OSError and socket issues related exceptions.
            command.errors = [exc_to_str(e)]
            self._handle_connect_error(e)
        except HTTPError as e:
            # This block catches most of the httpx Exceptions and logs a general message.
            command.errors = [exc_to_str(e)]
            anta_log_exception(e, f"An error occurred while issuing an eAPI request to {self.name}", logger)
        logger.debug("%s: %s", self.name, command)

async def copy(self, sources: list[Path], destination: Path, direction: Literal["to", "from"] = "from") -> None:
    """Copy files to and from the device using asyncssh.scp().

    The SSH connection is established transiently and closed safely via an `async with` context.

    Parameters
    ----------
    sources
        List of files to copy to or from the device.
    destination
        Local or remote destination when copying the files. Can be a folder.
    direction
        Defines if this coroutine copies files to or from the device.

    """
    async with asyncssh.connect(
        host=self._ssh_opts.host,
        port=self._ssh_opts.port,
        tunnel=self._ssh_opts.tunnel,
        family=self._ssh_opts.family,
        local_addr=self._ssh_opts.local_addr,
        options=self._ssh_opts,
    ) as conn:
        src: list[tuple[SSHClientConnection, Path]] | list[Path]
        dst: tuple[SSHClientConnection, Path] | Path
        if direction == "from":
            src = [(conn, file) for file in sources]
            dst = destination
            for file in sources:
                message = f"Copying '{file}' from device {self.name} to '{destination}' locally"
                logger.info(message)

        elif direction == "to":
            src = sources
            dst = conn, destination
            for file in src:
                message = f"Copying '{file}' to device {self.name} to '{destination}' remotely"
                logger.info(message)

        else:
            logger.critical("'direction' argument to copy() function is invalid: %s", direction)

            return
        await asyncssh.scp(src, dst)

async def disconnect(self) -> None:
    """Close the eAPI httpx client.

    Safe to call even if the client is already closed.
    Use `refresh()` to reconnect.
    """
    logger.debug("Disconnecting device %s", self.name)
    if not self._client.is_closed:
        await self._client.aclose()
    self.is_online = False
    self.established = False

async def refresh(self) -> None:
    """Update attributes of an AsyncEOSDevice instance.

    If the eAPI client has been closed (e.g. after a `disconnect()` call), it is
    automatically recreated before attempting to reach the device.

    Updates the following attributes:

    - `is_online`: True when the eAPI HTTP endpoint responds successfully.
    - `established`: True when a command execution succeeds.
    - `hw_model`: Hardware model parsed from `show version`.
    """
    logger.debug("Refreshing device %s", self.name)
    if self._client.is_closed:
        logger.debug("Recreating closed httpx client for device %s", self.name)
        self._client = self._create_client()
    try:
        self.is_online = await self._client.check_api_endpoint()
    except (EapiAuthenticationError, HTTPError) as e:
        self.is_online = False
        self.established = False
        logger.warning("An error occurred while attempting to connect to device %s: %s", self.name, exc_to_str(e))
        return

    show_version = AntaCommand(command="show version")
    await self._collect(show_version)
    if not show_version.collected:
        self.established = False
        logger.warning("Cannot get hardware information from device %s", self.name)
        return

    self.hw_model = show_version.json_output.get("modelName", None)
    if self.hw_model is None:
        self.established = False
        logger.critical("Cannot parse 'show version' returned by device %s", self.name)
    # in some cases it is possible that 'modelName' comes back empty
    elif self.hw_model == "":
        self.established = False
        logger.critical("Got an empty 'modelName' in the 'show version' returned by device %s", self.name)
    else:
        self.established = True