Using Scrapli as PyATS Connection Plugin

genie_scrapli.py

from pyats.connections import BaseConnection
from pyats.utils.secret_strings import to_plaintext
from scrapli import Scrapli as Scrapli_connector


class Scrapli(BaseConnection):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._is_connected = False
        platform_translation = {
            'comware': "hp_comware",
            'iosxe': 'cisco_iosxe',
            'nxos': 'cisco_nxos'
            'junos': 'juniper_junos'
        }
        if self.device.os in platform_translation:
            platform = platform_translation[self.device.os]
        else:
            platform = self.device.os
        
        device = {
           "host": self.connection_info.ip.exploded,
           "auth_username": self.connection_info.get('credentials')['default']['username'],
           "auth_password": to_plaintext(self.connection_info.get('credentials')['default']['password']),
           "auth_strict_key": False,
           "platform": platform,
        }
        
        if self.connection_info.get('protocol') is not None:
            device['transport'] = self.connection_info.get('protocol')
            
        self.conn = Scrapli_connector(**device)    

    @property
    def connected(self):
        return self._is_connected

    def connect(self, *args, **kwargs):
        self.conn.open()
        self._is_connected = True
        return self        

    def disconnect(self):
        if self.conn.isalive():
            self.conn.close()
        self._is_connected = False

    def execute(self, cmd):
        out = self.conn.send_command(cmd)
        return out.result
         
    def configure(self, cmd):
        out = self.conn.send_configs(cmd)
        return out.result

testbed.yaml

devices:
        SW_1:
                os: comware
                type: 'Switch'
                credentials:
                        default:
                                username: 'admin'
                                password: 'admin'
                connections:
                        a:
                                class: genie_scrapli.Scrapli
                                protocol: <transport>
                                ip: <ip_address>