diff --git a/src/targetcli/targetcli_shell.py b/src/targetcli/targetcli_shell.py index df4f7b0..3349655 100644 --- a/src/targetcli/targetcli_shell.py +++ b/src/targetcli/targetcli_shell.py @@ -40,6 +40,7 @@ 'cd', 'pwd', 'ls', 'set', 'get', 'help', 'refresh', 'status', 'clearconfig', 'restoreconfig', 'saveconfig', 'exit'] +NUM_PARTS = 2 class TargetCLI(ConfigShell): default_prefs = {'color_path': 'magenta', 'color_command': 'cyan', @@ -168,17 +169,45 @@ def call_daemon(shell, req, interactive): # get the actual data in chunks output = "" path = "" + raise_error = False while amount_received < amount_expected: data = sock.recv(1024) data = data.decode() amount_received += len(data) output += data - if get_pwd: + # Check if output is in new format (STDOUT/STDERR) + if "STDOUT:" in output and "STDERR:" in output: + parts = output.split("STDERR:") + if len(parts) == NUM_PARTS: + stdout_part = parts[0].split("STDOUT:")[1].strip() + stderr_content = parts[1].strip() + stdout_content = stdout_part + + if get_pwd: + output_split = stdout_content.splitlines() + lines = len(output_split) + for i in range(lines): + if i == lines-1: + path = str(output_split[i]) + else: + print(str(output_split[i]), end="\n") + + else: + # Print stdout content + if stdout_content: + print(stdout_content, end="\n") + # Print stderr content + if stderr_content: + print(stderr_content, end="\n", file=sys.stderr) + raise_error = True + + # Handle old format (plain text) + elif get_pwd: output_split = output.splitlines() lines = len(output_split) for i in range(lines): - if i == lines - 1: + if i == lines-1: path = str(output_split[i]) else: print(str(output_split[i]), end="\n") @@ -188,6 +217,9 @@ def call_daemon(shell, req, interactive): sock.send(b'-END@OF@DATA-') sock.close() + if raise_error and not interactive: + sys.exit(1) + return path def switch_to_daemon(shell, interactive): diff --git a/src/targetcli/targetclid.py b/src/targetcli/targetclid.py index 32c582a..0ea90c0 100644 --- a/src/targetcli/targetclid.py +++ b/src/targetcli/targetclid.py @@ -153,28 +153,51 @@ def client_thread(self, connection): connection.close() still_listen = False else: - self.con._stdout = self.con._stderr = f = tempfile.NamedTemporaryFile(mode='w', delete=False) + # Create separate tempfiles for stdout and stderr + stdout_file = tempfile.NamedTemporaryFile(mode='w', delete=False) + stderr_file = tempfile.NamedTemporaryFile(mode='w', delete=False) + + # Save original stdout/stderr + orig_stdout = self.con._stdout + orig_stderr = self.con._stderr + + # Redirect to tempfiles + self.con._stdout = stdout_file + self.con._stderr = stderr_file + try: # extract multiple commands delimited with '%' list_data = data.decode().split('%') for cmd in list_data: self.shell.run_cmdline(cmd) except Exception as e: - print(str(e), file=f) # push error to stream - - # Restore - self.con._stdout = self.con_stdout_ - self.con._stderr = self.con_stderr_ - f.close() - - with open(f.name) as f: - output = f.read() - var = struct.pack('i', len(output)) - connection.sendall(var) # length of string - if len(output): - connection.sendall(output.encode()) # actual string - - Path(f.name).unlink() + print(str(e), file=stderr_file) + # Restore original stdout/stderr + self.con._stdout = orig_stdout + self.con._stderr = orig_stderr + + # Close tempfiles + stdout_file.close() + stderr_file.close() + + # Read contents + with open(stdout_file.name) as f: + stdout_content = f.read() + with open(stderr_file.name) as f: + stderr_content = f.read() + + # Create combined output + combined_output = f"STDOUT: {stdout_content}\nSTDERR: {stderr_content}" + + # Send combined output + var = struct.pack('i', len(combined_output)) + connection.sendall(var) # length of string + if len(combined_output): + connection.sendall(combined_output.encode()) # actual string + + # Cleanup tempfiles + Path(stdout_file.name).unlink() + Path(stderr_file.name).unlink() def usage():