diff --git a/bale/interfaces/cli.py b/bale/interfaces/cli.py index 02cede1..39d2a90 100644 --- a/bale/interfaces/cli.py +++ b/bale/interfaces/cli.py @@ -34,6 +34,7 @@ class Cli: self.stderr: List[str] = [] self._terminate: asyncio.Event = asyncio.Event() self._busy: bool = False + self._truncated: bool = False self.prefix_line: str = "" self._stdout_terminals: List[Terminal] = [] self._stderr_terminals: List[Terminal] = [] @@ -70,8 +71,11 @@ class Cli: else: break - async def _controller(self, process: Process) -> None: + async def _controller(self, process: Process, max_output_lines) -> None: while process.returncode is None: + if max_output_lines > 0 and len(self.stderr) + len(self.stdout) > max_output_lines: + self._truncated = True + process.terminate() if self._terminate.is_set(): process.terminate() try: @@ -83,7 +87,7 @@ class Cli: def terminate(self) -> None: self._terminate.set() - async def execute(self, command: str) -> Result: + async def execute(self, command: str, max_output_lines: int = 0) -> Result: self._busy = True c = shlex.split(command, posix=False) try: @@ -92,13 +96,14 @@ class Cli: self.stdout.clear() self.stderr.clear() self._terminate.clear() + self._truncated = False terminated = False now = datetime.now().strftime("%Y/%m/%d %H:%M:%S") self.prefix_line = f"<{now}> {command}\n" for terminal in self._stdout_terminals: terminal.call_terminal_method("write", "\n" + self.prefix_line) await asyncio.gather( - self._controller(process=process), + self._controller(process=process, max_output_lines=max_output_lines), self._read_stdout(stream=process.stdout), self._read_stderr(stream=process.stderr), ) @@ -110,7 +115,9 @@ class Cli: finally: self._terminate.clear() self._busy = False - return Result(command=command, return_code=process.returncode, stdout_lines=self.stdout.copy(), stderr_lines=self.stderr.copy(), terminated=terminated) + return Result( + command=command, return_code=process.returncode, stdout_lines=self.stdout.copy(), stderr_lines=self.stderr.copy(), terminated=terminated, truncated=self._truncated + ) async def shell(self, command: str) -> Result: self._busy = True diff --git a/bale/interfaces/ssh.py b/bale/interfaces/ssh.py index 3581325..44b87ca 100644 --- a/bale/interfaces/ssh.py +++ b/bale/interfaces/ssh.py @@ -90,10 +90,10 @@ class Ssh(Cli): del self._config[self.host] self.write_config() - async def execute(self, command: str) -> Result: + async def execute(self, command: str, max_output_lines: int = 0) -> Result: self._base_cmd = f"{'' if self.use_key else f'sshpass -p {self.password} '} ssh -F {self._config_path} {self.host}" self._full_cmd = f"{self._base_cmd} {command}" - return await super().execute(self._full_cmd) + return await super().execute(self._full_cmd, max_output_lines) async def send_key(self) -> Result: await get_public_key(self._raw_path) diff --git a/bale/interfaces/zfs.py b/bale/interfaces/zfs.py index 4c50aeb..32c412e 100644 --- a/bale/interfaces/zfs.py +++ b/bale/interfaces/zfs.py @@ -81,7 +81,7 @@ class Zfs: command = command if len(command) < 160 else command[:160] + "..." el.notify(command) - async def execute(self, command: str, notify: bool = True) -> Result: + async def execute(self, command: str, max_output_lines: int = 0, notify: bool = True) -> Result: if notify: self.notify(command) return Result(command=command) @@ -244,10 +244,10 @@ class Ssh(ssh.Ssh, Zfs): def notify(self, command: str): super().notify(f"<{self.host}> {command}") - async def execute(self, command: str, notify: bool = True) -> Result: + async def execute(self, command: str, max_output_lines: int = 0, notify: bool = True) -> Result: if notify: self.notify(command) - result = await super().execute(command) + result = await super().execute(command, max_output_lines) if result.stderr != "": el.notify(result.stderr, type="negative") result.name = self.host diff --git a/bale/result.py b/bale/result.py index 010467a..4330ff7 100644 --- a/bale/result.py +++ b/bale/result.py @@ -13,6 +13,7 @@ class Result: stdout_lines: List[str] = field(default_factory=list) stderr_lines: List[str] = field(default_factory=list) terminated: bool = False + truncated: bool = False data: Any = None trace: str = "" cached: bool = False