#!/usr/bin/env python3 # -*- coding: utf-8 -*- # Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import json import os import re from fbnet.command_runner_asyncio.CommandRunner.ttypes import SessionType from . import utils from .base_service import ServiceObj from .command_session import SSHCommandSession from .options import Option from .ssh_netconf import SSHNetconf class VendorConfig: def __init__(self, defaults, session_names): self._cfg = {} self._session_names = session_names self.update(defaults) def __getattr__(self, attr): return self._cfg.get(attr) def update(self, cfg): for prop, val in cfg.items(): self._cfg[prop] = utils.canonicalize(val) if "supported_sessions" in cfg: # Refresh supported_sessions only if it's updated self._cfg["supported_sessions"] = { self._session_names[s] for s in self._cfg["supported_sessions"] } if "session_type" in cfg: self._cfg["session_type"] = self._session_names[self._cfg["session_type"]] # Default session type should be supported self._cfg["supported_sessions"].add(self._cfg["session_type"]) class DeviceVendor(ServiceObj): _DEFAULTS = { "cli_setup": [b"term len 0", b"term width 511"], "prompt_regex": [rb"[\w.]+[>#$]"], "cmd_timeout_sec": 30, "clear_command": b"\x15", "session_type": b"ssh", "supported_sessions": {b"ssh", b"netconf"}, "autocomplete": True, "port": 22, } _PROMPTS_RE = re.compile( # pyre-fixme[16]: `int` has no attribute `__iter__`. b"|".join([b"(%s)" % p for p in _DEFAULTS["prompt_regex"]]), re.M, ) _SESSION_NAMES = {b"ssh": SessionType.SSH, b"netconf": SessionType.SSH_NETCONF} _SESSION_TYPES = { SessionType.SSH: SSHCommandSession, SessionType.SSH_NETCONF: SSHNetconf, } def __init__(self, vendor_name, service): super().__init__(service, "DeviceVendor") self._vendor_name = vendor_name self._config = VendorConfig(self._DEFAULTS, self._SESSION_NAMES) self._prompt_re = self._PROMPTS_RE def __repr__(self): props = { "cli_setup": self._config.cli_setup, "prompt_regex": self._config.prompt_regex, "cmd_timeout_sec": self._config.cmd_timeout_sec, "autocomplete": self._config.autocomplete, } return "DeviceVendor(%s) %s" % (self.vendor_name, props) @classmethod def register_counters(cls, stats_mgr): for session_type in cls._SESSION_TYPES.values(): session_type.register_counters(stats_mgr) stats_mgr.register_counter("device_vendor.all_sessions") stats_mgr.register_counter("device_vendor.unsupported_session") def get_prompt_re(self, trailer=None): """ Get prompt regex for the device. Optionally a trailer can be specified. This is extra text expected after the prompt. Mostly useful for interactive command. E.g. when we get a list of completion, the intial command is inserted after the prompt """ if not trailer: return self._prompt_re return self._get_prompt_re(trailer) def get_port(self): return self._config.port @property def vendor_name(self): return self._vendor_name @property def cmd_timeout_sec(self): return self._config.cmd_timeout_sec @property def clear_command(self): return self._config.clear_command @property def exit_command(self): return self._config.exit_command @property def cli_setup(self): return self._config.cli_setup @property def session_type(self): return self._SESSION_TYPES[self._config.session_type] @property def autocomplete(self): return self._config.autocomplete def select_session_type(self, options): """ Select session type for given set of options. Users can override session type here, by specifying session_type in options. This needs to be implemented for vendors supporting multiple session types """ self.inc_counter("device_vendor.all_sessions") session_type = options.get("session_type", None) if session_type in self._config.supported_sessions: return self._SESSION_TYPES.get(session_type, self.session_type) else: if session_type is not None: self.logger.warning( "Device vendor {} does not support session {}".format( self._vendor_name, session_type ) ) self.inc_counter("device_vendor.unsupported_session") return self.session_type def update_config(self, vendor_config): self._config.update(vendor_config) self._update_prompts_re() def set_user_prompts(self, prompts): self._config.update({"user_prompts": prompts}) self._update_prompts_re() def _update_prompts_re(self): self._prompt_re = self._get_prompt_re() def _get_prompt_re(self, trailer=None): prompts = self._config.prompt_regex if self._config.shell_prompts: prompts += self._config.shell_prompts if self._config.user_prompts: prompts += self._config.user_prompts if self._config.bootstrap_prompts: prompts += self._config.bootstrap_prompts return self._build_prompt_re(prompts, trailer) @classmethod def _build_prompt_re(cls, prompts, trailer=None): all_prompts = (b"(%s)" % prompt for prompt in prompts) trailer = trailer or b"" # the prompt must be at the start of the line. # Also since we are sending one command at a time, it must also be the # last text in the text. Although still not perfect, this greatly # reduces the probability of this matching some random text in the # output. Not that we are matching at end of the text, not at the end of # each line in text (re.M is not specified) return re.compile( b"(?<=[\n\r])(?P" + b"|".join(all_prompts) + rb")\s*" + trailer + b"$", re.M, ) class DeviceVendors(ServiceObj): # User specified device vendor information device_vendors = Option( "--device_vendors", help="A JSON file containing vendor information", default=None, ) def __init__(self, service, name=None): super().__init__(service, name) self._vendors = {} self._load_vendors_data() @classmethod def register_counters(cls, stats_mgr): DeviceVendor.register_counters(stats_mgr) def get(self, name): return self._vendors.get(name) or self._createVendor(name) def _update_user_prompts(self, path, cfg): if cfg is not None: for vendor, prompts in cfg["prompt_regexs"].items(): self.get(vendor).set_user_prompts(prompts) def _update_device_vendors(self, path, cfg): # now load the vendor information for name, props in cfg["vendor_config"].items(): vendor = self.get(name) vendor.update_config(props) def load_vendors(self, path, json_str): cfg = json.loads(json_str) return self._update_device_vendors(path, cfg) def _load_device_vendors(self): """ Load device vendors specified on command line """ if self.device_vendors and os.path.exists(self.device_vendors): self.logger.info("loading local file") with open(self.device_vendors, "rb") as fh: jsonb = fh.read() return self.load_vendors(self.device_vendors, jsonb.decode("utf-8")) def _load_vendors_data(self): """ Load vendors information """ self._load_device_vendors() def _createVendor(self, name): vendor = DeviceVendor(name, self.service) self._vendors[name] = vendor return vendor