pypsrp 0.9.0rc1__py3-none-any.whl → 0.9.0rc2__py3-none-any.whl

pypsrp/_pool_manager.py

@@ -0,0 +1,142 @@
+# Copyright: (c) 2026, Jordan Borean (@jborean93) <jborean93@gmail.com>
+# MIT License (see LICENSE or https://opensource.org/licenses/MIT)
+
+from __future__ import annotations
+
+import contextvars
+import functools
+import types
+import typing as t
+
+import requests.adapters
+from requests.packages.urllib3.util.retry import Retry
+
+T = t.TypeVar("T")
+
+
+class NewConnectionDisabled(Exception):
+    """Raised when new connections are being made but have been disabled in the current context."""
+
+
+class DisableNewConnectionsContext:
+    """Context manager to disable new connections from being created."""
+
+    __slots__ = "_contextvar_token"
+
+    _contextvar: t.ClassVar[contextvars.ContextVar] = contextvars.ContextVar("NewConnectionContext")
+    _contextvar_token: contextvars.Token
+
+    @classmethod
+    def current(cls) -> DisableNewConnectionsContext | None:
+        """Returns the current context or None if not set."""
+        try:
+            return cls._contextvar.get()
+        except LookupError:
+            return None
+
+    def __enter__(self) -> DisableNewConnectionsContext:
+        self._contextvar_token = self.__class__._contextvar.set(self)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> bool | None:
+        self.__class__._contextvar.reset(self._contextvar_token)
+        del self._contextvar_token
+
+        return None
+
+
+class HTTPSAdapterWithKeyPassword(requests.adapters.HTTPAdapter):
+
+    def __init__(
+        self,
+        *args: t.Any,
+        _pypsrp_key_password: str | None = None,
+        **kwargs: t.Any,
+    ) -> None:
+        self.__key_password = _pypsrp_key_password
+        super().__init__(*args, **kwargs)
+
+    def init_poolmanager(
+        self,
+        connections,
+        maxsize,
+        block=False,
+        **pool_kwargs,
+    ):
+        return super().init_poolmanager(
+            connections,
+            maxsize,
+            block,
+            key_password=self.__key_password,
+            **pool_kwargs,
+        )
+
+
+def _wrap(
+    func: t.Callable[..., T],
+    *,
+    before: t.Callable[[], None] | None = None,
+    after: t.Callable[[T], None] | None = None,
+) -> t.Callable[..., T]:
+    """Wraps a function with before and after callables."""
+
+    @functools.wraps(func)
+    def wrapper(*args: t.Any, **kwargs: t.Any) -> T:
+        if before:
+            before()
+
+        res = func(*args, **kwargs)
+
+        if after:
+            after(res)
+
+        return res
+
+    return wrapper
+
+
+def _raise_if_connections_disabled():
+    if DisableNewConnectionsContext.current() is not None:
+        raise NewConnectionDisabled()
+
+
+def _create_new_connection_with_failure(connection):
+    connection.connect = _wrap(connection.connect, before=_raise_if_connections_disabled)
+
+
+def _create_connection_pool(pool):
+    # This is called in the urllib3 pool when a new connection is needed. We
+    # wrap it so we can wrap the connect method when a new connection is
+    # created.
+    pool.ConnectionCls = _wrap(pool.ConnectionCls, after=_create_new_connection_with_failure)
+
+
+def create_request_adapter(  # type: ignore[no-any-unimported]  # requests does not have typing stubs for urllib3
+    *,
+    max_retries: Retry,
+    key_password: str | None = None,
+) -> requests.adapters.HTTPAdapter:
+    """Creates a HTTPAdapter with support for disabling new connections via context."""
+    adapter: requests.adapters.HTTPAdapter
+    if key_password is not None:
+        adapter = HTTPSAdapterWithKeyPassword(
+            max_retries=max_retries,
+            _pypsrp_key_password=key_password,
+        )
+    else:
+        adapter = requests.adapters.HTTPAdapter(max_retries=max_retries)
+
+    # pool_classes_by_scheme stores the urllib3 pool types used when creating
+    # the connection pool for http/https. We wrap the __init__ method so we can
+    # inject our custom ConnectionCls that wraps the connect method when
+    # created.
+    pool_classes = adapter.poolmanager.pool_classes_by_scheme
+    for scheme, pool_cls in list(pool_classes.items()):
+        pool_classes[scheme] = _wrap(pool_cls, after=_create_connection_pool)
+
+    return adapter

pypsrp/negotiate.py

@@ -36,7 +36,7 @@ class HTTPNegotiateAuth(AuthBase):
         password: typing.Optional[str] = None,
         auth_provider: str = "negotiate",
         send_cbt: bool = True,
-        service: str = "WSMAN",
+        service: str = "host",
         delegate: bool = False,
         hostname_override: typing.Optional[str] = None,
         wrap_required: bool = False,
@@ -60,7 +60,7 @@ class HTTPNegotiateAuth(AuthBase):
         :param send_cbt: Try to bind the channel token (HTTPS only) to the auth
             process, default is True
         :param service: The service part of the SPN to authenticate with,
-            defaults to HTTP
+            defaults to host
         :param delegate: Whether to get an auth token that allows the token to
             be delegated to other servers, this is only used with Kerberos and
             defaults to False

pypsrp/powershell.py

@@ -1,6 +1,8 @@
 # Copyright: (c) 2018, Jordan Borean (@jborean93) <jborean93@gmail.com>
 # MIT License (see LICENSE or https://opensource.org/licenses/MIT)
 
+from __future__ import annotations
+
 import base64
 import logging
 import struct
@@ -11,9 +13,11 @@ import uuid
 import warnings
 import xml.etree.ElementTree as ET
 
+import requests.exceptions
 from cryptography.hazmat.primitives.asymmetric import padding, rsa
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
+from pypsrp import _pool_manager
 from pypsrp._utils import version_equal_or_newer
 from pypsrp.complex_objects import (
     ApartmentState,
@@ -104,6 +108,7 @@ class RunspacePool(object):
         session_key_timeout_ms: int = 60000,
         *,
         no_profile: bool = False,
+        idle_timeout: int | None = None,
     ) -> None:
         """
         Represents a Runspace pool on a remote host. This pool can contain
@@ -132,6 +137,9 @@ class RunspacePool(object):
             key transfer from the server
         :param no_profile: If True, the user profile will not be loaded and
             will use the machine defaults.
+        :param idle_timeout: The idle timeout, in seconds, for the runspace. The
+            runspace will be closed if no operations are performed in the time
+            specified. If None the server default will be used.
         """
         log.info("Initialising RunspacePool object for configuration %s" % configuration_name)
         # The below are defined in some way at
@@ -147,6 +155,7 @@ class RunspacePool(object):
             input_streams="stdin pr",
             output_streams="stdout",
             no_profile=no_profile,
+            idle_time_out=idle_timeout,
         )
         self.ci_table: typing.Dict = {}
         self.pipelines: typing.Dict[str, "PowerShell"] = {}
@@ -659,6 +668,47 @@ class RunspacePool(object):
         """
         return self._serializer.serialize(obj, metadata=metadata)
 
+    def is_alive(
+        self,
+        *,
+        timeout: int | None = None,
+    ) -> bool:
+        """Checks whether the RunspacePool is still alive.
+
+        :param timeout: Override the connection timeout defaults for the
+            request. The WSMan operation timeout will be set to this value
+            and the HTTP connect/read timeouts will be this value + 2 seconds.
+        :return: A bool True if the pool is still alive, False otherwise.
+        """
+        is_closed = False
+
+        try:
+            # We don't want to try and open a new connection if the socket
+            # was closed, we treat it as the pool is not alive.
+            with _pool_manager.DisableNewConnectionsContext():
+                state = self.shell._get_shell(timeout=timeout)
+
+            is_closed = state.get("State", "") == "Disconnected"
+        except WSManFaultError as exc:
+            if exc.code in [
+                0x80338029,  # ERROR_WSMAN_OPERATION_TIMEDOUT
+                0x8033805B,  # ERROR_WSMAN_UNEXPECTED_SELECTORS
+            ]:
+                is_closed = True
+            else:
+                raise
+
+        except (_pool_manager.NewConnectionDisabled, TimeoutError, requests.exceptions.ConnectionError):
+            # If a timeout or connection error occurs, treat the pool as closed.
+            is_closed = True
+
+        if is_closed:
+            # Ensures that close() doesn't try and close an already closed pool
+            self.state = RunspacePoolState.CLOSED
+            return False
+
+        return True
+
     def _receive(
         self,
         id: typing.Optional[str] = None,

pypsrp/shell.py

@@ -1,6 +1,8 @@
 # Copyright: (c) 2018, Jordan Borean (@jborean93) <jborean93@gmail.com>
 # MIT License (see LICENSE or https://opensource.org/licenses/MIT)
 
+from __future__ import annotations
+
 import base64
 import logging
 import types
@@ -315,6 +317,42 @@ class WinRS(object):
         ET.SubElement(signal, "{%s}Code" % rsp).text = code
         return self.wsman.signal(self.resource_uri, signal, selector_set=self._selector_set)
 
+    def _get_shell(
+        self,
+        *,
+        timeout: int | None = None,
+    ) -> dict[str, str | None]:
+        """Queries the WSMan host for the current shell state.
+
+        Returns the WSMan GetResponse fields as a dictionary.
+
+        The timeout parameter can be used to override the default WSMan timeout
+        for this operation. If set the HTTP connect and read timeout will be
+        set to this value + 2 seconds to ensure the request does not block for
+        a longer time.
+
+        :param timeout: Sets the WSMan operational timeout to this value in
+            seconds. Overrides the default timeout set on the WSMan instance.
+        :return: A dictionary containing the shell state fields from the raw XML
+        response.
+        """
+        rsp = NAMESPACES["rsp"]
+
+        resp = self.wsman.get(
+            "http://schemas.microsoft.com/wbem/wsman/1/windows/shell",
+            selector_set=self._selector_set,
+            timeout=timeout,
+        )
+
+        state = {}
+        shell_items = resp.find("rsp:Shell", namespaces=NAMESPACES)
+        if shell_items is not None:
+            for child in shell_items:
+                field_name = child.tag.replace(f"{{{rsp}}}", "")
+                state[field_name] = child.text
+
+        return state
+
     def _parse_shell_create(
         self,
         response: ET.Element,

pypsrp/wsman.py

@@ -1,7 +1,7 @@
 # Copyright: (c) 2018, Jordan Borean (@jborean93) <jborean93@gmail.com>
 # MIT License (see LICENSE or https://opensource.org/licenses/MIT)
 
-from __future__ import division
+from __future__ import annotations
 
 import ipaddress
 import logging
@@ -15,6 +15,7 @@ import xml.etree.ElementTree as ET
 import requests
 from requests.packages.urllib3.util.retry import Retry
 
+from pypsrp import _pool_manager
 from pypsrp._utils import get_hostname, to_string, to_unicode
 from pypsrp.encryption import WinRMEncryption
 from pypsrp.exceptions import (
@@ -40,7 +41,7 @@ log = logging.getLogger(__name__)
 SUPPORTED_AUTHS = ["basic", "certificate", "credssp", "kerberos", "negotiate", "ntlm"]
 
 AUTH_KWARGS: typing.Dict[str, typing.List[str]] = {
-    "certificate": ["certificate_key_pem", "certificate_pem"],
+    "certificate": ["certificate_key_pem", "certificate_pem", "certificate_key_password"],
     "credssp": ["credssp_auth_mechanism", "credssp_disable_tlsv1_2", "credssp_minimum_version"],
     "negotiate": ["negotiate_delegate", "negotiate_hostname_override", "negotiate_send_cbt", "negotiate_service"],
 }
@@ -471,6 +472,14 @@ class WSMan(object):
         s = NAMESPACES["s"]
         envelope = ET.Element("{%s}Envelope" % s)
 
+        http_timeout = None
+        if timeout:
+            # Failsafe if the operation timeout exceeds we don't want the
+            # request to hang indefinitely. The HTTP timeout needs to be more
+            # than the WSMan timeout as we cannot guarantee a WSMan response
+            # until the operation timeout is hit.
+            http_timeout = timeout + 2
+
         message_id, header = self._create_header(action, resource_uri, option_set, selector_set, timeout)
         message_id = f"uuid:{message_id}"
         envelope.append(header)
@@ -482,7 +491,11 @@ class WSMan(object):
         xml = ET.tostring(envelope, encoding="utf-8", method="xml")
 
         try:
-            response = self.transport.send(xml, retries_on_read_timeout=retries_on_read_timeout)
+            response = self.transport.send(
+                xml,
+                retries_on_read_timeout=retries_on_read_timeout,
+                timeout=http_timeout,
+            )
         except WinRMTransportError as err:
             try:
                 # try and parse the XML and get the WSManFault
@@ -787,6 +800,7 @@ class _TransportHTTP(object):
 
         self.certificate_key_pem: typing.Optional[str] = None
         self.certificate_pem: typing.Optional[str] = None
+        self.certificate_key_password: typing.Optional[str] = None
         for kwarg_list in AUTH_KWARGS.values():
             for kwarg in kwarg_list:
                 setattr(self, kwarg, kwargs.get(kwarg, None))
@@ -810,6 +824,7 @@ class _TransportHTTP(object):
         message: bytes,
         *,
         retries_on_read_timeout: int = 0,
+        timeout: int | float | tuple[int | float, int | float] | None = None,
     ) -> bytes:
         hostname = get_hostname(self.endpoint)
         if self.session is None:
@@ -824,9 +839,9 @@ class _TransportHTTP(object):
                 prep_request = self.session.prepare_request(request)
 
                 try:
-                    self._send_request(prep_request)
+                    self._send_request(prep_request, timeout=timeout)
                 except (requests.ReadTimeout, requests.ConnectTimeout) as e:
-                    log.exception("%s during initial authentication request - attempt %d", (type(e).__name__, attempt))
+                    log.exception("%s during initial authentication request - attempt %d", type(e).__name__, attempt)
                     if attempt == retries_on_read_timeout:
                         raise
 
@@ -867,9 +882,12 @@ class _TransportHTTP(object):
             request = requests.Request("POST", self.endpoint, data=payload, headers=headers)
             prep_request = self.session.prepare_request(request)
             try:
-                return self._send_request(prep_request)
+                return self._send_request(
+                    prep_request,
+                    timeout=timeout,
+                )
             except (requests.ReadTimeout, requests.ConnectTimeout) as e:
-                log.exception("%s during WSMan request - attempt %d", (type(e).__name__, attempt))
+                log.exception("%s during WSMan request - attempt %d", type(e).__name__, attempt)
                 if attempt == retries_on_read_timeout:
                     raise
 
@@ -880,8 +898,15 @@ class _TransportHTTP(object):
                 attempt += 1
                 time.sleep(self.reconnection_backoff * (2**attempt - 1))
 
-    def _send_request(self, request: requests.PreparedRequest) -> bytes:
-        response = self.session.send(request, timeout=(self.connection_timeout, self.read_timeout))  # type: ignore[union-attr] # This should not happen
+    def _send_request(
+        self,
+        request: requests.PreparedRequest,
+        timeout: int | float | tuple[int | float, int | float] | None = None,
+    ) -> bytes:
+        response = self.session.send(  # type: ignore[union-attr] # This should not happen
+            request,
+            timeout=timeout if timeout is not None else (self.connection_timeout, self.read_timeout),
+        )
 
         content_type = response.headers.get("content-type", "")
         if content_type.startswith("multipart/encrypted;") or content_type.startswith("multipart/x-multi-encrypted;"):
@@ -959,8 +984,14 @@ class _TransportHTTP(object):
             del retry_kwargs["status"]
             retries = Retry(**retry_kwargs)
 
-        session.mount("http://", requests.adapters.HTTPAdapter(max_retries=retries))
-        session.mount("https://", requests.adapters.HTTPAdapter(max_retries=retries))
+        session.mount(
+            "http://",
+            _pool_manager.create_request_adapter(max_retries=retries),
+        )
+        session.mount(
+            "https://",
+            _pool_manager.create_request_adapter(max_retries=retries, key_password=self.certificate_key_password),
+        )
 
         # set cert validation config
         session.verify = self.cert_validation
@@ -998,6 +1029,9 @@ class _TransportHTTP(object):
         if self.ssl is False:
             raise ValueError("For certificate auth, SSL must be used")
 
+        # requests does not expose the password through the cert tuple. If set
+        # it'll be passed to urllib3 through the custom adapter created for the
+        # session.
         session.cert = (self.certificate_pem, self.certificate_key_pem)
         session.headers["Authorization"] = "http://schemas.dmtf.org/wbem/wsman/1/wsman/secprofile/https/mutual"
 

{pypsrp-0.9.0rc1.dist-info → pypsrp-0.9.0rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pypsrp
-Version: 0.9.0rc1
+Version: 0.9.0rc2
 Summary: PowerShell Remoting Protocol and WinRM for Python
 Author-email: Jordan Borean <jborean93@gmail.com>
 License-Expression: MIT
@@ -18,7 +18,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: cryptography>=3.1
 Requires-Dist: pyspnego<1.0.0,>=0.7.0
-Requires-Dist: requests>=2.9.1
+Requires-Dist: requests>=2.27.0
 Provides-Extra: credssp
 Requires-Dist: requests-credssp>=2.0.0; extra == "credssp"
 Provides-Extra: kerberos
@@ -224,7 +224,8 @@ These are the options that can be used to setup `WSMan`;
 * `data_locale`: The `wsmv:DataLocale` value to set on each WSMan request. This specifies the format in which numerical data is presented in the response text, default is the value of `locale`
 * `reconnection_retries`: Number of retries on a connection problem, default is `0`
 * `reconnection_backoff`: Number of seconds to backoff in between reconnection attempts (first sleeps X, then sleeps 2*X, 4*X, 8*X, ...), default is `2.0`
-* `certificate_key_pem`: The path to the certificate key used in `certificate` authentication
+* `certificate_key_pem`: The path to the certificate key used in `certificate` authentication. The key can be in either a `PKCS#1` or `PKCS#8` format
+* `certificate_key_password`: The password for `certificate_key_pem` if it is encrypted
 * `certificate_pem`: The path to the certificate used in `certificate` authentication
 * `credssp_auth_mechanism`: The sub-auth mechanism used in CredSSP, default is `auto`, choices are `auto`, `ntlm`, or `kerberos`
 * `credssp_disable_tlsv1_2`: Whether to used CredSSP auth over the insecure TLSv1.0, default is `False`

{pypsrp-0.9.0rc1.dist-info → pypsrp-0.9.0rc2.dist-info}/RECORD

@@ -1,4 +1,5 @@
 pypsrp/__init__.py,sha256=9aButMuBNIEph0LfHsodNd8uSS7XKZj9IPHJsPd8hjY,958
+pypsrp/_pool_manager.py,sha256=yKUzdzkE4hEo2mqIC-kKPL9uWpIk9mixCQ1AP5rgn_o,4222
 pypsrp/_utils.py,sha256=y9sezXQPl0vXnr1beyZ8xcn-Q_zdg_7cg_lPk9QmjVQ,3540
 pypsrp/client.py,sha256=iJERbEXwwE4wswpqAL0Dcz8ACqLvXUT9Dcl_Qqa3830,13892
 pypsrp/complex_objects.py,sha256=PkJ1F1Byw5RI_Vpfy1KfcMFJdOfioOSsIbi50S_fZ3g,62666
@@ -6,17 +7,17 @@ pypsrp/encryption.py,sha256=4wOontH-eF8BBxTyHTX9iAF0uIM6m-z3N9poatspP1A,3870
 pypsrp/exceptions.py,sha256=SWdQ9UAPOIwr7B4twGqDoXLCEtXFK_5u3PYIVLvEIbA,4034
 pypsrp/host.py,sha256=XQgnkwzsZprNmAr1TPJpEBmj3xRSWFrYQFSIi870aNE,45065
 pypsrp/messages.py,sha256=mCuie1ywOQmkXltISPvrb6ITS6GUA0oSld2jfNLdBf4,34169
-pypsrp/negotiate.py,sha256=YU6mGuDd8krOapiRcY6qID8v-GMDbIypCYUc7eTkzmw,10952
-pypsrp/powershell.py,sha256=SjmD8LfbMtyoGfXpQmSV10ded4FO14J8vM5Qn37R0dI,67359
+pypsrp/negotiate.py,sha256=iVX2J-slABamesMajjTfmfqucDS4n5vVTTde_IDqpuE,10951
+pypsrp/powershell.py,sha256=NRDoFtOsYjIrJHMQ6r6vRq6uL2EtPhhem5ZA_bWoDnw,69296
 pypsrp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pypsrp/serializer.py,sha256=49rBxfo3fh511mU3uzoK9vFvIOrBtyX3J-CUugbezc0,33767
-pypsrp/shell.py,sha256=EoGtElIYvNFPjBvAIMuLaHPyX0NgGtBN265Vt6YqkhQ,16748
-pypsrp/wsman.py,sha256=V90G93O2bi9fPP09fPhyheYN4_G_8zo_Sxo1ea2pqPI,47500
+pypsrp/shell.py,sha256=80mh35c-UXjsy_MFG9J3GGYLYqvk2FppUgrFHkWh6xg,18065
+pypsrp/wsman.py,sha256=KkRg9UDDq1yxY4TKK9mGI-vS6px6ftQDPhmeojjItKg,48728
 pypsrp/pwsh_scripts/__init__.py,sha256=VckRBWWDUdh_LTwPspE2FXSo8TaT50QnZ1aZ-yvxJmU,139
 pypsrp/pwsh_scripts/copy.ps1,sha256=sTy4jcnh0cV5Y0l2-dd1znc-eM_kgE-okk51xZzdIQk,5374
 pypsrp/pwsh_scripts/fetch.ps1,sha256=7gvPxURnksZetYgjWRjbIDIWSJYv2xGXZjRymvNUj6w,1989
-pypsrp-0.9.0rc1.dist-info/licenses/LICENSE,sha256=8uJLgs8Wb_Us_8XaTvOOn7Yf0dwzFi9pcFHVxfgJKoE,1079
-pypsrp-0.9.0rc1.dist-info/METADATA,sha256=RsbvzIn44QmLDG7Pi5QCULyxtJ8Z2kZByhI7dkU6Fmo,23485
-pypsrp-0.9.0rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pypsrp-0.9.0rc1.dist-info/top_level.txt,sha256=9sFQD0PMIG07xbaR4j-1Fq7NaD9kwInn_HTc3Gqznq0,7
-pypsrp-0.9.0rc1.dist-info/RECORD,,
+pypsrp-0.9.0rc2.dist-info/licenses/LICENSE,sha256=8uJLgs8Wb_Us_8XaTvOOn7Yf0dwzFi9pcFHVxfgJKoE,1079
+pypsrp-0.9.0rc2.dist-info/METADATA,sha256=_zCveXjKqq1_OiCmjLtge2Zr3v1G-S2EFrF33PTd6QA,23630
+pypsrp-0.9.0rc2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+pypsrp-0.9.0rc2.dist-info/top_level.txt,sha256=9sFQD0PMIG07xbaR4j-1Fq7NaD9kwInn_HTc3Gqznq0,7
+pypsrp-0.9.0rc2.dist-info/RECORD,,

{pypsrp-0.9.0rc1.dist-info → pypsrp-0.9.0rc2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any