lookingglass/hyperglass/execution/drivers/_common.py
2021-09-12 15:06:34 -07:00

93 lines
3.2 KiB
Python

"""Base Connection Class."""
# Standard Library
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Dict, Union, Sequence
# Project
from hyperglass.log import log
from hyperglass.plugins import OutputPluginManager
from hyperglass.models.api import Query
from hyperglass.parsing.nos import scrape_parsers, structured_parsers
from hyperglass.parsing.common import parsers
from hyperglass.models.config.devices import Device
# Local
from ._construct import Construct
if TYPE_CHECKING:
# Project
from hyperglass.compat._sshtunnel import SSHTunnelForwarder
class Connection(ABC):
"""Base transport driver class."""
def __init__(self, device: Device, query_data: Query) -> None:
"""Initialize connection to device."""
self.device = device
self.query_data = query_data
self.query_type = self.query_data.query_type
self.query_target = self.query_data.query_target
self._query = Construct(device=self.device, query=self.query_data)
self.query = self._query.queries()
self.plugin_manager = OutputPluginManager()
@abstractmethod
def setup_proxy(self: "Connection") -> "SSHTunnelForwarder":
"""Return a preconfigured sshtunnel.SSHTunnelForwarder instance."""
pass
async def parsed_response( # noqa: C901 ("too complex")
self, output: Sequence[str]
) -> Union[str, Sequence[Dict]]:
"""Send output through common parsers."""
log.debug("Pre-parsed responses:\n{}", output)
parsed = ()
response = None
structured_nos = structured_parsers.keys()
structured_query_types = structured_parsers.get(self.device.nos, {}).keys()
scrape_nos = scrape_parsers.keys()
scrape_query_types = scrape_parsers.get(self.device.nos, {}).keys()
if not self.device.structured_output:
_parsed = ()
for func in parsers:
for response in output:
_output = func(commands=self.query, output=response)
_parsed += (_output,)
if self.device.nos in scrape_nos and self.query_type in scrape_query_types:
func = scrape_parsers[self.device.nos][self.query_type]
for response in _parsed:
_output = func(response)
parsed += (_output,)
else:
parsed += _parsed
response = "\n\n".join(parsed)
elif (
self.device.structured_output
and self.device.nos in structured_nos
and self.query_type not in structured_query_types
):
for func in parsers:
for response in output:
_output = func(commands=self.query, output=response)
parsed += (_output,)
response = "\n\n".join(parsed)
elif (
self.device.structured_output
and self.device.nos in structured_nos
and self.query_type in structured_query_types
):
func = structured_parsers[self.device.nos][self.query_type]
response = func(output)
if response is None:
response = "\n\n".join(output)
log.debug("Post-parsed responses:\n{}", response)
return response