Vty.py

import json

import asyncssh

from CapybaraNetty import CapybaraNetty


class Vty(CapybaraNetty):
    def __init__(self, host_ip, host_user):
        super().__init__()
        self.host_ip = host_ip
        self.host_user = host_user

    def ssh_is_alive(self, **kwargs):
        return super(Vty, self).ssh_is_alive(self.host_ip, self.host_user)

    def build_vty_exec(self, commands: list, echo_commands=False):
        additional_vtysh_flags = ""
        if echo_commands:
            additional_vtysh_flags += " -E "
        base = f"sudo vtysh {additional_vtysh_flags} -c 'configure term'"
        exec_commands = []
        for command in commands:
            self.__logger.info(f"{self.log_format_ip(self.host_ip)} CMD: {command}")
            exec_commands.append(f" -c '{command}' ")
        return base + "".join(exec_commands)

    async def get_routes(self):
        route_json, stderr = await self.ssh_exec(
            self.build_vty_exec(["do show ip route json"]),
        )
        # It returns list of strings which combined makes up JSON - need to join it up
        routes = json.loads("".join(route_json))
        # Routes now has dict of lists, where d[prefix][0] is one of the possible routes for prefix, change this so
        # it is easier to process
        # Now it will be list of dicts, do note that this means that there cannot be two routes for same prefix,
        # but for PoC this is fine
        routes_converted = [dict(zip(routes, col)) for col in zip(*routes.values())][0]
        routes_converted = list(routes_converted.values())
        return routes_converted

    # Add a route - make note of generate_note_cmd signature - some attributes are location sensitive
    # Make sure commands are batched up otherwise new routes will be slow to apply since 1 command = 1 ssh connection
    async def add_routes(self, routes: list):
        commands = []
        if len(routes) > 0:
            for route in routes:
                commands.append(self.generate_route_config_cmd(**route))
            await self.ssh_exec(
                self.build_vty_exec(commands),
            )

    async def del_routes(self, routes: list):
        commands = []
        if len(routes) > 0:
            for route in routes:
                commands.append(f"no {self.generate_route_config_cmd(**route)} tag 2")
            return await self.ssh_exec(
                self.build_vty_exec(commands, echo_commands=True),
            )
        return None, None

    async def show_route(self, prefix):
        stdout, stderr = await self.ssh_exec(self.build_vty_exec([f"do show ip route {prefix}"]))
        return stdout

    def generate_route_config_cmd(self, prefix, gateway, interface, distance=1, **kwargs):
        additional_cmd_args = ""
        for key, value in kwargs.items():
            additional_cmd_args += f"{key} {value} "
        return f"ip route {prefix} {gateway} {interface} {distance}  {additional_cmd_args}"

    # This includes information from the configuration itself - useful for removing routes
    # Provide arguments to get route CMD which match the provided arguments
    async def get_config_routes(self, prefix=None, gateway=None, interface=None, **kwargs):
        def filter(route: str):
            route = route[9:]  # remove `ip route ` part from cmd
            if len(route) == 0:
                return False
            if prefix and route.split()[0] != prefix:
                return False
            if gateway and route.split()[1] != gateway:
                return False
            if interface and interface not in route.split()[2:]:
                return False
            for key, value in kwargs.items():
                if f"{key} {value}" not in route:
                    return False
            return True

        routes_config, stderr = await self.ssh_exec(
            self.build_vty_exec(["do show running-config | include ip route"]),
        )
        routes = []
        for route in routes_config:
            # Remove newlines at end
            route = route.rstrip()
            if filter(route):
                routes.append(route)
        return routes

    async def del_config_routes(self, route_cmds):
        no_route_cmds = []
        for route in route_cmds:
            no_route_cmds.append(f"no {route}")
        await self.ssh_exec(
            self.build_vty_exec(no_route_cmds, echo_commands=True),
        )

    async def get_interfaces(self):
        interface_json, stderr = await self.ssh_exec(
            self.build_vty_exec(["do show interface json"]),
        )
        # It returns list of strings which combined makes up JSON - need to join it up
        interfaces = json.loads("".join(interface_json))
        return interfaces

    async def ssh_exec(self, command, **kwargs):
        return await super(Vty, self).ssh_exec(self.host_ip, self.host_user, command=command)