#!/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 asyncio import re from typing import ( List, TYPE_CHECKING, Any, AnyStr, Dict, NamedTuple, Optional, Pattern, Tuple, Union, ) from fbnet.command_runner.exceptions import ( PermissionErrorException, RuntimeErrorException, ) from .command_session import SSHCommandSession from .device_info import IPInfo from .options import Option if TYPE_CHECKING: from asyncio import AbstractEventLoop from fbnet.command_runner.service import FcrServiceBase from .command_session import ResponseMatch from .device_info import DeviceInfo class ConsoleInfo(NamedTuple): """ Information about the console """ contype: str host: str server: str port: Union[str, int] def __repr__(self) -> str: """ pretty representation of console information """ return "host:{s.host} {s.contype}: {s.server}:{s.port}".format(s=self) class ConsoleCommandSession(SSHCommandSession): """ A command session that runs over a console connections. Currently we only support SSH connection to the console server """ # The three default class variables below let you preset the prompts in-case vendor is unknown. # If you would like to hard-code or preset the prompts for specific devices, you can do so # by hardcoding _CONFIG_CONSOLE_PROMPTS_RE_DICT and _CONFIG_INTERACT_PROMPTS_RE_DICT _DEFAULT_INTERACT_PROMPTS: Dict[bytes, bytes] = { b"Y": rb"Do you acknowledge\? \(Y/N\)\?" } _DEFAULT_CONSOLE_PROMPTS: Dict[bytes, bytes] = { # For login we need to ignore output like: # Last login: Mon May 8 13:53:17 on ttyS0 b"login": rb".*((?%](?!Login incorrect)", b"interact_prompts": rb"Do you acknowledge\? \(Y/N\)\?", # Ignore these prompts during login attempts b"ignore_prompts": b"( to cli \\])|(who is on this device.\\]\\r\\n)|(Press RETURN to get started\r\n)", } _DEFAULT_CONSOLE_PROMPTS_RE: Optional[Pattern] = None _DEFAULT_LOGOUT_COMMAND: bytes = b"exit" _CONSOLE_EXPECT_DELAY: int = 5 _CONSOLE_LOGIN_TIMEOUT_S = Option( "--console_login_timeout_s", help="The upper bound of the time (in seconds) that FCR waits for a " "console server to login to the target device (only applies when a " "console is used to connect to a device). (default: %(default)s)", type=int, default=60, ) _CONFIG_CONSOLE_PROMPTS_RE_DICT: Optional[Dict[bytes, Pattern]] = None _CONFIG_INTERACT_PROMPTS_RE_DICT: Optional[Dict[bytes, Dict[bytes, bytes]]] = None def __init__( self, service: "FcrServiceBase", devinfo: "DeviceInfo", options: Dict[str, Any], loop: "AbstractEventLoop", ) -> None: super().__init__(service, devinfo, options, loop) self._console: str = options["console"] @classmethod def _build_and_set_prompts_re_dict( cls, vendor_prompts: Optional[Dict[bytes, Dict[bytes, bytes]]] = None, vendor_interact_prompts: Optional[Dict[bytes, Dict[bytes, bytes]]] = None, ) -> None: """ This method takes in a dictionary of a dictionary of vendor specific prompts, builds each set of vendor prompts and adds them to the instance dictionary of regexs CONFIG_CONSOLE_PROMPTS_RE_DICT to be matched against console prompts. Also sets the interact_prompts dictionary, if applicable (unneccessary to compile regexes until there is a match). """ cls._CONFIG_CONSOLE_PROMPTS_RE_DICT = {} cls._CONFIG_INTERACT_PROMPTS_RE_DICT = {} if vendor_prompts: for vendor, prompts in vendor_prompts.items(): cls._CONFIG_CONSOLE_PROMPTS_RE_DICT[ # pyre-ignore vendor ] = cls._build_individual_prompt_re(prompts) if vendor_interact_prompts: cls._CONFIG_INTERACT_PROMPTS_RE_DICT = vendor_interact_prompts @classmethod def _build_individual_prompt_re( cls, prompts: Optional[Dict[bytes, bytes]] = None ) -> Pattern: """ This function takes in a dictionary of regexes of a single vendor to match prompts against, then compiles all the prompts into a single grouped regex and returns the regex. """ prompts_re = [] if prompts: prompts_re = [ b"(?P<%s>%s)" % (group, regex) for group, regex in prompts.items() ] prompt_re = b"|".join(prompts_re) return re.compile(prompt_re + rb"\s*$") @classmethod def get_default_console_prompt_re(cls) -> Pattern: # This check ensures that _DEFAULT_CONSOLE_PROMPTS_RE is not None, ignoring pyre warning if not cls._DEFAULT_CONSOLE_PROMPTS_RE: cls._DEFAULT_CONSOLE_PROMPTS_RE = cls._build_individual_prompt_re( cls._DEFAULT_CONSOLE_PROMPTS ) return cls._DEFAULT_CONSOLE_PROMPTS_RE # pyre-ignore @classmethod def get_prompt_re(cls, vendor_name: Optional[str] = None) -> Pattern: """ This method takes in a vendor name and retrieves the pre-compiled group regex corresponding to that vendor from the _CONFIG_CONSOLE_PROMPTS_RE_DICT. If the vendor is not given or does not already exist in the dictionary, return the default (hard-coded) regex. If the default regex is not already built, build it and then set it. """ # This check ensure that _CONFIG_CONSOLE_PROMPTS_RE_DICT is not None, ignoring pyre warning if cls._CONFIG_CONSOLE_PROMPTS_RE_DICT and vendor_name: return cls._CONFIG_CONSOLE_PROMPTS_RE_DICT.get( # pyre-ignore vendor_name.encode("utf-8"), cls.get_default_console_prompt_re() ) return cls.get_default_console_prompt_re() async def dest_info(self) -> Tuple[List[IPInfo], int, str, str]: console = await self.get_console_info() self.logger.info(f"{str(console)}") # By default we assume a console is directly specified by the user, # so we don't want to raise error messages that it is not pingable # Set is_pingable to True since check_ip method always returns True is_pingable = True return ( [IPInfo(console.server, is_pingable)], console.port, self._username, self._password, ) async def expect( self, regex: Pattern, timeout: int = _CONSOLE_EXPECT_DELAY ) -> "ResponseMatch": try: return await asyncio.wait_for( self.wait_prompt(prompt_re=regex, timeout=timeout), timeout, loop=self._loop, ) except asyncio.TimeoutError as ex: data = [] # This statement ensures that _stream_reader is not none if self._stream_reader: data = await self._stream_reader.drain() raise asyncio.TimeoutError( "Timeout during waiting for prompt." f"Currently received data: {data[-200:]}" ) from ex def send(self, data: AnyStr, end: bytes = b"\n") -> None: """ send some data and optionally wait for some data """ if isinstance(data, str): data = data.encode("utf8") # This check ensure _stream_writer is not none, ignoring the pyre warning if self._stream_writer: self._stream_writer.write(data + end) # pyre-ignore async def _try_login( # noqa C901 self, username: Optional[str] = None, passwd: Optional[str] = None, kickstart: bool = False, username_tried: bool = False, pwd_tried: bool = False, get_response_timeout: int = _CONSOLE_EXPECT_DELAY, ) -> None: """ A helper function that tries to login into the device. kickstart gives the option to send a clear line to the console before entering the username and password to prevent false TimeoutError, since sometimes user needs to hit an Enter before logging in. username_tried and pwd_tried are an indicator showing that the FCR has received a login error from the console after sending the username and password to the console, if this triggers, we would raise a PermissionError stating that the console has failed to login, this will also prevent false TimeoutError to happen. """ try: res = await self._get_response(timeout=get_response_timeout) except asyncio.TimeoutError: if kickstart and username: # We likely didn't get anything from the console. Try sending a # newline to kickstart the login process self.logger.debug("kickstart console login") # Clear the current line and send a newline self._send_clearline() return await self._try_login( username=username, passwd=passwd, username_tried=username_tried, pwd_tried=pwd_tried, ) else: raise if res.groupdict.get("ignore_prompts"): # If we match anything in the ignore prompts, set a \r\n self._send_newline(end=b"") await asyncio.sleep(0.2) # Let the console catch up # Now again try to login. return await self._try_login( username=username, passwd=passwd, username_tried=username_tried, pwd_tried=pwd_tried, ) elif res.groupdict.get("login"): if username_tried: raise PermissionErrorException( "Login failure, possibly incorrect username or password, " "or device refuses to login." ) # The device is requesting login information # If we don't have a username, then likely we already sent a # username. The consoles are slow, we may have send extra # carriage returns, resulting in multiple login prompts. We will # simply ignore the subsequent login prompts. if username is not None: # TODO: It seems the original author of this session used the # arguments `username` and `passwd` to indicate whether the # username or password have been sent to the device (it will # set the corresponding argument to None after sending the # credential). This seems to be duplicated with the arguments # `username_tried` and `pwd_tried` we added later. If so, we'll # need to remove one of them self.send(self._username) # if we don't have username, we are likely waiting for password return await self._try_login( passwd=passwd, username_tried=True, pwd_tried=pwd_tried ) elif res.groupdict.get("passwd"): if pwd_tried: raise PermissionErrorException( "Login failure, possibly incorrect username or password, " "or device refuses to login." ) if passwd is None: # passwd information not available # Likely we have alreay sent the password. Bail out instead # of getting stuck in a loop. raise RuntimeErrorException("Failed to login: Missing password") self.send(self._password) return await self._try_login( username_tried=username_tried, pwd_tried=True, get_response_timeout=self._CONSOLE_LOGIN_TIMEOUT_S, ) elif res.groupdict.get("interact_prompts"): # E.g. send 'Y' to to satisfy the post login prompt on Nokia device self._interact_prompts_action(res.groupdict.get("interact_prompts")) return await self._try_login( username_tried=username_tried, pwd_tried=pwd_tried ) elif res.groupdict.get("prompt"): # Finally we matched a prompt. we are done return self._send_newline() else: raise RuntimeErrorException("Matched no group: %s" % (res.groupdict)) async def _get_response(self, timeout: int) -> "ResponseMatch": # A small delay to avoid having to match extraneous input await asyncio.sleep(0.1) res = await self.expect( self.get_prompt_re(self._devinfo.vendor_name), timeout=timeout ) return res async def _try_logout(self, kick_shutdown: bool = False) -> None: """ Run logout command and wait for the login prompt to show up (the login prompt indicates that it successfully logs out and is waiting for the next login). This is to ensure we cleanly diconnect from the device after running the command """ if not self._stream_writer: return logout_cmd = ( self._devinfo.vendor_data.exit_command or self._DEFAULT_LOGOUT_COMMAND ) self.logger.info(f"Logout from device, running logout command: {logout_cmd}") self._stream_writer.write(logout_cmd + b"\n") # Make sure we logout of the system while True: try: res = await self._get_response(timeout=self._CONSOLE_EXPECT_DELAY) except asyncio.TimeoutError: if kick_shutdown: self.logger.info( "Received first timeout while matching console prompt, " "sending a new line character to the console" ) self._send_newline() return await self._try_logout() else: self.logger.exception("Console session timed out while logging out") return except Exception: self.logger.exception("Console session log out failure") return if res.groupdict.get("ignore_prompts"): # If we match anything in the ignore prompts, set a \r\n self._send_newline(end=b"") await asyncio.sleep(0.2) # Let the console catch up elif res.groupdict.get("login"): self.logger.info("Logout successfully") return else: self.logger.error(f"Get unexpected prompt when logging out: {res}") return async def _close(self) -> None: await self._try_logout(kick_shutdown=True) await super()._close() def _send_clearline(self) -> None: self.send(b"\x15\r\n") def _send_newline(self, end: bytes = b"\n") -> None: self.send(b"\r", end) def _interact_prompts_action(self, prompt_match: AnyStr) -> None: # Check if to use default interactive prompts or configured prompts vendor_interact_prompts = self._DEFAULT_CONSOLE_PROMPTS if ( self._CONFIG_INTERACT_PROMPTS_RE_DICT and self._devinfo.vendor_name.encode("utf8") in self._CONFIG_INTERACT_PROMPTS_RE_DICT ): vendor_interact_prompts = self._CONFIG_INTERACT_PROMPTS_RE_DICT.get( self._devinfo.vendor_name.encode("utf8") ) interact_prompts = [ b"(?P<%s>%s)" % (group, regex) # pyre-fixme[16]: `Optional` has no attribute `items`. for group, regex in vendor_interact_prompts.items() ] interact_prompts_re = b"|".join(interact_prompts) interact_prompt_match = re.match( interact_prompts_re, prompt_match # pyre-ignore ) # Ignoring the pyre warning since the logic in _try_login ensures we # find a match for action in interact_prompt_match.groupdict(): # pyre-ignore self.send(action) async def _setup_connection(self) -> None: if self._opts.get("raw_session"): await asyncio.sleep(1) else: # Since this is a normal session, try to login to device. await self._try_login(self._username, self._password, kickstart=True) # Now send the setup commands await super()._setup_connection() async def get_console_info(self) -> ConsoleInfo: """ By default we assume a console is directly specified by the user. Depending on your system, you may want to get this information from your local database. In such case you can override this method according to your needs """ con_srv, con_port = self._console.split(":") return ConsoleInfo("CON", self.hostname, con_srv, con_port) async def _run_command( self, command: bytes, timeout: Optional[int] = None, prompt_re: Optional[Pattern] = None, ) -> bytes: if self._opts.get("raw_session"): # This statement ensures _stream_reader is not none if self._stream_reader: await self._stream_reader.drain() self.send(command) resp = await self.wait_prompt(prompt_re) return resp.data + resp.matched else: return await super()._run_command(command, timeout, prompt_re)