Neels Hofmeyr | 106865a | 2020-11-27 08:19:42 +0100 | [diff] [blame] | 1 | # osmo_gsm_tester: VTY connection |
| 2 | # |
| 3 | # Copyright (C) 2020 by sysmocom - s.f.m.c. GmbH |
| 4 | # |
| 5 | # Author: Neels Hofmeyr <neels@hofmeyr.de> |
| 6 | # |
| 7 | # This program is free software: you can redistribute it and/or modify |
| 8 | # it under the terms of the GNU General Public License as |
| 9 | # published by the Free Software Foundation, either version 3 of the |
| 10 | # License, or (at your option) any later version. |
| 11 | # |
| 12 | # This program is distributed in the hope that it will be useful, |
| 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 15 | # GNU General Public License for more details. |
| 16 | # |
| 17 | # You should have received a copy of the GNU General Public License |
| 18 | # along with this program. If not, see <http://www.gnu.org/licenses/>. |
| 19 | |
| 20 | import socket |
| 21 | import struct |
| 22 | import re |
| 23 | import time |
| 24 | import sys |
| 25 | |
| 26 | from ..core import log |
| 27 | from ..core.event_loop import MainLoop |
| 28 | |
| 29 | class VtyInterfaceExn(Exception): |
| 30 | pass |
| 31 | |
| 32 | class OsmoVty(log.Origin): |
| 33 | '''Suggested usage: |
| 34 | with OsmoVty(...) as vty: |
| 35 | vty.cmds('enable', 'configure network', 'net') |
| 36 | response = vty.cmd('foo 1 2 3') |
| 37 | print('\n'.join(response)) |
| 38 | |
| 39 | Using 'with' ensures that the connection is closed again. |
| 40 | There should not be nested 'with' statements on this object. |
Neels Hofmeyr | af4e231 | 2020-11-27 08:20:56 +0100 | [diff] [blame] | 41 | |
| 42 | Note that test env objects (like tenv.bsc()) may keep a VTY connected until the test exits. A 'with' should not |
| 43 | be used on those. |
Neels Hofmeyr | 106865a | 2020-11-27 08:19:42 +0100 | [diff] [blame] | 44 | ''' |
| 45 | |
| 46 | ############## |
| 47 | # PROTECTED |
| 48 | ############## |
| 49 | |
| 50 | def __init__(self, host, port, prompt=None): |
| 51 | super().__init__(log.C_BUS, 'Vty', host=host, port=port) |
| 52 | self.host = host |
| 53 | self.port = port |
| 54 | self.sck = None |
| 55 | self.prompt = prompt |
| 56 | self.re_prompt = None |
| 57 | self.this_node = None |
| 58 | self.this_prompt_char = None |
| 59 | self.last_node = None |
| 60 | self.last_prompt_char = None |
| 61 | |
| 62 | def try_connect(self): |
| 63 | '''Do a connection attempt, return True when successful, False otherwise. |
| 64 | Does not raise exceptions, but logs them to the debug log.''' |
| 65 | assert self.sck is None |
| 66 | try: |
| 67 | self.dbg('Connecting') |
| 68 | sck = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 69 | try: |
| 70 | sck.connect((self.host, self.port)) |
| 71 | except: |
| 72 | sck.close() |
| 73 | raise |
| 74 | # set self.sck only after the connect was successful |
| 75 | self.sck = sck |
| 76 | return True |
| 77 | except: |
| 78 | self.dbg('Failed to connect', sys.exc_info()[0]) |
| 79 | return False |
| 80 | |
| 81 | def _command(self, command_str, timeout=10, strict=True): |
| 82 | '''Send a command and return the response.''' |
| 83 | # (copied from https://git.osmocom.org/python/osmo-python-tests/tree/osmopy/osmo_interact/vty.py) |
| 84 | self.dbg('Sending', command_str=command_str) |
| 85 | self.sck.send(command_str.encode()) |
| 86 | |
| 87 | waited_since = time.time() |
| 88 | received_lines = [] |
| 89 | last_line = '' |
| 90 | |
| 91 | # (not using MainLoop.wait() to accumulate received responses across |
| 92 | # iterations) |
| 93 | while True: |
| 94 | new_data = self.sck.recv(4096).decode('utf-8') |
| 95 | |
| 96 | last_line = "%s%s" % (last_line, new_data) |
| 97 | |
| 98 | if last_line: |
| 99 | # Separate the received response into lines. |
| 100 | # But note: the VTY logging currently separates with '\n\r', not '\r\n', |
| 101 | # see _vty_output() in libosmocore logging_vty.c. |
| 102 | # So we need to jump through hoops to not separate 'abc\n\rdef' as |
| 103 | # [ 'abc', '', 'def' ]; but also not to convert '\r\n\r\n' to '\r\n\n' ('\r{\r\n}\n') |
| 104 | # Simplest is to just drop all the '\r' and only care about the '\n'. |
| 105 | last_line = last_line.replace('\r', '') |
| 106 | lines = last_line.splitlines() |
| 107 | if last_line.endswith('\n'): |
| 108 | received_lines.extend(lines) |
| 109 | last_line = "" |
| 110 | else: |
| 111 | # if pkt buffer ends in the middle of a line, we need to keep |
| 112 | # last non-finished line: |
| 113 | received_lines.extend(lines[:-1]) |
| 114 | last_line = lines[-1] |
| 115 | |
| 116 | match = self.re_prompt.match(last_line) |
| 117 | if not match: |
| 118 | if time.time() - waited_since > timeout: |
| 119 | raise IOError("Failed to read data (did the app crash?)") |
| 120 | MainLoop.sleep(.1) |
| 121 | continue |
| 122 | |
| 123 | self.last_node = self.this_node |
| 124 | self.last_prompt_char = self.this_prompt_char |
| 125 | self.this_node = match.group(1) or None |
| 126 | self.this_prompt_char = match.group(2) |
| 127 | break |
| 128 | |
| 129 | # expecting to have received the command we sent as echo, remove it |
| 130 | clean_command_str = command_str.strip() |
| 131 | if clean_command_str.endswith('?'): |
| 132 | clean_command_str = clean_command_str[:-1] |
| 133 | if received_lines and received_lines[0] == clean_command_str: |
| 134 | received_lines = received_lines[1:] |
| 135 | if len(received_lines) > 1: |
| 136 | self.dbg('Received\n|', '\n| '.join(received_lines), '\n') |
| 137 | elif len(received_lines) == 1: |
| 138 | self.dbg('Received', repr(received_lines[0])) |
| 139 | |
| 140 | if received_lines == ['% Unknown command.']: |
| 141 | errmsg = 'VTY reports unknown command: %r' % command_str |
| 142 | if strict: |
| 143 | raise VtyInterfaceExn(errmsg) |
| 144 | else: |
| 145 | self.log('ignoring error:', errmsg) |
| 146 | |
| 147 | return received_lines |
| 148 | |
| 149 | ######################## |
| 150 | # PUBLIC - INTERNAL API |
| 151 | ######################## |
| 152 | |
| 153 | def connect(self, timeout=30): |
| 154 | '''Connect to the VTY self.host and self.port, retry for 'timeout' seconds. |
| 155 | connect() and disconnect() are called implicitly when using the 'with' statement. |
| 156 | See class OsmoVty's doc. |
| 157 | ''' |
| 158 | MainLoop.wait(self.try_connect, timestep=3, timeout=timeout) |
| 159 | self.sck.setblocking(1) |
| 160 | |
| 161 | # read first prompt |
| 162 | # (copied from https://git.osmocom.org/python/osmo-python-tests/tree/osmopy/osmo_interact/vty.py) |
| 163 | self.this_node = None |
| 164 | self.this_prompt_char = '>' # slight cheat for initial prompt char |
| 165 | self.last_node = None |
| 166 | self.last_prompt_char = None |
| 167 | |
| 168 | data = self.sck.recv(4096) |
| 169 | if not self.prompt: |
| 170 | b = data |
| 171 | b = b[b.rfind(b'\n') + 1:] |
| 172 | while b and (b[0] < ord('A') or b[0] > ord('z')): |
| 173 | b = b[1:] |
| 174 | prompt_str = b.decode('utf-8') |
| 175 | if '>' in prompt_str: |
| 176 | self.prompt = prompt_str[:prompt_str.find('>')] |
| 177 | self.dbg(prompt=self.prompt) |
| 178 | if not self.prompt: |
| 179 | raise VtyInterfaceExn('Could not find application name; needed to decode prompts.' |
| 180 | ' Initial data was: %r' % data) |
| 181 | self.re_prompt = re.compile('^%s(?:\(([\w-]*)\))?([#>]) (.*)$' % self.prompt) |
| 182 | |
| 183 | def disconnect(self): |
| 184 | '''Disconnect. |
| 185 | connect() and disconnect() are called implicitly when using the 'with' statement. |
| 186 | See class OsmoVty's doc. |
| 187 | ''' |
| 188 | if self.sck is None: |
| 189 | return |
| 190 | self.dbg('Disconnecting') |
| 191 | self.sck.close() |
| 192 | self.sck = None |
| 193 | |
| 194 | ################### |
| 195 | # PUBLIC (test API included) |
| 196 | ################### |
| 197 | |
| 198 | def cmd(self, command_str, timeout=10, strict=True): |
| 199 | '''Send one VTY command and return its response. |
| 200 | Return a list of strings, one string per line, without line break characters: |
| 201 | [ 'first line', 'second line', 'third line' ] |
| 202 | When strict==False, do not raise exceptions on '% Unknown command'. |
| 203 | If the connection is not yet open, briefly connect for only this command and disconnect again. If it is open, |
| 204 | use the open connection and leave it open. |
| 205 | ''' |
| 206 | # allow calling for both already connected VTY as well as establishing |
| 207 | # a connection just for this command. |
| 208 | if self.sck is None: |
| 209 | with self: |
| 210 | return self.cmd(command_str, timeout, strict) |
| 211 | |
| 212 | # (copied from https://git.osmocom.org/python/osmo-python-tests/tree/osmopy/osmo_interact/vty.py) |
| 213 | command_str = command_str or '\r' |
| 214 | if command_str[-1] not in '?\r\t': |
| 215 | command_str = command_str + '\r' |
| 216 | |
| 217 | received_lines = self._command(command_str, timeout, strict) |
| 218 | |
| 219 | # send escape to cancel the '?' command line |
| 220 | if command_str[-1] == '?': |
| 221 | self._command('\x03', timeout) |
| 222 | |
| 223 | return received_lines |
| 224 | |
| 225 | def cmds(self, *cmds, timeout=10, strict=True): |
| 226 | '''Send a series of commands and return each command's response: |
| 227 | cmds('foo', 'bar', 'baz') --> [ ['foo line 1','foo line 2'], ['bar line 1'], ['baz line 1']] |
| 228 | When strict==False, do not raise exceptions on '% Unknown command'. |
| 229 | If the connection is not yet open, briefly connect for only these commands and disconnect again. If it is |
| 230 | open, use the open connection and leave it open. |
| 231 | ''' |
| 232 | # allow calling for both already connected VTY as well as establishing |
| 233 | # a connection just for this command. |
| 234 | if self.sck is None: |
| 235 | with self: |
| 236 | return self.cmds(*cmds, timeout=timeout, strict=strict) |
| 237 | |
| 238 | responses = [] |
| 239 | for cmd in cmds: |
| 240 | responses.append(self.cmd(cmd, timeout, strict)) |
| 241 | return responses |
| 242 | |
| 243 | def __enter__(self): |
| 244 | self.connect() |
| 245 | return self |
| 246 | |
| 247 | def __exit__(self, *exc_info): |
| 248 | self.disconnect() |
| 249 | |
| 250 | # vim: expandtab tabstop=4 shiftwidth=4 |