@@ -581,24 +581,40 @@ class Machine:
581581 + "'{}' but it is in state ‘{}’".format(require_state, state)
582582 )
583583
584− def execute(self, command: str) -> Tuple[int, str]:
584+ def _next_newline_closed_block_from_shell(self) -> str:
585+ assert self.shell
586+ output_buffer = []
587+ while True:
588+ # This receives up to 4096 bytes from the socket
589+ chunk = self.shell.recv(4096)
590+ if not chunk:
591+ # Probably a broken pipe, return the output we have
592+ break
593+
594+ decoded = chunk.decode()
595+ output_buffer += [decoded]
596+ if decoded[-1] == "\n":
597+ break
598+ return "".join(output_buffer)
599+
600+ def execute(self, command: str, check_return: bool = True) -> Tuple[int, str]:
585601 self.connect()
586602
587− out_command = "( set -euo pipefail; {} ); echo '|!=EOF' $?\n".format(command)
603+ out_command = f"( set -euo pipefail; {command} ) | (base64 --wrap 0; echo)\n"
588604 assert self.shell
589605 self.shell.send(out_command.encode())
590606
591− output = ""
592− status_code_pattern = re.compile(r"(.*)\|\!=EOF\s+(\d+)")
607+ # Get the output
608+ output = base64.b64decode(self._next_newline_closed_block_from_shell())
593609
594− while True:
595− chunk = self.shell.recv(4096).decode(errors="ignore")
596− match = status_code_pattern.match(chunk)
597− if match:
598− output += match[1]
599− status_code = int(match[2])
600− return (status_code, output)
601− output += chunk
610+ if not check_return:
611+ return (-1, output.decode())
612+
613+ # Get the return code
614+ self.shell.send("echo ${PIPESTATUS[0]}\n".encode())
615+ rc = int(self._next_newline_closed_block_from_shell().strip())
616+
617+ return (rc, output.decode())
602618
603619 def shell_interact(self) -> None:
604620 """Allows you to interact with the guest shell