eos-downloader/eos_downloader/eos.py

#!/usr/bin/python
# coding: utf-8 -*-
# flake8: noqa: F811

"""
Specific EOS inheritance from object_download
"""

import os
import xml.etree.ElementTree as ET
from typing import List, Union

import rich
from loguru import logger
from rich import console

from eos_downloader.models.version import (
    BASE_BRANCH_STR,
    BASE_VERSION_STR,
    REGEX_EOS_VERSION,
    RTYPE_FEATURE,
    EosVersion,
)
from eos_downloader.object_downloader import ObjectDownloader

# logger = logging.getLogger(__name__)

console = rich.get_console()


class EOSDownloader(ObjectDownloader):
    """
    EOSDownloader Object to download EOS images from Arista.com website

    Supercharge ObjectDownloader to support EOS specific actions

    Parameters
    ----------
    ObjectDownloader : ObjectDownloader
        Base object
    """

    eos_versions: Union[List[EosVersion], None] = None

    @staticmethod
    def _disable_ztp(file_path: str) -> None:
        """
        _disable_ztp Method to disable ZTP in EOS image

        Create a file in the EOS image to disable ZTP process during initial boot

        Parameters
        ----------
        file_path : str
            Path where EOS image is located
        """
        logger.info("Mounting volume to disable ZTP")
        console.print("🚀 Mounting volume to disable ZTP")
        raw_folder = os.path.join(file_path, "raw")
        os.system(f"rm -rf {raw_folder}")
        os.system(f"mkdir -p {raw_folder}")
        os.system(
            f'guestmount -a {os.path.join(file_path, "hda.qcow2")} -m /dev/sda2 {os.path.join(file_path, "raw")}'
        )
        ztp_file = os.path.join(file_path, "raw/zerotouch-config")
        with open(ztp_file, "w", encoding="ascii") as zfile:
            zfile.write("DISABLE=True")
        logger.info(f"Unmounting volume in {file_path}")
        os.system(f"guestunmount {os.path.join(file_path, 'raw')}")
        os.system(f"rm -rf {os.path.join(file_path, 'raw')}")
        logger.info(f"Volume has been successfully unmounted at {file_path}")

    def _parse_xml_for_version(
        self,
        root_xml: ET.ElementTree,
        xpath: str = './/dir[@label="Active Releases"]/dir/dir/[@label]',
    ) -> List[EosVersion]:
        """
        Extract list of available EOS versions from Arista.com website

        Create a list of EosVersion object for all versions available on Arista.com

        Args:
            root_xml (ET.ElementTree): XML file with all versions available
            xpath (str, optional): XPATH to use to extract EOS version. Defaults to './/dir[@label="Active Releases"]/dir/dir/[@label]'.

        Returns:
            List[EosVersion]: List of EosVersion representing all available EOS versions
        """
        # XPATH: .//dir[@label="Active Releases"]/dir/dir/[@label]
        if self.eos_versions is None:
            logger.debug(f"Using xpath {xpath}")
            eos_versions = []
            for node in root_xml.findall(xpath):
                if "label" in node.attrib and node.get("label") is not None:
                    label = node.get("label")
                    if label is not None and REGEX_EOS_VERSION.match(label):
                        eos_version = EosVersion.from_str(label)
                        eos_versions.append(eos_version)
                        logger.debug(f"Found {label} - {eos_version}")
            logger.debug(f"List of versions found on arista.com is: {eos_versions}")
            self.eos_versions = eos_versions
        else:
            logger.debug(
                "receiving instruction to download versions, but already available"
            )
        return self.eos_versions

    def _get_branches(self, with_rtype: str = RTYPE_FEATURE) -> List[str]:
        """
        Extract all EOS branches available from arista.com

        Call self._parse_xml_for_version and then build list of available branches

        Args:
            rtype (str, optional): Release type to find. Can be M or F, default to F

        Returns:
            List[str]: A lsit of string that represent all availables EOS branches
        """
        root = self.get_folder_tree()
        versions = self._parse_xml_for_version(root_xml=root)
        return list(
            {version.branch for version in versions if version.rtype == with_rtype}
        )

    def latest_branch(self, rtype: str = RTYPE_FEATURE) -> EosVersion:
        """
        Get latest branch from semver standpoint

        Args:
            rtype (str, optional): Release type to find. Can be M or F, default to F

        Returns:
            EosVersion: Latest Branch object
        """
        selected_branch = EosVersion.from_str(BASE_BRANCH_STR)
        for branch in self._get_branches(with_rtype=rtype):
            branch = EosVersion.from_str(branch)
            selected_branch = max(selected_branch, branch)
        return selected_branch

    def get_eos_versions(
        self, branch: Union[str, None] = None, rtype: Union[str, None] = None
    ) -> List[EosVersion]:
        """
        Get a list of available EOS version available on arista.com

        If a branch is provided, only version in this branch are listed.
        Otherwise, all versions are provided.

        Args:
            branch (str, optional): An EOS branch to filter. Defaults to None.
            rtype (str, optional): Release type to find. Can be M or F, default to F

        Returns:
            List[EosVersion]: A list of versions available
        """
        root = self.get_folder_tree()
        result = []
        for version in self._parse_xml_for_version(root_xml=root):
            if branch is None and (version.rtype == rtype or rtype is None):
                result.append(version)
            elif (
                branch is not None
                and version.is_in_branch(branch)
                and version.rtype == rtype
            ):
                result.append(version)
        return result

    def latest_eos(
        self, branch: Union[str, None] = None, rtype: str = RTYPE_FEATURE
    ) -> EosVersion:
        """
        Get latest version of EOS

        If a branch is provided, only version in this branch are listed.
        Otherwise, all versions are provided.
        You can select what type of version to consider: M or F

        Args:
            branch (str, optional): An EOS branch to filter. Defaults to None.
            rtype (str, optional): An EOS version type to filter, Can be M or F. Defaults to None.

        Returns:
            EosVersion: latest version selected
        """
        selected_version = EosVersion.from_str(BASE_VERSION_STR)
        if branch is None:
            latest_branch = self.latest_branch(rtype=rtype)
        else:
            latest_branch = EosVersion.from_str(branch)
        for version in self.get_eos_versions(
            branch=str(latest_branch.branch), rtype=rtype
        ):
            if version > selected_version:
                if rtype is not None and version.rtype == rtype:
                    selected_version = version
                if rtype is None:
                    selected_version = version
        return selected_version