| 1 | import asyncio |
| 2 | import paramiko |
| 3 | import time |
| 4 | import re |
| 5 | from typing import Tuple |
| 6 | from helpers.log import Log |
| 7 | from helpers.print_style import PrintStyle |
| 8 | # from helpers.strings import calculate_valid_match_lengths |
| 9 | |
| 10 | # Paramiko caches this optional import error; its traceback retains our tool-loading stack. |
| 11 | if paramiko.config.invoke_import_error: |
| 12 | paramiko.config.invoke_import_error.__traceback__ = None |
| 13 | |
| 14 | |
| 15 | # Injected into every new SSH shell to keep it safe for non-interactive use. |
| 16 | # Pagers (more/less) would otherwise block forever waiting for input that |
| 17 | # never arrives and spin at 100% CPU; see issue #1697. |
| 18 | PAGER_DISABLE_COMMAND = "export GIT_PAGER=cat; export PAGER=cat" |
| 19 | |
| 20 | |
| 21 | class SSHInteractiveSession: |
| 22 | |
| 23 | # end_comment = "# @@==>> SSHInteractiveSession End-of-Command <<==@@" |
| 24 | # ps1_label = "SSHInteractiveSession CLI>" |
| 25 | |
| 26 | def __init__( |
| 27 | self, logger: Log, hostname: str, port: int, username: str, password: str, cwd: str|None = None |
| 28 | ): |
| 29 | self.logger = logger |
| 30 | self.hostname = hostname |
| 31 | self.port = port |
| 32 | self.username = username |
| 33 | self.password = password |
| 34 | self.client = paramiko.SSHClient() |
| 35 | self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) |
| 36 | self.shell = None |
| 37 | self.full_output = b"" |
| 38 | self.last_command = b"" |
| 39 | self.trimmed_command_length = 0 # Initialize trimmed_command_length |
| 40 | self.cwd = cwd |
| 41 | self._exit_code: int | None = None |
| 42 | |
| 43 | def __del__(self): |
| 44 | for resource in (getattr(self, "shell", None), getattr(self, "client", None)): |
| 45 | try: |
| 46 | if resource: |
| 47 | resource.close() |
| 48 | except Exception: |
| 49 | pass |
| 50 | |
| 51 | async def connect(self, keepalive_interval: int = 5): |
| 52 | """ |
| 53 | Establish the SSH connection and start an interactive shell. |
| 54 | |
| 55 | Parameters |
| 56 | ---------- |
| 57 | keepalive_interval : int |
| 58 | Interval in **seconds** between keep-alive packets sent by Paramiko. |
| 59 | A value ≤ 0 disables Paramiko’s keep-alive feature. |
| 60 | """ |
| 61 | errors = 0 |
| 62 | while True: |
| 63 | try: |
| 64 | # --- establish TCP/SSH session --------------------------------- |
| 65 | self.client.connect( |
| 66 | self.hostname, |
| 67 | self.port, |
| 68 | self.username, |
| 69 | self.password, |
| 70 | allow_agent=False, |
| 71 | look_for_keys=False, |
| 72 | ) |
| 73 | |
| 74 | # --------- NEW: enable transport-level keep-alives ------------- |
| 75 | transport = self.client.get_transport() |
| 76 | if transport and keepalive_interval > 0: |
| 77 | # sends an SSH_MSG_IGNORE every <keepalive_interval> seconds |
| 78 | transport.set_keepalive(keepalive_interval) |
| 79 | # ---------------------------------------------------------------- |
| 80 | |
| 81 | # invoke interactive shell |
| 82 | self.shell = self.client.invoke_shell(width=100, height=50) |
| 83 | self._exit_code = None |
| 84 | |
| 85 | # disable systemd/OSC prompt metadata and disable local echo |
| 86 | initial_command = f"unset PROMPT_COMMAND PS0; stty -echo; {PAGER_DISABLE_COMMAND}" |
| 87 | if self.cwd: |
| 88 | initial_command = f"cd {self.cwd}; {initial_command}" |
| 89 | self.shell.send(f"{initial_command}\n".encode()) |
| 90 | |
| 91 | # wait for initial prompt/output to settle |
| 92 | while True: |
| 93 | full, part = await self.read_output() |
| 94 | if full and not part: |
| 95 | return |
| 96 | time.sleep(0.1) |
| 97 | |
| 98 | except Exception as e: |
| 99 | errors += 1 |
| 100 | if errors < 3: |
| 101 | PrintStyle.standard(f"SSH Connection attempt {errors}...") |
| 102 | self.logger.log( |
| 103 | type="info", |
| 104 | content=f"SSH Connection attempt {errors}...", |
| 105 | ) |
| 106 | time.sleep(5) |
| 107 | else: |
| 108 | raise e |
| 109 | |
| 110 | async def close(self): |
| 111 | if self.shell: |
| 112 | self.shell.close() |
| 113 | if self.client: |
| 114 | self.client.close() |
| 115 | |
| 116 | async def send_command(self, command: str): |
| 117 | if not self.shell: |
| 118 | raise Exception("Shell not connected") |
| 119 | self.full_output = b"" |
| 120 | # if len(command) > 10: # if command is long, add end_comment to split output |
| 121 | # command = (command + " \\\n" +SSHInteractiveSession.end_comment + "\n") |
| 122 | # else: |
| 123 | command = command + "\n" |
| 124 | self.last_command = command.encode() |
| 125 | self.trimmed_command_length = 0 |
| 126 | self.shell.send(self.last_command) |
| 127 | |
| 128 | def is_terminated(self) -> bool: |
| 129 | if not self.shell: |
| 130 | return True |
| 131 | try: |
| 132 | transport = self.client.get_transport() |
| 133 | if not transport or not transport.is_active(): |
| 134 | return True |
| 135 | return self.shell.closed or self.shell.exit_status_ready() |
| 136 | except Exception: |
| 137 | return True |
| 138 | |
| 139 | def get_exit_code(self) -> int | None: |
| 140 | if self._exit_code is not None: |
| 141 | return self._exit_code |
| 142 | try: |
| 143 | if self.shell and self.shell.exit_status_ready(): |
| 144 | self._exit_code = self.shell.recv_exit_status() |
| 145 | except Exception: |
| 146 | return None |
| 147 | return self._exit_code |
| 148 | |
| 149 | async def read_output( |
| 150 | self, timeout: float = 0, reset_full_output: bool = False |
| 151 | ) -> Tuple[str, str]: |
| 152 | if not self.shell: |
| 153 | raise Exception("Shell not connected") |
| 154 | |
| 155 | if reset_full_output: |
| 156 | self.full_output = b"" |
| 157 | partial_output = b"" |
| 158 | leftover = b"" |
| 159 | start_time = time.time() |
| 160 | |
| 161 | while self.shell.recv_ready() and ( |
| 162 | timeout <= 0 or time.time() - start_time < timeout |
| 163 | ): |
| 164 | |
| 165 | # data = self.shell.recv(1024) |
| 166 | data = self.receive_bytes() |
| 167 | |
| 168 | # # Trim own command from output |
| 169 | # if ( |
| 170 | # self.last_command |
| 171 | # and len(self.last_command) > self.trimmed_command_length |
| 172 | # ): |
| 173 | # command_to_trim = self.last_command[self.trimmed_command_length :] |
| 174 | # data_to_trim = leftover + data |
| 175 | |
| 176 | # trim_com, trim_out = calculate_valid_match_lengths( |
| 177 | # command_to_trim, |
| 178 | # data_to_trim, |
| 179 | # deviation_threshold=8, |
| 180 | # deviation_reset=2, |
| 181 | # ignore_patterns=[ |
| 182 | # rb"\x1b\[\?\d{4}[a-zA-Z](?:> )?", # ANSI escape sequences |
| 183 | # rb"\r", # Carriage return |
| 184 | # rb">\s", # Greater-than symbol |
| 185 | # ], |
| 186 | # debug=False, |
| 187 | # ) |
| 188 | |
| 189 | # leftover = b"" |
| 190 | # if trim_com > 0 and trim_out > 0: |
| 191 | # data = data_to_trim[trim_out:] |
| 192 | # leftover = data |
| 193 | # self.trimmed_command_length += trim_com |
| 194 | |
| 195 | partial_output += data |
| 196 | self.full_output += data |
| 197 | await asyncio.sleep(0.1) # Prevent busy waiting |
| 198 | |
| 199 | # Decode once at the end |
| 200 | decoded_partial_output = partial_output.decode("utf-8", errors="replace") |
| 201 | decoded_full_output = self.full_output.decode("utf-8", errors="replace") |
| 202 | |
| 203 | decoded_partial_output = clean_string(decoded_partial_output) |
| 204 | decoded_full_output = clean_string(decoded_full_output) |
| 205 | |
| 206 | return decoded_full_output, decoded_partial_output |
| 207 | |
| 208 | def receive_bytes(self, num_bytes=1024): |
| 209 | if not self.shell: |
| 210 | raise Exception("Shell not connected") |
| 211 | # Receive initial chunk of data |
| 212 | shell = self.shell |
| 213 | data = self.shell.recv(num_bytes) |
| 214 | |
| 215 | # Helper function to ensure that we receive exactly `num_bytes` |
| 216 | def recv_all(num_bytes): |
| 217 | data = b"" |
| 218 | while len(data) < num_bytes: |
| 219 | chunk = shell.recv(num_bytes - len(data)) |
| 220 | if not chunk: |
| 221 | break # Connection might be closed or no more data |
| 222 | data += chunk |
| 223 | return data |
| 224 | |
| 225 | # Check if the last byte(s) form an incomplete multi-byte UTF-8 sequence |
| 226 | if len(data) > 0: |
| 227 | last_byte = data[-1] |
| 228 | |
| 229 | # Check if the last byte is part of a multi-byte UTF-8 sequence (continuation byte) |
| 230 | if (last_byte & 0b11000000) == 0b10000000: # It's a continuation byte |
| 231 | # Now, find the start of this sequence by checking earlier bytes |
| 232 | for i in range( |
| 233 | 2, 5 |
| 234 | ): # Look back up to 4 bytes (since UTF-8 is up to 4 bytes long) |
| 235 | if len(data) - i < 0: |
| 236 | break |
| 237 | byte = data[-i] |
| 238 | |
| 239 | # Detect the leading byte of a multi-byte sequence |
| 240 | if (byte & 0b11100000) == 0b11000000: # 2-byte sequence (110xxxxx) |
| 241 | data += recv_all(1) # Need 1 more byte to complete |
| 242 | break |
| 243 | elif ( |
| 244 | byte & 0b11110000 |
| 245 | ) == 0b11100000: # 3-byte sequence (1110xxxx) |
| 246 | data += recv_all(2) # Need 2 more bytes to complete |
| 247 | break |
| 248 | elif ( |
| 249 | byte & 0b11111000 |
| 250 | ) == 0b11110000: # 4-byte sequence (11110xxx) |
| 251 | data += recv_all(3) # Need 3 more bytes to complete |
| 252 | break |
| 253 | |
| 254 | return data |
| 255 | |
| 256 | def clean_string(input_string): |
| 257 | # Remove ANSI escape codes |
| 258 | ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") |
| 259 | cleaned = ansi_escape.sub("", input_string) |
| 260 | |
| 261 | # remove null bytes |
| 262 | cleaned = cleaned.replace("\x00", "") |
| 263 | |
| 264 | # remove ipython \r\r\n> sequences from the start |
| 265 | cleaned = re.sub(r'^[ \r]*(?:\r*\n>[ \r]*)*', '', cleaned) |
| 266 | # also remove any amount of '> ' sequences from the start |
| 267 | cleaned = re.sub(r'^(>\s*)+', '', cleaned) |
| 268 | |
| 269 | # Replace '\r\n' with '\n' |
| 270 | cleaned = cleaned.replace("\r\n", "\n") |
| 271 | |
| 272 | # remove leading \r and spaces |
| 273 | cleaned = cleaned.lstrip("\r ") |
| 274 | |
| 275 | # Split the string by newline characters to process each segment separately |
| 276 | lines = cleaned.split("\n") |
| 277 | |
| 278 | for i in range(len(lines)): |
| 279 | # Handle carriage returns '\r' by splitting and taking the last part |
| 280 | parts = [part for part in lines[i].split("\r") if part.strip()] |
| 281 | if parts: |
| 282 | lines[i] = parts[ |
| 283 | -1 |
| 284 | ].rstrip() # Overwrite with the last part after the last '\r' |
| 285 | |
| 286 | return "\n".join(lines) |