Win PTY implementation for ssh connections in windows

Author: ofiriluz

ark_sdk_python/common/connections/ssh/ark_pty_ssh_win_connection.py

@@ -0,0 +1,233 @@
+import os
+import re
+import subprocess
+import threading
+import time
+from shutil import which
+from typing import Any, Final, Optional
+
+from overrides import overrides
+
+from ark_sdk_python.common.connections.ark_connection import ArkConnection
+from ark_sdk_python.common.connections.ssh.ark_ssh_connection import SSH_PORT
+from ark_sdk_python.models import ArkException
+from ark_sdk_python.models.common.connections import ArkConnectionCommand, ArkConnectionDetails, ArkConnectionResult
+
+
+class ArkPTYSSHWinConnection(ArkConnection):
+    __ANSI_COLOR_STRIPPER: Final[Any] = re.compile(r'\x1b[^m]*m|\x1b\[\?2004[hl]')
+    __ANSI_ESCAPE_STRIPPER: Final[Any] = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
+    __PROMPT_REGEX: Final[Any] = re.compile(r"[#$]")
+    __PASSWORD_PROMPT_REGEX: Final[Any] = re.compile(
+        r"(?i)(?:password:)|(?:passphrase for key)|(?:^Please enter your password[\s\S]+.+ >:$)"
+    )
+    __EXIT_MAKRER: Final[str] = '__ARK_RC_MARK__'
+    __DEFAULT_COL_SIZE: Final[int] = 80
+    __DEFAULT_ROW_SIZE: Final[int] = 24
+    __DEFAULT_NONEAGER_TIMEOEUT: Final[float] = 1.0
+    __DEFAULT_PROMPT_OVERALL_TIMEOUT: Final[float] = 20.0
+
+    def __init__(self):
+        super().__init__()
+        self.__is_connected: bool = False
+        self.__is_suspended: bool = False
+        self.__pty: Optional[Any] = None
+        self.__output_lock: threading.Lock = threading.Lock()
+        self.__buffer: str = ''
+
+    def __strip_ansi(self, input: str) -> str:
+        input = input.strip()
+        input = ArkPTYSSHWinConnection.__ANSI_COLOR_STRIPPER.sub('', input)
+        input = ArkPTYSSHWinConnection.__ANSI_ESCAPE_STRIPPER.sub('', input)
+        return input
+
+    def __reset_buffer(self) -> None:
+        with self.__output_lock:
+            self.__buffer = ''
+
+    def __read_until_latest_prompt(
+        self,
+        password: Optional[str] = None,
+        noneager_timeout: float = __DEFAULT_NONEAGER_TIMEOEUT,
+        overall_timeout: float = __DEFAULT_PROMPT_OVERALL_TIMEOUT,
+        login_prompt: bool = False,
+        expected_prompt: Any = __PROMPT_REGEX,
+    ) -> None:
+        buffer = ''
+        prompt_found_at = -1
+        start_time = time.time()
+        noneager_start_time = start_time
+
+        while True:
+            char = self.__pty.read(1)
+            if not char:
+                if ArkPTYSSHWinConnection.__PASSWORD_PROMPT_REGEX.search(self.__strip_ansi(buffer)) and not password and login_prompt:
+                    raise RuntimeError(f'Password prompt with no password given [{buffer}]')
+                if (time.time() - noneager_start_time > noneager_timeout) and prompt_found_at != -1:
+                    break
+                if (time.time() - start_time) > overall_timeout and prompt_found_at == -1:
+                    raise RuntimeError(f'Timeout while waiting for prompt [{buffer}]')
+                time.sleep(0.05)
+                continue
+
+            buffer += char
+
+            with self.__output_lock:
+                self.__buffer += char
+
+            if ArkPTYSSHWinConnection.__PASSWORD_PROMPT_REGEX.search(self.__strip_ansi(buffer)) and login_prompt:
+                if not password:
+                    raise RuntimeError(f'Password prompt with no password given [{buffer}]')
+                self.__pty.write(password + '\n')
+                buffer = ""
+                prompt_found_at = -1
+                continue
+
+            if expected_prompt.search(buffer):
+                prompt_found_at = len(buffer)
+
+            noneager_start_time = time.time()
+
+        if prompt_found_at != -1:
+            with self.__output_lock:
+                self.__buffer = buffer
+
+    @overrides
+    def connect(self, connection_details: ArkConnectionDetails) -> None:
+        """
+        Performs SSH connection with given details or keys
+        Saves the ssh session to be used for command executions
+        Done using windows pty
+
+        Args:
+            connection_details (ArkConnectionDetails): _description_
+
+        Raises:
+            ArkException: _description_
+        """
+        import winpty
+
+        if self.__is_connected:
+            return
+
+        address = connection_details.address
+        user = connection_details.credentials.user
+        port = connection_details.port or SSH_PORT
+        password = None
+        key_path = None
+        if connection_details.credentials.password:
+            password = connection_details.credentials.password.get_secret_value()
+        elif connection_details.credentials.private_key_filepath:
+            key_path = connection_details.credentials.private_key_filepath
+
+        ssh_args = [f'{user}@{address}', '-p', str(port), '-o', 'StrictHostKeyChecking=no']
+        if key_path:
+            ssh_args.extend(['-i', str(key_path)])
+        try:
+            self.__pty = winpty.PTY(
+                cols=ArkPTYSSHWinConnection.__DEFAULT_COL_SIZE,
+                rows=ArkPTYSSHWinConnection.__DEFAULT_ROW_SIZE,
+                backend=winpty.enums.Backend.WinPTY,
+            )
+            ssh_full_cmd = which('ssh', path=os.environ.get('PATH', os.defpath))
+            self.__pty.spawn(ssh_full_cmd, cmdline=' ' + subprocess.list2cmdline(ssh_args))
+            self.__read_until_latest_prompt(password, login_prompt=True)
+            self.__is_connected = True
+        except Exception as ex:
+            raise ArkException(f'Failed to ssh connect [{str(ex)}]') from ex
+
+    @overrides
+    def disconnect(self) -> None:
+        """
+        Disconnects the ssh session
+        """
+        if self.__pty:
+            self.__pty.write('exit\n')
+            self.__pty = None
+        self.__is_connected = False
+        self.__is_suspended = False
+
+    @overrides
+    def suspend_connection(self) -> None:
+        """
+        Suspends execution of ssh commands
+        """
+        self.__is_suspended = True
+
+    @overrides
+    def restore_connection(self) -> None:
+        """
+        Restores execution of ssh commands
+        """
+        self.__is_suspended = False
+
+    @overrides
+    def is_suspended(self) -> bool:
+        """
+        Checks whether ssh commands can be executed or not
+
+        Returns:
+            bool: _description_
+        """
+        return self.__is_suspended
+
+    @overrides
+    def is_connected(self) -> bool:
+        """
+        Checks whether theres a ssh session connected
+
+        Returns:
+            bool: _description_
+        """
+        return self.__is_connected
+
+    @overrides
+    def run_command(self, command: ArkConnectionCommand) -> ArkConnectionResult:
+        """
+        Runs a command over ssh session via pty, returning the result accordingly
+        stderr is not supported, only stdout is returned, and it'll contain everything including stderr
+
+        Args:
+            command (ArkConnectionCommand): _description_
+
+        Raises:
+            ArkException: _description_
+
+        Returns:
+            ArkConnectionResult: _description_
+        """
+        if not self.__is_connected or self.__is_suspended:
+            raise ArkException('Cannot run command while not connected or suspended')
+
+        self._logger.debug(f'Running command [{command.command}]')
+
+        self.__reset_buffer()
+        self.__pty.write(command.command + "\n")
+        self.__read_until_latest_prompt()
+
+        with self.__output_lock:
+            stdout = self.__buffer.strip()
+
+        self.__reset_buffer()
+        self.__pty.write(f'echo {ArkPTYSSHWinConnection.__EXIT_MAKRER}_$?;\n')
+        self.__read_until_latest_prompt()
+
+        with self.__output_lock:
+            lines = self.__buffer.strip().splitlines()
+
+        rc = None
+        for line in lines:
+            line = self.__strip_ansi(line)
+            if line.startswith(f'{ArkPTYSSHWinConnection.__EXIT_MAKRER}_'):
+                try:
+                    rc = int(line[len(ArkPTYSSHWinConnection.__EXIT_MAKRER) + 1 :])  # +1 for the underscore
+                    break
+                except ValueError:
+                    continue
+        if rc is None:
+            raise ArkException(f"Failed to parse exit code from output - [{self.__buffer}]")
+        if rc != command.expected_rc and command.raise_on_error:
+            raise ArkException(f'Failed to execute command [{command.command}] - [{rc}] - [{stdout}]')
+        self._logger.debug(f'Command rc: [{rc}]')
+        self._logger.debug(f'Command stdout: [{stdout}]')
+        return ArkConnectionResult(stdout=stdout, rc=rc)

ark_sdk_python/examples/validate_ssh_connection.py

@@ -1,41 +1,54 @@
 import argparse
+import platform
 import tempfile
 
 from ark_sdk_python.auth import ArkISPAuth
 from ark_sdk_python.common.ark_jwt_utils import ArkJWTUtils
 from ark_sdk_python.common.connections.ssh.ark_pty_ssh_connection import ArkPTYSSHConnection
+from ark_sdk_python.common.connections.ssh.ark_pty_ssh_win_connection import ArkPTYSSHWinConnection
 from ark_sdk_python.models.auth import ArkAuthMethod, ArkAuthProfile, ArkSecret
-from ark_sdk_python.models.auth.ark_auth_method import IdentityServiceUserArkAuthMethodSettings
+from ark_sdk_python.models.auth.ark_auth_method import (
+    ArkAuthMethodsDescriptionMap,
+    IdentityArkAuthMethodSettings,
+    IdentityServiceUserArkAuthMethodSettings,
+)
 from ark_sdk_python.models.common.connections.ark_connection_command import ArkConnectionCommand
 from ark_sdk_python.models.common.connections.ark_connection_credentials import ArkConnectionCredentials
 from ark_sdk_python.models.common.connections.ark_connection_details import ArkConnectionDetails
 from ark_sdk_python.models.services.sia.sso.ark_sia_sso_get_ssh_key import ArkSIASSOGetSSHKey
 from ark_sdk_python.services.sia.sso import ArkSIASSOService
 
 
-def login_to_identity_security_platform(service_user: str, service_token: str, application_name: str) -> ArkISPAuth:
+def login_to_identity_security_platform(user: str, secret: str, is_service_user: bool, application_name: str) -> ArkISPAuth:
     """
-    This will perform login to the tenant with the given service user credentials
+    This will perform login to the tenant with the given user credentials
+    Where the user is normal or service user based on the boolean
     Caching the logged in token is disabled and will perform a full login on each call
     Will return an identity security platform authenticated class
 
     Args:
-        service_user (str): _description_
-        service_token (str): _description_
+        user (str): _description_
+        secret (str): _description_
         application_name (str): _description_
 
     Returns:
         ArkISPAuth: _description_
     """
-    print('Logging in to the tenant')
     isp_auth = ArkISPAuth(cache_authentication=False)
+    auth_method = ArkAuthMethod.IdentityServiceUser if is_service_user else ArkAuthMethod.Identity
+    auth_methods_settings = (
+        IdentityServiceUserArkAuthMethodSettings(identity_authorization_application=application_name)
+        if is_service_user
+        else IdentityArkAuthMethodSettings()
+    )
+    print(f'Logging in to the tenant with [{ArkAuthMethodsDescriptionMap[auth_method]}] user type and user [{user}]')
     isp_auth.authenticate(
         auth_profile=ArkAuthProfile(
-            username=service_user,
-            auth_method=ArkAuthMethod.IdentityServiceUser,
-            auth_method_settings=IdentityServiceUserArkAuthMethodSettings(identity_authorization_application=application_name),
+            username=user,
+            auth_method=auth_method,
+            auth_method_settings=auth_methods_settings,
         ),
-        secret=ArkSecret(secret=service_token),
+        secret=ArkSecret(secret=secret),
     )
     print('Logged in successfully')
     return isp_auth
@@ -92,7 +105,10 @@ def connect_and_validate_connection(ssh_key_path: str, proxy_address: str, conne
         command (str): _description_
     """
     print(f'Connecting to proxy [{proxy_address}] and connection string [{connection_string}]')
-    ssh_connection = ArkPTYSSHConnection()
+    if platform.system() == 'Windows':
+        ssh_connection = ArkPTYSSHWinConnection()
+    else:
+        ssh_connection = ArkPTYSSHConnection()
     ssh_connection.connect(
         ArkConnectionDetails(
             address=proxy_address,
@@ -112,8 +128,9 @@ def connect_and_validate_connection(ssh_key_path: str, proxy_address: str, conne
 if __name__ == '__main__':
     # Construct an argument parser for CLI parameters for the script
     parser = argparse.ArgumentParser()
-    parser.add_argument('--service-user', required=True, help='Service user to login and perform the operation with')
-    parser.add_argument('--service-token', required=True, help='Service user token to use for logging in and connecting')
+    parser.add_argument('--user', required=True, help='User to login and perform the operation with')
+    parser.add_argument('--secret', required=True, help='User secret to use for logging in and connecting')
+    parser.add_argument('--is-service-user', action='store_true', help='Whether this user is a service user and not a normal user')
     parser.add_argument(
         '--service-application',
         default='__idaptive_cybr_user_oidc',
@@ -139,7 +156,7 @@ def connect_and_validate_connection(ssh_key_path: str, proxy_address: str, conne
         # - Generate an MFA Caching SSH key
         # - Construct the ssh proxy address
         # - Connect in SSH to the proxy / target with the MFA Caching SSH Key and perform the command
-        isp_auth = login_to_identity_security_platform(args.service_user, args.service_token, args.service_application)
+        isp_auth = login_to_identity_security_platform(args.user, args.secret, args.is_service_user, args.service_application)
         ssh_key_path = generate_mfa_caching_ssh_key(isp_auth, temp_folder)
         proxy_address = construct_ssh_proxy_address(isp_auth)
         connect_and_validate_connection(ssh_key_path, proxy_address, args.connection_string, args.test_command)