async def wrapper(self, *args, **kwargs): for i, arg in enumerate(args): if arg is None: raise ttypes.SessionException( message=f"The ({i + 1})th argument cannot be None.") if isinstance(arg, ttypes.Device): _check_device(arg) if isinstance(arg, ttypes.Session): _check_session(arg) for key in kwargs: if key == "command" and not kwargs[key]: raise ttypes.SessionException( message="Required argument (command) cannot be None.") if key == "device": _check_device(kwargs[key]) if key == "session": _check_session(kwargs[key]) if key == "device_to_commands" or key == "device_to_configlets": if not kwargs[key]: raise ttypes.SessionException( message=f"Required argument ({key}) cannot be None.") for device in kwargs[key]: _check_device(device) return await fn(self, *args, **kwargs)
def _check_device(arg): if not arg: raise ttypes.SessionException( message="Required argument (device) cannot be None." ) if not arg.hostname: raise ttypes.SessionException(message="device.hostname cannot be empty.")
def _check_session(arg): if not arg: raise ttypes.SessionException( message="Required argument (session) cannot be None." ) if not arg.hostname: raise ttypes.SessionException(message="session.hostname cannot be empty.") if not arg.id: raise ttypes.SessionException(message="session.id cannot be empty.") if not arg.name: raise ttypes.SessionException(message="session.name cannot be empty.")
async def close_session(self, tsession, client_ip, client_port): try: session = CommandSession.get(tsession.id, client_ip, client_port) await session.close() except Exception as e: raise ttypes.SessionException(message="close_session failed: %r" % (e)) from e
async def _open_session( self, device, open_timeout, idle_timeout, client_ip, client_port, uuid, raw_session=False, ): options = self._get_options( device, client_ip, client_port, open_timeout, idle_timeout, raw_session=raw_session, ) try: devinfo = await self._lookup_device(device) session = await devinfo.setup_session(self.service, device, options, loop=self.loop) return ttypes.Session(id=session.id, name=session.hostname, hostname=device.hostname) except Exception as e: raise ttypes.SessionException(message="open_session failed: %r" % e) from e
async def _remote_task(chunk): # Run the chunk of commands on remote instance self.incrementCounter("bulk_run.remote") retry_count = 0 while True: try: return await self._bulk_run_remote(chunk, timeout, open_timeout, client_ip, client_port, uuid) except ttypes.InstanceOverloaded as ioe: # Instance we ran the call on was overloaded. We can retry # the command again, hopefully on a different instance self.incrementCounter("bulk_run.remote.overload_error") self.logger.error("Instance Overloaded: %d: %s", retry_count, ioe) if retry_count > self.BULK_RETRY_LIMIT: # Fail the calls return self._bulk_failure(chunk, str(ioe)) # Stagger the retries delay = random.uniform(self.BULK_RETRY_DELAY_MIN, self.BULK_RETRY_DELAY_MAX) await asyncio.sleep(delay) retry_count += 1 except Exception as e: raise ttypes.SessionException( message="bulk_run_remote failed: %r" % (e)) from e
async def close_session(self, tsession, client_ip, client_port, uuid): uuid = self._generate_new_uuid(old_uuid=uuid) try: session = CommandSession.get(tsession.id, client_ip, client_port) await session.close() except Exception as e: raise ttypes.SessionException(message="close_session failed: %r" % (e)) from e
async def _run_session( self, tsession, command, timeout, client_ip, client_port, uuid, prompt_re=None ): try: session = CommandSession.get(tsession.id, client_ip, client_port) return await self._run_command(session, command, timeout, uuid, prompt_re) except Exception as e: raise ttypes.SessionException(message="run_session failed: %r" % (e)) from e
async def _run_commands( self, commands, device, timeout, open_timeout, client_ip, client_port, uuid, return_exceptions=False, ): options = self._get_options(device, client_ip, client_port, open_timeout, timeout) if device.command_prompts: options["command_prompts"] = { c.encode(): p.encode() for c, p in device.command_prompts.items() } command = commands[0] devinfo = None session = None try: devinfo = await self._lookup_device(device) async with devinfo.create_session(self.service, device, options, loop=self.loop) as session: results = [] for command in commands: result = await self._run_command(session, command, timeout, uuid) results.append(result) return results except Exception as e: self._record_error(e, command, uuid, options, devinfo, session) if not isinstance(e, ttypes.SessionException): e = ttypes.SessionException(message="%r" % e) if return_exceptions: e.message = self.add_debug_info_to_error_message( # noqa error_msg=e.message, uuid=uuid # noqa ) return [ ttypes.CommandResult(output="", status="%r" % e, command=command) ] else: # raise from the original place so we have full stacktrace raise e
async def run_raw_session(self, tsession, command, timeout, prompt_regex, client_ip, client_port): if not prompt_regex: raise ttypes.SessionException(message="prompt_regex not specified") prompt_re = re.compile(prompt_regex.encode('utf8'), re.M) return await self._run_session(tsession, command, timeout, client_ip, client_port, prompt_re)
async def run_raw_session(self, tsession, command, timeout, prompt_regex, client_ip, client_port, uuid): uuid = self._generate_new_uuid(old_uuid=uuid) if not prompt_regex: raise ttypes.SessionException(message="prompt_regex not specified") prompt_re = re.compile(prompt_regex.encode("utf8"), re.M) return await self._run_session(tsession, command, timeout, client_ip, client_port, uuid, prompt_re)
async def _run_commands(self, commands, device, timeout, open_timeout, client_ip, client_port, return_exceptions=False): options = self._get_options(device, client_ip, client_port, open_timeout, timeout) if device.command_prompts: options['command_prompts'] = { c.encode(): p.encode() for c, p in device.command_prompts.items() } command = "" try: devinfo = await self._lookup_device(device) async with devinfo.create_session(self.service, device, options, loop=self.loop) as session: results = [] for command in commands: result = await self._run_command(session, command, timeout) results.append(result) return results except Exception as e: if not isinstance(e, ttypes.SessionException): e = ttypes.SessionException(message='%r' % e) if return_exceptions: return [ ttypes.CommandResult(output='', status='%r' % e, command=command) ] else: raise e from e