device.py

import glob
import os
import importlib
import logging

from utils.info_packet import info_packet
from utils.netconf_utils import nc_config, nc_get_schema, nc_get_subscriptions
from utils.yanglint import extract_yang_with_lxml

def get_available_vendors(vendors_dir="./vendors"):
    vendor_files = glob.glob(os.path.join(vendors_dir, "vendor_*.py"))
    return sorted(
        [
            os.path.basename(f)[7:-3]
            for f in vendor_files
            if not f.endswith("__init__.py") and not f.endswith("base.py")
        ]
    )

def load_vendor_class(vendor_name, xml_dir, config):
    vendor_mod = f"vendors.vendor_{vendor_name.lower()}"
    mod = importlib.import_module(vendor_mod)
    class_name = f"Vendor{vendor_name.capitalize()}"
    return getattr(mod, class_name)(xml_dir, config)

def get_schema_to_file(conn, module_name, results_dir):
    xml = nc_get_schema(conn, module_name)
    module_file = os.path.join(results_dir, f"{module_name}.yang")
    yang_schema = extract_yang_with_lxml(xml)
    if yang_schema:
        with open(module_file, "w") as f:
            f.write(yang_schema)
        print(f"YANG schema for {module_name} saved to {module_file}.")
        logging.info(f"YANG schema for {module_name} saved to {module_file}.")
    else:
        logging.error(f"Failed to extract YANG schema for {module_name}.")

def restore_device_yp_config(conn, vendor, backup_file=None, results_dir=None):
    info_packet("--- Restore YANG Push configuration from backup ---")
    if backup_file is None and results_dir is not None:
        backup_file = os.path.join(results_dir, "yp_config_backup.xml")
    elif backup_file is not None and results_dir is None:
        backup_file = backup_file
    else:
        logging.error("No backup file specified. Cannot restore configuration.")
        return
    if not os.path.exists(backup_file):
        logging.error("Backup file not found. Cannot restore configuration.")
        return

    with open(backup_file, "r") as f:
        xml_content = f.read()

    try:
        nc_config(conn, xml_content)
        logging.info("YANG Push configuration restored successfully.")
    except Exception as e:
        logging.error(f"Failed to restore YANG Push configuration: {e}")

def backup_device_yp_config(conn, vendor, backup_file=None, results_dir=None):
    info_packet("--- Backup current YANG Push configuration ---")
    if backup_file is None and results_dir is not None:
        backup_file = os.path.join(results_dir, "yp_config_backup.xml")
    elif backup_file is not None and results_dir is None:
        backup_file = backup_file
    else:
        logging.error("No backup file specified. Cannot save configuration.")
        return
    try:
        result = nc_get_subscriptions(conn, vendor)
    except Exception as e:
        logging.error(f"Failed to retrieve YANG Push configuration: {e}")
        return

    if result:
        with open(backup_file, "w") as f:
            f.write(result)
        print(f"YANG Push configuration backup saved to {backup_file}.")
        logging.info(f"YANG Push configuration backup saved to {backup_file}.")
    else:
        logging.error("Failed to retrieve YANG Push configuration.")