Module pandare.panda_expect
Custom library for interacting/expecting data via serial-like FDs
Expand source code
""" Custom library for interacting/expecting data via serial-like FDs"""
import os
import re
import select
import sys
import string
from time import monotonic
from errno import EAGAIN, EWOULDBLOCK
from colorama import Fore, Style
class TimeoutExpired(Exception): pass
class Expect(object):
'''
Class to manage typing commands into consoles and waiting for responses.
Designed to be used with the qemu monitor and serial consoles for Linux guests.
'''
def __init__(self, name, filelike=None, expectation=None, logfile_base=None, consume_first=False, unansi=False):
'''
To debug, set logfile_base to something like '/tmp/log' and then look at logs written to /tmp/log_monitor.txt and /tmp/log_serial.txt. Or directyl access
'''
self.name = name
self.logfile = None
if logfile_base:
self.set_logging(f"{logfile_base}_{name}.txt")
if filelike is None: # Must later use connect(filelike)
self.fd = None
else:
self.connect(filelike)
self.prior_lines = []
self.current_line = bytearray()
self.last_msg = None
self.running = True
self.cleared = False
self.update_expectation(expectation)
# If consumed_first is false, we'll consume a message before anything else. Requires self.expectation to be set
self.consumed_first = not consume_first
self.use_unansi = unansi
def update_expectation(self, expectation):
if isinstance(expectation, bytes):
expectation = expectation.decode()
self.last_prompt = expectation # approximation
self.expectation_re = re.compile(expectation)
self.expectation_ends_re = re.compile(r'(.*)' + expectation)
def set_logging(self, name):
self.logfile = open(name, "wb")
def connect(self, filelike):
if type(filelike) == int:
self.fd = filelike
else:
self.fd = filelike.fileno()
self.poller = select.poll()
self.poller.register(self.fd, select.POLLIN)
def is_connected(self):
return self.fd != None
def __del__(self):
if self.logfile:
self.logfile.close()
def abort(self):
self.running = False
def consume_partial(self):
'''
Get the message so far and reset.
To ensure that we're not consuming the final line at we just don't clear
'''
result = self.get_partial()
self.prior_lines = []
return result
def get_partial(self):
'''
Get the message
'''
if len(self.prior_lines):
return "\n".join(self.prior_lines)
return ""
def unansi(self):
'''
Take the string in self.current_line and any prior lines in self.prior_lines and render ANSI.
prior lines should be plain strings while current_line may contain escapes
Given a string with ansi control codes, emulate behavior to generate the resulting string.
First we split input into a list of ('fn', [args]) / ('text', ['foo']) ansi commands then
evaluate the commands to render real text output
See https://notes.burke.libbey.me/ansi-escape-codes/ and
http://ascii-table.com/ansi-escape-sequences-vt-100.php for ansi escape code details
'''
# Join prior lines into a single text element in our reformatted list
reformatted = []
if len(self.prior_lines):
reformatted = [('text', ['\n'.join(self.prior_lines)])]
# Then split current line into the tuple format describe above
msg = self.current_line
if isinstance(msg, str):
msg = msg.encode()
# Check for simple case where no ansi is in line - if so just copy into prior_lines
if b'\x1b' not in msg:
text = "".join([chr(x) for x in msg]).strip()
self.prior_lines.append(text)
self.current_line = bytearray()
return
start_args = re.compile(br"^(\d+);")
last_arg = re.compile(rb"^(\d+)")
last_text = ""
idx = 0 # XXX: mutates during loop
while idx < len(msg):
if msg[idx] != 0x1b:
last_text += chr(msg[idx])
else:
if len(last_text):
reformatted.append(('text', [last_text]))
last_text = ""
if idx+3 <= len(msg) and msg[idx+1] == ord('['):
args = []
shift = idx+2
arg_s = msg[shift:]
while start_args.match(arg_s):
arg = start_args.match(arg_s).groups()[0].decode()
args.append(arg)
shift += len(arg)+1 # for ;
arg_s = msg[shift:]
# Last arg is just #
if last_arg.match(arg_s):
arg = last_arg.match(arg_s).groups()[0].decode()
shift += len(arg)
args.append(arg)
arg_s = msg[shift:]
# Next is one char for cmd
cmd = chr(msg[shift])
reformatted.append((cmd, args))
idx = shift # final char
idx += 1
if len(last_text):
reformatted.append(('text', [last_text]))
# Now render it!
# Note the very first line will \r to beginning, then 'C' forward to go past expect prompt
lines_out = [" "*len(self.last_prompt)] # Starting point - it's an approximation since we don't know real current prompt
cur_line = 0
line_pos = len(self.last_prompt)
store_ptr = (0, 0)
def _dump(lines_out, cur_line, line_pos):
print("-"*100)
for idx, line in enumerate(lines_out):
print("'", line, "'")
if cur_line == idx:
print((" "*(line_pos-1) if line_pos > 0 else "") + "^")
print("="*100)
for idx, (typ, args) in enumerate(reformatted):
#print(typ, args)
if typ == 'text':
n = args[0]
for idx, char in enumerate(n):
if char == '\n':
cur_line += 1
while cur_line >= len(lines_out):
lines_out.append("")
if char == '\r':
line_pos = 0
continue # Don't clobber old char
line = list(lines_out[cur_line])
if (line_pos) >= len(line):
line.append(char)
else:
#if line[line_pos] != ' ':
# print("Replace", repr(line[line_pos]) , "with", repr(char))
line[line_pos] = char
lines_out[cur_line] = "".join(line)
if char not in ['\n', '\r']:
line_pos += 1
else:
args[:] = [int(x) for x in args]
if typ == 'A':
if not len(args):
# Incomplete
continue
n = args[0]
if cur_line - n < 0: # Need to shift
cur_line = 0
else:
cur_line -= n
assert(cur_line >= 0)
elif typ == 'B':
if not len(args):
# Incomplete
continue
n = args[0]
cur_line += n
while cur_line >= len(lines_out):
lines_out.append(" "*100)
elif typ == 'D':
n = 1 # Default move left 1
if len(args):
n = args[0]
line_pos -= n
if line_pos < 0:
line_pos = 0
assert(line_pos >= 0)
elif typ == 'C': # Move right
n = 1 # Default move 1
if len(args):
n = args[0]
line_pos += n
if line_pos > len(lines_out[cur_line])-1:
line_pos = len(lines_out)-1
assert(line_pos >= 0)
elif typ == 'J':
# Optional arg 0, 1, 2
n = 0 # default
if len(args):
n = args[0]
if n == 0:
# clear down
lines_out = lines_out[:cur_line+1]
elif n == 1:
# clear up
lines_out = lines_out[cur_line:]
elif n == 2:
# clear everything
lines_out = [""]
cur_line = 0
line_pos = 0
store_ptr = (0, 0)
elif typ == 'K':
# Optional arg 0, 1, 2
n = 0 # default
if len(args):
n = args[0]
# HURISTIC-y hack: linux loves to have a line 123456 then do K(0) 6\r\n
# so if the next line is text and the text is [eol]\r\n where [eol] matches the end of this
# line - align cur_pos
if len(reformatted) > idx+1: # Have another message
(next_typ, next_args) = reformatted[idx+1]
if next_typ == 'text':
if '\r\n' in next_args[0]:
next_lines = next_args[0].split("\r\n")
if lines_out[cur_line].strip().endswith(next_lines[0]):
# Its the buggy case. Just align it such that we clear the text
# that's about to get echoed
line_pos = line_pos - len(next_lines[0])
if n == 0:
# clear right of cursor
lines_out[cur_line] = lines_out[cur_line][:line_pos]
elif n == 1:
# clear left of cursor
lines_out[cur_line] = (" "*line_pos)+lines_out[cur_line][line_pos:]
elif n == 2:
# clear whole line
lines_out[cur_line] = " "*len(lines_out[cur_line])
elif typ == 'H':
n = args[0]-1
m = args[1]-1
cur_line = n
line_pos = m
while cur_line >= len(lines_out):
lines_out.append("")
while line_pos > len(lines_out[cur_line]):
lines_out[cur_line] += " "
elif typ == 'T':
# Scroll window down
pass
elif typ == 'S':
# Scroll window up
pass
elif typ == 's':
store_ptr = (cur_line, line_pos)
elif typ == 'u':
(cur_line, line_pos) = store_ptr
elif typ == 'm':
# alter character attributes - just ignore
pass
else:
raise ValueError(f"Unsupported ANSI command {typ}")
#_dump(lines_out, cur_line, line_pos)
# Done processing - update variables
self.prior_lines = lines_out[:-1] # Strings
self.current_line = bytearray() # Bytearray
if len(lines_out[-1].strip()):
self.current_line.append(lines_out[-1])
def expect(self, expectation=None, timeout=30):
'''
Assumptions: as you send a command, the guest may send back
The same command + ansi control codes.
The epxectation value will show up on the start of a line.
We add characters into current_line as we recv them. At each newline we
1) Render ANSI control characters in current line (may affect prior lines)
2) Check if the line we just parsed matches the provided expectation (if so we're done)
3) Append current_line into prior_lines
'''
if expectation:
raise ValueError("Deprecated interface - must set expectation in class init")
if self.fd is None:
raise RuntimeError("Must connect() prior to expect()")
self.current_line = bytearray()
start_time = monotonic()
time_passed = 0
while (timeout is None or time_passed < timeout) and self.running:
if timeout is not None:
time_passed = (monotonic() - start_time)
time_left = timeout - time_passed
else:
time_left = float("inf")
ready = self.poller.poll(min(time_left, 1))
# Debug - flush debug logs
if self.logfile:
self.logfile.flush()
if self.fd in [fd for (fd, _) in ready]:
try:
char = os.read(self.fd, 1)
except OSError as e:
if e.errno in [EAGAIN, EWOULDBLOCK]:
continue
else: raise
self.current_line.extend(char)
# Debugging - log current line to file
if self.logfile:
self.logfile.write(("\n\n" + repr(self.prior_lines) + " Current line = " + repr(self.current_line)).encode())
# Translate the current_line buffer into plaintext, then determine if we're finished (bc we see new prompt)
# note this drops the echo'd command
if self.current_line.endswith(b"\n"):
# End of line - need to potentially unansi and move into prior_lines
if self.use_unansi:
self.unansi()
else:
self.prior_lines.append(self.current_line[:-1].decode(errors='ignore'))
self.current_line = bytearray()
# Now we have command\nresults..........\nprompt
#self.logfile.write(b"\n UNANSIs to: " + repr(self.prior_lines).encode()+b"\n")
#lines = [x.replace("\r", "") for x in plaintext.split("\n")]
# Check current line to see if it ends with prompt (indicating we finished)
# current_line is a bytearray. Need it as a string
current_line_s = self.current_line.decode(errors='ignore')
end_match = self.expectation_ends_re.match(current_line_s)
if end_match is not None:
# This line matches the end regex - it's either like root@host:... or it's [output]root@host:...
# We'll use self.expectation_re on the current line to identify where the prompt is and grab any final output
final_output = end_match.groups(1)[0]
if len(final_output):
self.prior_lines.append(final_output)
current_line_s = current_line_s[len(final_output):]
# Note we may have a line like [output]root@.... in which case we need to identify where the prompt was
self.last_prompt = current_line_s
# Drop command we sent - note it won't be a direct match with last_cmd because of weird escape codes
# which are based on guest line position when printed - i.e., it would only be an exact
# match if we knew and included the prompt when the command was run. Let's just always drop it
if len(self.prior_lines) > 1:
self.prior_lines = self.prior_lines[1:]
else:
self.prior_lines = []
plaintext = "\n".join(self.prior_lines)
self.prior_lines = []
return plaintext
if not self.running: # Aborted
return None
if self.logfile:
self.logfile.flush()
full_buffer = self.prior_lines + [self.current_line]
raise TimeoutExpired(f"{self.name} Read message \n{full_buffer}\n")
def send(self, msg):
if not self.consumed_first: # Before we send anything, consume header
pre = self.expect("")
self.consumed_first = True
# Newlines will call problems
assert len(msg.decode(errors='ignore').split("\n")) <= 2, "Multiline cmds unsupported"
self.last_msg = msg.decode(errors='ignore').replace("\n", "")
os.write(self.fd, msg)
if self.logfile:
self.logfile.write(msg)
self.logfile.flush()
def send_eol(self): # Just send an EOL
if self.last_msg:
self.last_msg+="\n"
os.write(self.fd, b"\n")
if self.logfile:
self.logfile.write(b"\n")
self.logfile.flush()
def sendline(self, msg=b""):
assert(self.fd is not None), "Must connect before sending"
self.send(msg + b"\n")Classes
class Expect (name, filelike=None, expectation=None, logfile_base=None, consume_first=False, unansi=False)-
Class to manage typing commands into consoles and waiting for responses.
Designed to be used with the qemu monitor and serial consoles for Linux guests.
To debug, set logfile_base to something like '/tmp/log' and then look at logs written to /tmp/log_monitor.txt and /tmp/log_serial.txt. Or directyl access
Expand source code
class Expect(object): ''' Class to manage typing commands into consoles and waiting for responses. Designed to be used with the qemu monitor and serial consoles for Linux guests. ''' def __init__(self, name, filelike=None, expectation=None, logfile_base=None, consume_first=False, unansi=False): ''' To debug, set logfile_base to something like '/tmp/log' and then look at logs written to /tmp/log_monitor.txt and /tmp/log_serial.txt. Or directyl access ''' self.name = name self.logfile = None if logfile_base: self.set_logging(f"{logfile_base}_{name}.txt") if filelike is None: # Must later use connect(filelike) self.fd = None else: self.connect(filelike) self.prior_lines = [] self.current_line = bytearray() self.last_msg = None self.running = True self.cleared = False self.update_expectation(expectation) # If consumed_first is false, we'll consume a message before anything else. Requires self.expectation to be set self.consumed_first = not consume_first self.use_unansi = unansi def update_expectation(self, expectation): if isinstance(expectation, bytes): expectation = expectation.decode() self.last_prompt = expectation # approximation self.expectation_re = re.compile(expectation) self.expectation_ends_re = re.compile(r'(.*)' + expectation) def set_logging(self, name): self.logfile = open(name, "wb") def connect(self, filelike): if type(filelike) == int: self.fd = filelike else: self.fd = filelike.fileno() self.poller = select.poll() self.poller.register(self.fd, select.POLLIN) def is_connected(self): return self.fd != None def __del__(self): if self.logfile: self.logfile.close() def abort(self): self.running = False def consume_partial(self): ''' Get the message so far and reset. To ensure that we're not consuming the final line at we just don't clear ''' result = self.get_partial() self.prior_lines = [] return result def get_partial(self): ''' Get the message ''' if len(self.prior_lines): return "\n".join(self.prior_lines) return "" def unansi(self): ''' Take the string in self.current_line and any prior lines in self.prior_lines and render ANSI. prior lines should be plain strings while current_line may contain escapes Given a string with ansi control codes, emulate behavior to generate the resulting string. First we split input into a list of ('fn', [args]) / ('text', ['foo']) ansi commands then evaluate the commands to render real text output See https://notes.burke.libbey.me/ansi-escape-codes/ and http://ascii-table.com/ansi-escape-sequences-vt-100.php for ansi escape code details ''' # Join prior lines into a single text element in our reformatted list reformatted = [] if len(self.prior_lines): reformatted = [('text', ['\n'.join(self.prior_lines)])] # Then split current line into the tuple format describe above msg = self.current_line if isinstance(msg, str): msg = msg.encode() # Check for simple case where no ansi is in line - if so just copy into prior_lines if b'\x1b' not in msg: text = "".join([chr(x) for x in msg]).strip() self.prior_lines.append(text) self.current_line = bytearray() return start_args = re.compile(br"^(\d+);") last_arg = re.compile(rb"^(\d+)") last_text = "" idx = 0 # XXX: mutates during loop while idx < len(msg): if msg[idx] != 0x1b: last_text += chr(msg[idx]) else: if len(last_text): reformatted.append(('text', [last_text])) last_text = "" if idx+3 <= len(msg) and msg[idx+1] == ord('['): args = [] shift = idx+2 arg_s = msg[shift:] while start_args.match(arg_s): arg = start_args.match(arg_s).groups()[0].decode() args.append(arg) shift += len(arg)+1 # for ; arg_s = msg[shift:] # Last arg is just # if last_arg.match(arg_s): arg = last_arg.match(arg_s).groups()[0].decode() shift += len(arg) args.append(arg) arg_s = msg[shift:] # Next is one char for cmd cmd = chr(msg[shift]) reformatted.append((cmd, args)) idx = shift # final char idx += 1 if len(last_text): reformatted.append(('text', [last_text])) # Now render it! # Note the very first line will \r to beginning, then 'C' forward to go past expect prompt lines_out = [" "*len(self.last_prompt)] # Starting point - it's an approximation since we don't know real current prompt cur_line = 0 line_pos = len(self.last_prompt) store_ptr = (0, 0) def _dump(lines_out, cur_line, line_pos): print("-"*100) for idx, line in enumerate(lines_out): print("'", line, "'") if cur_line == idx: print((" "*(line_pos-1) if line_pos > 0 else "") + "^") print("="*100) for idx, (typ, args) in enumerate(reformatted): #print(typ, args) if typ == 'text': n = args[0] for idx, char in enumerate(n): if char == '\n': cur_line += 1 while cur_line >= len(lines_out): lines_out.append("") if char == '\r': line_pos = 0 continue # Don't clobber old char line = list(lines_out[cur_line]) if (line_pos) >= len(line): line.append(char) else: #if line[line_pos] != ' ': # print("Replace", repr(line[line_pos]) , "with", repr(char)) line[line_pos] = char lines_out[cur_line] = "".join(line) if char not in ['\n', '\r']: line_pos += 1 else: args[:] = [int(x) for x in args] if typ == 'A': if not len(args): # Incomplete continue n = args[0] if cur_line - n < 0: # Need to shift cur_line = 0 else: cur_line -= n assert(cur_line >= 0) elif typ == 'B': if not len(args): # Incomplete continue n = args[0] cur_line += n while cur_line >= len(lines_out): lines_out.append(" "*100) elif typ == 'D': n = 1 # Default move left 1 if len(args): n = args[0] line_pos -= n if line_pos < 0: line_pos = 0 assert(line_pos >= 0) elif typ == 'C': # Move right n = 1 # Default move 1 if len(args): n = args[0] line_pos += n if line_pos > len(lines_out[cur_line])-1: line_pos = len(lines_out)-1 assert(line_pos >= 0) elif typ == 'J': # Optional arg 0, 1, 2 n = 0 # default if len(args): n = args[0] if n == 0: # clear down lines_out = lines_out[:cur_line+1] elif n == 1: # clear up lines_out = lines_out[cur_line:] elif n == 2: # clear everything lines_out = [""] cur_line = 0 line_pos = 0 store_ptr = (0, 0) elif typ == 'K': # Optional arg 0, 1, 2 n = 0 # default if len(args): n = args[0] # HURISTIC-y hack: linux loves to have a line 123456 then do K(0) 6\r\n # so if the next line is text and the text is [eol]\r\n where [eol] matches the end of this # line - align cur_pos if len(reformatted) > idx+1: # Have another message (next_typ, next_args) = reformatted[idx+1] if next_typ == 'text': if '\r\n' in next_args[0]: next_lines = next_args[0].split("\r\n") if lines_out[cur_line].strip().endswith(next_lines[0]): # Its the buggy case. Just align it such that we clear the text # that's about to get echoed line_pos = line_pos - len(next_lines[0]) if n == 0: # clear right of cursor lines_out[cur_line] = lines_out[cur_line][:line_pos] elif n == 1: # clear left of cursor lines_out[cur_line] = (" "*line_pos)+lines_out[cur_line][line_pos:] elif n == 2: # clear whole line lines_out[cur_line] = " "*len(lines_out[cur_line]) elif typ == 'H': n = args[0]-1 m = args[1]-1 cur_line = n line_pos = m while cur_line >= len(lines_out): lines_out.append("") while line_pos > len(lines_out[cur_line]): lines_out[cur_line] += " " elif typ == 'T': # Scroll window down pass elif typ == 'S': # Scroll window up pass elif typ == 's': store_ptr = (cur_line, line_pos) elif typ == 'u': (cur_line, line_pos) = store_ptr elif typ == 'm': # alter character attributes - just ignore pass else: raise ValueError(f"Unsupported ANSI command {typ}") #_dump(lines_out, cur_line, line_pos) # Done processing - update variables self.prior_lines = lines_out[:-1] # Strings self.current_line = bytearray() # Bytearray if len(lines_out[-1].strip()): self.current_line.append(lines_out[-1]) def expect(self, expectation=None, timeout=30): ''' Assumptions: as you send a command, the guest may send back The same command + ansi control codes. The epxectation value will show up on the start of a line. We add characters into current_line as we recv them. At each newline we 1) Render ANSI control characters in current line (may affect prior lines) 2) Check if the line we just parsed matches the provided expectation (if so we're done) 3) Append current_line into prior_lines ''' if expectation: raise ValueError("Deprecated interface - must set expectation in class init") if self.fd is None: raise RuntimeError("Must connect() prior to expect()") self.current_line = bytearray() start_time = monotonic() time_passed = 0 while (timeout is None or time_passed < timeout) and self.running: if timeout is not None: time_passed = (monotonic() - start_time) time_left = timeout - time_passed else: time_left = float("inf") ready = self.poller.poll(min(time_left, 1)) # Debug - flush debug logs if self.logfile: self.logfile.flush() if self.fd in [fd for (fd, _) in ready]: try: char = os.read(self.fd, 1) except OSError as e: if e.errno in [EAGAIN, EWOULDBLOCK]: continue else: raise self.current_line.extend(char) # Debugging - log current line to file if self.logfile: self.logfile.write(("\n\n" + repr(self.prior_lines) + " Current line = " + repr(self.current_line)).encode()) # Translate the current_line buffer into plaintext, then determine if we're finished (bc we see new prompt) # note this drops the echo'd command if self.current_line.endswith(b"\n"): # End of line - need to potentially unansi and move into prior_lines if self.use_unansi: self.unansi() else: self.prior_lines.append(self.current_line[:-1].decode(errors='ignore')) self.current_line = bytearray() # Now we have command\nresults..........\nprompt #self.logfile.write(b"\n UNANSIs to: " + repr(self.prior_lines).encode()+b"\n") #lines = [x.replace("\r", "") for x in plaintext.split("\n")] # Check current line to see if it ends with prompt (indicating we finished) # current_line is a bytearray. Need it as a string current_line_s = self.current_line.decode(errors='ignore') end_match = self.expectation_ends_re.match(current_line_s) if end_match is not None: # This line matches the end regex - it's either like root@host:... or it's [output]root@host:... # We'll use self.expectation_re on the current line to identify where the prompt is and grab any final output final_output = end_match.groups(1)[0] if len(final_output): self.prior_lines.append(final_output) current_line_s = current_line_s[len(final_output):] # Note we may have a line like [output]root@.... in which case we need to identify where the prompt was self.last_prompt = current_line_s # Drop command we sent - note it won't be a direct match with last_cmd because of weird escape codes # which are based on guest line position when printed - i.e., it would only be an exact # match if we knew and included the prompt when the command was run. Let's just always drop it if len(self.prior_lines) > 1: self.prior_lines = self.prior_lines[1:] else: self.prior_lines = [] plaintext = "\n".join(self.prior_lines) self.prior_lines = [] return plaintext if not self.running: # Aborted return None if self.logfile: self.logfile.flush() full_buffer = self.prior_lines + [self.current_line] raise TimeoutExpired(f"{self.name} Read message \n{full_buffer}\n") def send(self, msg): if not self.consumed_first: # Before we send anything, consume header pre = self.expect("") self.consumed_first = True # Newlines will call problems assert len(msg.decode(errors='ignore').split("\n")) <= 2, "Multiline cmds unsupported" self.last_msg = msg.decode(errors='ignore').replace("\n", "") os.write(self.fd, msg) if self.logfile: self.logfile.write(msg) self.logfile.flush() def send_eol(self): # Just send an EOL if self.last_msg: self.last_msg+="\n" os.write(self.fd, b"\n") if self.logfile: self.logfile.write(b"\n") self.logfile.flush() def sendline(self, msg=b""): assert(self.fd is not None), "Must connect before sending" self.send(msg + b"\n")Methods
def abort(self)-
Expand source code
def abort(self): self.running = False def connect(self, filelike)-
Expand source code
def connect(self, filelike): if type(filelike) == int: self.fd = filelike else: self.fd = filelike.fileno() self.poller = select.poll() self.poller.register(self.fd, select.POLLIN) def consume_partial(self)-
Get the message so far and reset. To ensure that we're not consuming the final line at we just don't clear
Expand source code
def consume_partial(self): ''' Get the message so far and reset. To ensure that we're not consuming the final line at we just don't clear ''' result = self.get_partial() self.prior_lines = [] return result def expect(self, expectation=None, timeout=30)-
Assumptions: as you send a command, the guest may send back The same command + ansi control codes. The epxectation value will show up on the start of a line.
We add characters into current_line as we recv them. At each newline we 1) Render ANSI control characters in current line (may affect prior lines) 2) Check if the line we just parsed matches the provided expectation (if so we're done) 3) Append current_line into prior_lines
Expand source code
def expect(self, expectation=None, timeout=30): ''' Assumptions: as you send a command, the guest may send back The same command + ansi control codes. The epxectation value will show up on the start of a line. We add characters into current_line as we recv them. At each newline we 1) Render ANSI control characters in current line (may affect prior lines) 2) Check if the line we just parsed matches the provided expectation (if so we're done) 3) Append current_line into prior_lines ''' if expectation: raise ValueError("Deprecated interface - must set expectation in class init") if self.fd is None: raise RuntimeError("Must connect() prior to expect()") self.current_line = bytearray() start_time = monotonic() time_passed = 0 while (timeout is None or time_passed < timeout) and self.running: if timeout is not None: time_passed = (monotonic() - start_time) time_left = timeout - time_passed else: time_left = float("inf") ready = self.poller.poll(min(time_left, 1)) # Debug - flush debug logs if self.logfile: self.logfile.flush() if self.fd in [fd for (fd, _) in ready]: try: char = os.read(self.fd, 1) except OSError as e: if e.errno in [EAGAIN, EWOULDBLOCK]: continue else: raise self.current_line.extend(char) # Debugging - log current line to file if self.logfile: self.logfile.write(("\n\n" + repr(self.prior_lines) + " Current line = " + repr(self.current_line)).encode()) # Translate the current_line buffer into plaintext, then determine if we're finished (bc we see new prompt) # note this drops the echo'd command if self.current_line.endswith(b"\n"): # End of line - need to potentially unansi and move into prior_lines if self.use_unansi: self.unansi() else: self.prior_lines.append(self.current_line[:-1].decode(errors='ignore')) self.current_line = bytearray() # Now we have command\nresults..........\nprompt #self.logfile.write(b"\n UNANSIs to: " + repr(self.prior_lines).encode()+b"\n") #lines = [x.replace("\r", "") for x in plaintext.split("\n")] # Check current line to see if it ends with prompt (indicating we finished) # current_line is a bytearray. Need it as a string current_line_s = self.current_line.decode(errors='ignore') end_match = self.expectation_ends_re.match(current_line_s) if end_match is not None: # This line matches the end regex - it's either like root@host:... or it's [output]root@host:... # We'll use self.expectation_re on the current line to identify where the prompt is and grab any final output final_output = end_match.groups(1)[0] if len(final_output): self.prior_lines.append(final_output) current_line_s = current_line_s[len(final_output):] # Note we may have a line like [output]root@.... in which case we need to identify where the prompt was self.last_prompt = current_line_s # Drop command we sent - note it won't be a direct match with last_cmd because of weird escape codes # which are based on guest line position when printed - i.e., it would only be an exact # match if we knew and included the prompt when the command was run. Let's just always drop it if len(self.prior_lines) > 1: self.prior_lines = self.prior_lines[1:] else: self.prior_lines = [] plaintext = "\n".join(self.prior_lines) self.prior_lines = [] return plaintext if not self.running: # Aborted return None if self.logfile: self.logfile.flush() full_buffer = self.prior_lines + [self.current_line] raise TimeoutExpired(f"{self.name} Read message \n{full_buffer}\n") def get_partial(self)-
Get the message
Expand source code
def get_partial(self): ''' Get the message ''' if len(self.prior_lines): return "\n".join(self.prior_lines) return "" def is_connected(self)-
Expand source code
def is_connected(self): return self.fd != None def send(self, msg)-
Expand source code
def send(self, msg): if not self.consumed_first: # Before we send anything, consume header pre = self.expect("") self.consumed_first = True # Newlines will call problems assert len(msg.decode(errors='ignore').split("\n")) <= 2, "Multiline cmds unsupported" self.last_msg = msg.decode(errors='ignore').replace("\n", "") os.write(self.fd, msg) if self.logfile: self.logfile.write(msg) self.logfile.flush() def send_eol(self)-
Expand source code
def send_eol(self): # Just send an EOL if self.last_msg: self.last_msg+="\n" os.write(self.fd, b"\n") if self.logfile: self.logfile.write(b"\n") self.logfile.flush() def sendline(self, msg=b'')-
Expand source code
def sendline(self, msg=b""): assert(self.fd is not None), "Must connect before sending" self.send(msg + b"\n") def set_logging(self, name)-
Expand source code
def set_logging(self, name): self.logfile = open(name, "wb") def unansi(self)-
Take the string in self.current_line and any prior lines in self.prior_lines and render ANSI. prior lines should be plain strings while current_line may contain escapes
Given a string with ansi control codes, emulate behavior to generate the resulting string.
First we split input into a list of ('fn', [args]) / ('text', ['foo']) ansi commands then evaluate the commands to render real text output
See https://notes.burke.libbey.me/ansi-escape-codes/ and http://ascii-table.com/ansi-escape-sequences-vt-100.php for ansi escape code details
Expand source code
def unansi(self): ''' Take the string in self.current_line and any prior lines in self.prior_lines and render ANSI. prior lines should be plain strings while current_line may contain escapes Given a string with ansi control codes, emulate behavior to generate the resulting string. First we split input into a list of ('fn', [args]) / ('text', ['foo']) ansi commands then evaluate the commands to render real text output See https://notes.burke.libbey.me/ansi-escape-codes/ and http://ascii-table.com/ansi-escape-sequences-vt-100.php for ansi escape code details ''' # Join prior lines into a single text element in our reformatted list reformatted = [] if len(self.prior_lines): reformatted = [('text', ['\n'.join(self.prior_lines)])] # Then split current line into the tuple format describe above msg = self.current_line if isinstance(msg, str): msg = msg.encode() # Check for simple case where no ansi is in line - if so just copy into prior_lines if b'\x1b' not in msg: text = "".join([chr(x) for x in msg]).strip() self.prior_lines.append(text) self.current_line = bytearray() return start_args = re.compile(br"^(\d+);") last_arg = re.compile(rb"^(\d+)") last_text = "" idx = 0 # XXX: mutates during loop while idx < len(msg): if msg[idx] != 0x1b: last_text += chr(msg[idx]) else: if len(last_text): reformatted.append(('text', [last_text])) last_text = "" if idx+3 <= len(msg) and msg[idx+1] == ord('['): args = [] shift = idx+2 arg_s = msg[shift:] while start_args.match(arg_s): arg = start_args.match(arg_s).groups()[0].decode() args.append(arg) shift += len(arg)+1 # for ; arg_s = msg[shift:] # Last arg is just # if last_arg.match(arg_s): arg = last_arg.match(arg_s).groups()[0].decode() shift += len(arg) args.append(arg) arg_s = msg[shift:] # Next is one char for cmd cmd = chr(msg[shift]) reformatted.append((cmd, args)) idx = shift # final char idx += 1 if len(last_text): reformatted.append(('text', [last_text])) # Now render it! # Note the very first line will \r to beginning, then 'C' forward to go past expect prompt lines_out = [" "*len(self.last_prompt)] # Starting point - it's an approximation since we don't know real current prompt cur_line = 0 line_pos = len(self.last_prompt) store_ptr = (0, 0) def _dump(lines_out, cur_line, line_pos): print("-"*100) for idx, line in enumerate(lines_out): print("'", line, "'") if cur_line == idx: print((" "*(line_pos-1) if line_pos > 0 else "") + "^") print("="*100) for idx, (typ, args) in enumerate(reformatted): #print(typ, args) if typ == 'text': n = args[0] for idx, char in enumerate(n): if char == '\n': cur_line += 1 while cur_line >= len(lines_out): lines_out.append("") if char == '\r': line_pos = 0 continue # Don't clobber old char line = list(lines_out[cur_line]) if (line_pos) >= len(line): line.append(char) else: #if line[line_pos] != ' ': # print("Replace", repr(line[line_pos]) , "with", repr(char)) line[line_pos] = char lines_out[cur_line] = "".join(line) if char not in ['\n', '\r']: line_pos += 1 else: args[:] = [int(x) for x in args] if typ == 'A': if not len(args): # Incomplete continue n = args[0] if cur_line - n < 0: # Need to shift cur_line = 0 else: cur_line -= n assert(cur_line >= 0) elif typ == 'B': if not len(args): # Incomplete continue n = args[0] cur_line += n while cur_line >= len(lines_out): lines_out.append(" "*100) elif typ == 'D': n = 1 # Default move left 1 if len(args): n = args[0] line_pos -= n if line_pos < 0: line_pos = 0 assert(line_pos >= 0) elif typ == 'C': # Move right n = 1 # Default move 1 if len(args): n = args[0] line_pos += n if line_pos > len(lines_out[cur_line])-1: line_pos = len(lines_out)-1 assert(line_pos >= 0) elif typ == 'J': # Optional arg 0, 1, 2 n = 0 # default if len(args): n = args[0] if n == 0: # clear down lines_out = lines_out[:cur_line+1] elif n == 1: # clear up lines_out = lines_out[cur_line:] elif n == 2: # clear everything lines_out = [""] cur_line = 0 line_pos = 0 store_ptr = (0, 0) elif typ == 'K': # Optional arg 0, 1, 2 n = 0 # default if len(args): n = args[0] # HURISTIC-y hack: linux loves to have a line 123456 then do K(0) 6\r\n # so if the next line is text and the text is [eol]\r\n where [eol] matches the end of this # line - align cur_pos if len(reformatted) > idx+1: # Have another message (next_typ, next_args) = reformatted[idx+1] if next_typ == 'text': if '\r\n' in next_args[0]: next_lines = next_args[0].split("\r\n") if lines_out[cur_line].strip().endswith(next_lines[0]): # Its the buggy case. Just align it such that we clear the text # that's about to get echoed line_pos = line_pos - len(next_lines[0]) if n == 0: # clear right of cursor lines_out[cur_line] = lines_out[cur_line][:line_pos] elif n == 1: # clear left of cursor lines_out[cur_line] = (" "*line_pos)+lines_out[cur_line][line_pos:] elif n == 2: # clear whole line lines_out[cur_line] = " "*len(lines_out[cur_line]) elif typ == 'H': n = args[0]-1 m = args[1]-1 cur_line = n line_pos = m while cur_line >= len(lines_out): lines_out.append("") while line_pos > len(lines_out[cur_line]): lines_out[cur_line] += " " elif typ == 'T': # Scroll window down pass elif typ == 'S': # Scroll window up pass elif typ == 's': store_ptr = (cur_line, line_pos) elif typ == 'u': (cur_line, line_pos) = store_ptr elif typ == 'm': # alter character attributes - just ignore pass else: raise ValueError(f"Unsupported ANSI command {typ}") #_dump(lines_out, cur_line, line_pos) # Done processing - update variables self.prior_lines = lines_out[:-1] # Strings self.current_line = bytearray() # Bytearray if len(lines_out[-1].strip()): self.current_line.append(lines_out[-1]) def update_expectation(self, expectation)-
Expand source code
def update_expectation(self, expectation): if isinstance(expectation, bytes): expectation = expectation.decode() self.last_prompt = expectation # approximation self.expectation_re = re.compile(expectation) self.expectation_ends_re = re.compile(r'(.*)' + expectation)
class TimeoutExpired (*args, **kwargs)-
Common base class for all non-exit exceptions.
Expand source code
class TimeoutExpired(Exception): passAncestors
- builtins.Exception
- builtins.BaseException