diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py index 0b9589d9dd02f..98f9190bc7c96 100644 --- a/airflow/providers/microsoft/psrp/hooks/psrp.py +++ b/airflow/providers/microsoft/psrp/hooks/psrp.py @@ -16,103 +16,245 @@ # specific language governing permissions and limitations # under the License. -from time import sleep +from contextlib import contextmanager +from copy import copy +from logging import DEBUG, ERROR, INFO, WARNING +from typing import Any, Callable, Dict, Iterator, Optional +from weakref import WeakKeyDictionary -from pypsrp.messages import ErrorRecord, InformationRecord, ProgressRecord +from pypsrp.messages import MessageType from pypsrp.powershell import PowerShell, PSInvocationState, RunspacePool from pypsrp.wsman import WSMan from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook +INFORMATIONAL_RECORD_LEVEL_MAP = { + MessageType.DEBUG_RECORD: DEBUG, + MessageType.ERROR_RECORD: ERROR, + MessageType.VERBOSE_RECORD: INFO, + MessageType.WARNING_RECORD: WARNING, +} -class PSRPHook(BaseHook): +OutputCallback = Callable[[str], None] + + +class PsrpHook(BaseHook): """ Hook for PowerShell Remoting Protocol execution. - The hook must be used as a context manager. + When used as a context manager, the runspace pool is reused between shell + sessions. + + :param psrp_conn_id: Required. The name of the PSRP connection. + :type psrp_conn_id: str + :param logging_level: + Logging level for message streams which are received during remote execution. + The default is to include all messages in the task log. + :type logging_level: int + :param operation_timeout: Override the default WSMan timeout when polling the pipeline. + :type operation_timeout: float + :param runspace_options: + Optional dictionary which is passed when creating the runspace pool. See + :py:class:`~pypsrp.powershell.RunspacePool` for a description of the + available options. + :type runspace_options: dict + :param wsman_options: + Optional dictionary which is passed when creating the `WSMan` client. See + :py:class:`~pypsrp.wsman.WSMan` for a description of the available options. + :type wsman_options: dict + :param on_output_callback: + Optional callback function to be called whenever an output response item is + received during job status polling. + :type on_output_callback: OutputCallback + :param exchange_keys: + If true (default), automatically initiate a session key exchange when the + hook is used as a context manager. + :type exchange_keys: bool + + You can provide an alternative `configuration_name` using either `runspace_options` + or by setting this key as the extra fields of your connection. """ - _client = None - _poll_interval = 1 + _conn = None + _configuration_name = None + _wsman_ref: "WeakKeyDictionary[RunspacePool, WSMan]" = WeakKeyDictionary() - def __init__(self, psrp_conn_id: str): + def __init__( + self, + psrp_conn_id: str, + logging_level: int = DEBUG, + operation_timeout: Optional[int] = None, + runspace_options: Optional[Dict[str, Any]] = None, + wsman_options: Optional[Dict[str, Any]] = None, + on_output_callback: Optional[OutputCallback] = None, + exchange_keys: bool = True, + ): self.conn_id = psrp_conn_id + self._logging_level = logging_level + self._operation_timeout = operation_timeout + self._runspace_options = runspace_options or {} + self._wsman_options = wsman_options or {} + self._on_output_callback = on_output_callback + self._exchange_keys = exchange_keys def __enter__(self): - conn = self.get_connection(self.conn_id) - - self.log.info("Establishing WinRM connection %s to host: %s", self.conn_id, conn.host) - self._client = WSMan( - conn.host, - ssl=True, - auth="ntlm", - encryption="never", - username=conn.login, - password=conn.password, - cert_validation=False, - ) - self._client.__enter__() + conn = self.get_conn() + self._wsman_ref[conn].__enter__() + conn.__enter__() + if self._exchange_keys: + conn.exchange_keys() + self._conn = conn return self def __exit__(self, exc_type, exc_value, traceback): try: - self._client.__exit__(exc_type, exc_value, traceback) + self._conn.__exit__(exc_type, exc_value, traceback) + self._wsman_ref[self._conn].__exit__(exc_type, exc_value, traceback) finally: - self._client = None + del self._conn - def invoke_powershell(self, script: str) -> PowerShell: - with RunspacePool(self._client) as pool: - ps = PowerShell(pool) - ps.add_script(script) + def get_conn(self) -> RunspacePool: + """ + Returns a runspace pool. + + The returned object must be used as a context manager. + """ + conn = self.get_connection(self.conn_id) + self.log.info("Establishing WinRM connection %s to host: %s", self.conn_id, conn.host) + + extra = conn.extra_dejson.copy() + + def apply_extra(d, keys): + d = d.copy() + for key in keys: + value = extra.pop(key, None) + if value is not None: + d[key] = value + return d + + wsman_options = apply_extra( + self._wsman_options, + ( + "auth", + "cert_validation", + "connection_timeout", + "locale", + "read_timeout", + "reconnection_retries", + "reconnection_backoff", + "ssl", + ), + ) + wsman = WSMan(conn.host, username=conn.login, password=conn.password, **wsman_options) + runspace_options = apply_extra(self._runspace_options, ("configuration_name",)) + + if extra: + raise AirflowException(f"Unexpected extra configuration keys: {', '.join(sorted(extra))}") + pool = RunspacePool(wsman, **runspace_options) + self._wsman_ref[pool] = wsman + return pool + + @contextmanager + def invoke(self) -> Iterator[PowerShell]: + """ + Context manager that yields a PowerShell object to which commands can be + added. Upon exit, the commands will be invoked. + """ + logger = copy(self.log) + logger.setLevel(self._logging_level) + local_context = self._conn is None + if local_context: + self.__enter__() + try: + assert self._conn is not None + ps = PowerShell(self._conn) + yield ps ps.begin_invoke() + streams = [ - (ps.output, self._log_output), - (ps.streams.debug, self._log_record), - (ps.streams.information, self._log_record), - (ps.streams.error, self._log_record), + ps.output, + ps.streams.debug, + ps.streams.error, + ps.streams.information, + ps.streams.progress, + ps.streams.verbose, + ps.streams.warning, ] offsets = [0 for _ in streams] # We're using polling to make sure output and streams are # handled while the process is running. while ps.state == PSInvocationState.RUNNING: - sleep(self._poll_interval) - ps.poll_invoke() + ps.poll_invoke(timeout=self._operation_timeout) - for (i, (stream, handler)) in enumerate(streams): + for i, stream in enumerate(streams): offset = offsets[i] while len(stream) > offset: - handler(stream[offset]) + record = stream[offset] + + # Records received on the output stream during job + # status polling are handled via an optional callback, + # while the other streams are simply logged. + if stream is ps.output: + if self._on_output_callback is not None: + self._on_output_callback(record) + else: + self._log_record(logger.log, record) offset += 1 offsets[i] = offset # For good measure, we'll make sure the process has - # stopped running. + # stopped running in any case. ps.end_invoke() + self.log.info("Invocation state: %s", str(PSInvocationState(ps.state))) if ps.streams.error: raise AirflowException("Process had one or more errors") + finally: + if local_context: + self.__exit__(None, None, None) - self.log.info("Invocation state: %s", str(PSInvocationState(ps.state))) - return ps + def invoke_cmdlet(self, name: str, use_local_scope=None, **parameters: Dict[str, str]) -> PowerShell: + """Invoke a PowerShell cmdlet and return session.""" + with self.invoke() as ps: + ps.add_cmdlet(name, use_local_scope=use_local_scope) + ps.add_parameters(parameters) + return ps - def _log_output(self, message: str): - self.log.info("%s", message) + def invoke_powershell(self, script: str) -> PowerShell: + """Invoke a PowerShell script and return session.""" + with self.invoke() as ps: + ps.add_script(script) + return ps - def _log_record(self, record): - # TODO: Consider translating some or all of these records into - # normal logging levels, using `log(level, msg, *args)`. - if isinstance(record, ErrorRecord): - self.log.info("Error: %s", record) - return + def _log_record(self, log, record): + message_type = record.MESSAGE_TYPE + if message_type == MessageType.ERROR_RECORD: + log(INFO, "%s: %s", record.reason, record) + if record.script_stacktrace: + for trace in record.script_stacktrace.split('\r\n'): + log(INFO, trace) - if isinstance(record, InformationRecord): - self.log.info("Information: %s", record.message_data) - return + level = INFORMATIONAL_RECORD_LEVEL_MAP.get(message_type) + if level is not None: + try: + message = str(record.message) + except BaseException as exc: + # See https://github.com/jborean93/pypsrp/pull/130 + message = str(exc) - if isinstance(record, ProgressRecord): - self.log.info("Progress: %s (%s)", record.activity, record.description) - return + # Sometimes a message will have a trailing \r\n sequence such as + # the tracing output of the Set-PSDebug cmdlet. + message = message.rstrip() - self.log.info("Unsupported record type: %s", type(record).__name__) + if record.command_name is None: + log(level, "%s", message) + else: + log(level, "%s: %s", record.command_name, message) + elif message_type == MessageType.INFORMATION_RECORD: + log(INFO, "%s (%s): %s", record.computer, record.user, record.message_data) + elif message_type == MessageType.PROGRESS_RECORD: + log(INFO, "Progress: %s (%s)", record.activity, record.description) + else: + log(WARNING, "Unsupported message type: %s", message_type) diff --git a/airflow/providers/microsoft/psrp/operators/psrp.py b/airflow/providers/microsoft/psrp/operators/psrp.py index bedf4b2423cda..dc63881109e37 100644 --- a/airflow/providers/microsoft/psrp/operators/psrp.py +++ b/airflow/providers/microsoft/psrp/operators/psrp.py @@ -16,30 +16,78 @@ # specific language governing permissions and limitations # under the License. -from typing import TYPE_CHECKING, List, Optional, Sequence +from logging import DEBUG +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence + +from jinja2.nativetypes import NativeEnvironment +from pypsrp.powershell import Command +from pypsrp.serializer import TaggedValue from airflow.exceptions import AirflowException from airflow.models import BaseOperator -from airflow.providers.microsoft.psrp.hooks.psrp import PSRPHook +from airflow.providers.microsoft.psrp.hooks.psrp import PsrpHook +from airflow.settings import json + + +# TODO: Replace with airflow.utils.helpers.exactly_one in Airflow 2.3. +def exactly_one(*args): + return len(set(filter(None, args))) == 1 + if TYPE_CHECKING: from airflow.utils.context import Context -class PSRPOperator(BaseOperator): +class PsrpOperator(BaseOperator): """PowerShell Remoting Protocol operator. + Use one of the 'command', 'cmdlet', or 'powershell' arguments. + + The 'securestring' template filter can be used to tag a value for + serialization into a `System.Security.SecureString` (applicable only + for DAGs which have `render_template_as_native_obj=True`). + + When using the `cmdlet` or `powershell` arguments and when `do_xcom_push` + is enabled, the command output is converted to JSON by PowerShell using + the `ConvertTo-Json + `__ cmdlet such + that the operator return value is serializable to an XCom value. + :param psrp_conn_id: connection id :param command: command to execute on remote host. (templated) :param powershell: powershell to execute on remote host. (templated) + :param cmdlet: + cmdlet to execute on remote host (templated). Also used as the default + value for `task_id`. + :param parameters: + When using the `cmdlet` or `powershell` arguments, use this parameter to + provide parameters (templated). Note that a parameter with a value of `None` + becomes an *argument* (i.e., switch). + :param logging_level: + Logging level for message streams which are received during remote execution. + The default is to include all messages in the task log. + :param runspace_options: + optional dictionary which is passed when creating the runspace pool. See + :py:class:`~pypsrp.powershell.RunspacePool` for a description of the + available options. + :param wsman_options: + optional dictionary which is passed when creating the `WSMan` client. See + :py:class:`~pypsrp.wsman.WSMan` for a description of the available options. + :param psrp_session_init: + Optional command which will be added to the pipeline when a new PowerShell + session has been established, prior to invoking the action specified using + the `cmdlet`, `command`, or `powershell` parameters. """ template_fields: Sequence[str] = ( + "cmdlet", "command", + "parameters", "powershell", ) template_fields_renderers = {"command": "powershell", "powershell": "powershell"} - ui_color = "#901dd2" + ui_color = "#c2e2ff" def __init__( self, @@ -47,20 +95,75 @@ def __init__( psrp_conn_id: str, command: Optional[str] = None, powershell: Optional[str] = None, + cmdlet: Optional[str] = None, + parameters: Optional[Dict[str, str]] = None, + logging_level: int = DEBUG, + runspace_options: Optional[Dict[str, Any]] = None, + wsman_options: Optional[Dict[str, Any]] = None, + psrp_session_init: Optional[Command] = None, **kwargs, ) -> None: + args = {command, powershell, cmdlet} + if not exactly_one(*args): + raise ValueError("Must provide exactly one of 'command', 'powershell', or 'cmdlet'") + if parameters and not cmdlet: + raise ValueError("Parameters only allowed with 'cmdlet'") + if cmdlet: + kwargs.setdefault('task_id', cmdlet) super().__init__(**kwargs) - if not (command or powershell): - raise ValueError("Must provide either 'command' or 'powershell'") self.conn_id = psrp_conn_id self.command = command self.powershell = powershell + self.cmdlet = cmdlet + self.parameters = parameters + self.logging_level = logging_level + self.runspace_options = runspace_options + self.wsman_options = wsman_options + self.psrp_session_init = psrp_session_init + + def execute(self, context: "Context") -> Optional[List[Any]]: + with PsrpHook( + self.conn_id, + logging_level=self.logging_level, + runspace_options=self.runspace_options, + wsman_options=self.wsman_options, + on_output_callback=self.log.info if not self.do_xcom_push else None, + ) as hook, hook.invoke() as ps: + if self.psrp_session_init is not None: + ps.add_command(self.psrp_session_init) + if self.command: + ps.add_script(f"cmd.exe /c @'\n{self.command}\n'@") + else: + if self.cmdlet: + ps.add_cmdlet(self.cmdlet) + else: + ps.add_script(self.powershell) + if self.parameters: + ps.add_parameters(self.parameters) + if self.do_xcom_push: + ps.add_cmdlet("ConvertTo-Json") - def execute(self, context: "Context") -> List[str]: - with PSRPHook(self.conn_id) as hook: - ps = hook.invoke_powershell( - f"cmd.exe /c @'\n{self.command}\n'@" if self.command else self.powershell - ) if ps.had_errors: raise AirflowException("Process failed") - return ps.output + + if not self.do_xcom_push: + return None + + return [json.loads(output) for output in ps.output] + + def get_template_env(self): + # Create a template environment overlay in order to leave the underlying + # environment unchanged. + env = super().get_template_env().overlay() + native = isinstance(env, NativeEnvironment) + + def securestring(value: str): + if not native: + raise AirflowException( + "Filter 'securestring' not applicable to non-native " "templating environment" + ) + return TaggedValue("SS", value) + + env.filters["securestring"] = securestring + + return env diff --git a/airflow/providers/microsoft/psrp/provider.yaml b/airflow/providers/microsoft/psrp/provider.yaml index d60b3cd44cf54..e71ff17c553ec 100644 --- a/airflow/providers/microsoft/psrp/provider.yaml +++ b/airflow/providers/microsoft/psrp/provider.yaml @@ -19,15 +19,16 @@ package-name: apache-airflow-providers-microsoft-psrp name: PowerShell Remoting Protocol (PSRP) description: | + This package provides remote execution capabilities via the `PowerShell Remoting Protocol (PSRP) - `__ + `__. versions: - 1.0.1 - 1.0.0 additional-dependencies: - - pypsrp>=0.5.0 + - pypsrp>=0.8.0 integrations: - integration-name: Windows PowerShell Remoting Protocol diff --git a/docs/apache-airflow-providers-microsoft-psrp/index.rst b/docs/apache-airflow-providers-microsoft-psrp/index.rst index be35ad3446bc0..1872be0a23e2b 100644 --- a/docs/apache-airflow-providers-microsoft-psrp/index.rst +++ b/docs/apache-airflow-providers-microsoft-psrp/index.rst @@ -22,6 +22,12 @@ Content ------- +.. toctree:: + :maxdepth: 1 + :caption: Guides + + Operators + .. toctree:: :maxdepth: 1 :caption: References @@ -72,8 +78,8 @@ PIP requirements ============= ================== PIP package Version required ============= ================== -``pypsrp`` ``>=0.5.0`` -``pypsrp`` ``~=0.5`` +``pypsrp`` ``>=0.8.0`` +``pypsrp`` ``~=0.8`` ============= ================== .. include:: ../../airflow/providers/microsoft/psrp/CHANGELOG.rst diff --git a/docs/apache-airflow-providers-microsoft-psrp/operators/index.rst b/docs/apache-airflow-providers-microsoft-psrp/operators/index.rst new file mode 100644 index 0000000000000..ffa739dbc0295 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-psrp/operators/index.rst @@ -0,0 +1,111 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + +Microsoft PSRP Operators +======================================= + +The +`PowerShell Remoting Protocol (PSRP) +`__ +protocol is required whenever a user wants to execute commands on a Windows +server from a client using a native PowerShell runspace. + +The :class:`~airflow.providers.microsoft.psrp.operators.psrp.PsrpOperator` +operator implements such client capabilities, enabling the +scheduling of Windows jobs from Airflow. Internally, it makes use of +the `pypsrp `__ client library. + +Compared to +:class:`~airflow.providers.microsoft.winrm.operators.winrm.WinRMOperator`, +using PSRP extends the remoting capabilities in Windows, providing +better session control and close integration with the PowerShell +ecosystem (i.e., .NET Runspace interface): + +* Run multiple commands in a single session +* Reuse the runspace to create multiple sessions +* Work with PowerShell objects instead of just text +* Use constrained endpoints using JEA (Just-Enough-Administration) +* Ability to use the .NET Runspace interface + + +Using the Operator +^^^^^^^^^^^^^^^^^^ + +When instantiating the +:class:`~airflow.providers.microsoft.psrp.operators.psrp.PsrpOperator` +operator, you must provide a cmdlet, command or script using one of the +following named arguments: + +.. list-table:: Provide one of the following arguments to the operator + :widths: 10 15 20 + :header-rows: 1 + + * - Argument name + - Description + - Examples + * - cmdlet + - Invoke a PowerShell cmdlet. + - ``Copy-Item``, ``Restart-Computer`` + * - command + - Carries out the specified command using the + `cmd `__ command interpreter. + - ``robocopy C:\Logfiles\* C:\Drawings /S /E`` + * - powershell + - Run a PowerShell script. + - ``Copy-Item -Path "C:\Logfiles\*" -Destination "C:\Drawings" -Recurse`` + + +Output +###### + +PowerShell provides multiple output streams. + +In general, the operator logs a record using the built-in logging +mechanism for records that arrive on these streams using a job status +polling mechanism. The success stream (i.e., stdout or shell output) +is handled differently, as explained in the following: + +When :doc:`XComs ` are enabled and when +the operator is used with a native PowerShell cmdlet or script, the +shell output is converted to JSON using the ``ConvertTo-Json`` cmdlet +and then decoded on the client-side by the operator such that the +operator's return value is compatible with the serialization required +by XComs. + +When XComs are not enabled (that is, ``do_xcom_push`` is set to +false), the shell output is instead logged like the other output +streams and will appear in the task instance log. + + +Secure strings +############## + +The operator adds a template filter ``securestring`` which will encrypt +the value and make it available in the remote session as a +`SecureString +`__ +type. This ensures for example that the value is not accidentally +logged. + +Using the template filter requires the DAG to be configured to +:ref:`render fields as native objects +` (the default is to coerce all +values into strings which won't work here because we need a value +which has been tagged to be serialized as a secure string). Use +``render_template_as_native_obj=True`` to enable this. diff --git a/docs/apache-airflow/concepts/operators.rst b/docs/apache-airflow/concepts/operators.rst index 3d489c2f3bc5d..9752287daa1ce 100644 --- a/docs/apache-airflow/concepts/operators.rst +++ b/docs/apache-airflow/concepts/operators.rst @@ -148,6 +148,8 @@ You can pass custom options to the Jinja ``Environment`` when creating your DAG. See the `Jinja documentation `_ to find all available options. +.. _concepts:templating-native-objects: + Rendering Fields as Native Python Objects ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 95f8bfb0bec89..4ac6a960b1b67 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -575,6 +575,7 @@ cloudwatch cls cmake cmd +cmdlet cmdline cmds cname @@ -1195,6 +1196,7 @@ reidentify reinit reinitialization relativedelta +remoting renewer replicaSet repo @@ -1223,6 +1225,7 @@ rtype ru runAsUser runnable +runspace runtime sagemaker salesforce @@ -1243,6 +1246,7 @@ sdk secretRef secretRefs securable +securestring securityManager seealso seedlist diff --git a/setup.py b/setup.py index ae8b7e72b5ac2..9fe0fcff8c4bc 100644 --- a/setup.py +++ b/setup.py @@ -450,7 +450,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version pandas_requirement, ] psrp = [ - 'pypsrp~=0.5', + 'pypsrp~=0.8', ] qubole = [ 'qds-sdk>=1.10.4', diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py index 5a161adfaba3e..67bf63164dc1d 100644 --- a/tests/providers/microsoft/psrp/hooks/test_psrp.py +++ b/tests/providers/microsoft/psrp/hooks/test_psrp.py @@ -17,55 +17,184 @@ # under the License. # -import unittest -from unittest.mock import MagicMock, call, patch +from logging import DEBUG, ERROR, INFO, WARNING +from unittest import TestCase +from unittest.mock import MagicMock, Mock, call, patch -from pypsrp.messages import InformationRecord +from parameterized import parameterized +from pypsrp.messages import MessageType from pypsrp.powershell import PSInvocationState +from pytest import raises +from airflow.exceptions import AirflowException from airflow.models import Connection -from airflow.providers.microsoft.psrp.hooks.psrp import PSRPHook +from airflow.providers.microsoft.psrp.hooks.psrp import PsrpHook CONNECTION_ID = "conn_id" +DUMMY_STACKTRACE = [ + r"at Invoke-Foo, C:\module.psm1: line 113", + r"at Invoke-Bar, C:\module.psm1: line 125", +] -class TestPSRPHook(unittest.TestCase): - @patch( - f"{PSRPHook.__module__}.{PSRPHook.__name__}.get_connection", - return_value=Connection( - login='username', - password='password', - host='remote_host', - ), - ) - @patch(f"{PSRPHook.__module__}.WSMan") - @patch(f"{PSRPHook.__module__}.PowerShell") - @patch(f"{PSRPHook.__module__}.RunspacePool") - @patch("logging.Logger.info") - def test_invoke_powershell(self, log_info, runspace_pool, powershell, ws_man, get_connection): - with PSRPHook(CONNECTION_ID) as hook: - ps = powershell.return_value = MagicMock() - ps.state = PSInvocationState.RUNNING - ps.output = [] - ps.streams.debug = [] - ps.streams.information = [] - ps.streams.error = [] - - def poll_invoke(): - ps.output.append("") - ps.streams.debug.append(MagicMock(spec=InformationRecord, message_data="")) - ps.state = PSInvocationState.COMPLETED - - def end_invoke(): - ps.streams.error = [] - - ps.poll_invoke.side_effect = poll_invoke - ps.end_invoke.side_effect = end_invoke - - hook.invoke_powershell("foo") +class MockPowerShell(MagicMock): + had_errors = False + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.state = PSInvocationState.NOT_STARTED + + def poll_invoke(self, timeout=None): + self.state = PSInvocationState.COMPLETED + self.output.append("output") + + def informational(message_type, message, **kwargs): + kwargs.setdefault("command_name", "command") + return Mock(MESSAGE_TYPE=message_type, message=message, **kwargs) + + self.streams.debug.append(informational(MessageType.DEBUG_RECORD, "debug1")) + self.streams.debug.append(informational(MessageType.DEBUG_RECORD, "debug2\r\n", command_name=None)) + self.streams.verbose.append(informational(MessageType.VERBOSE_RECORD, "verbose")) + self.streams.warning.append(informational(MessageType.WARNING_RECORD, "warning")) + self.streams.information.append( + Mock( + MESSAGE_TYPE=MessageType.INFORMATION_RECORD, + computer="computer", + user="user", + message_data="information", + ) + ) + self.streams.progress.append( + Mock(MESSAGE_TYPE=MessageType.PROGRESS_RECORD, activity="activity", description="description") + ) + + if self.had_errors: + self.streams.error.append( + Mock( + MESSAGE_TYPE=MessageType.ERROR_RECORD, + command_name="command", + message="error", + reason="reason", + script_stacktrace="\r\n".join(DUMMY_STACKTRACE), + ) + ) + + def begin_invoke(self): + self.state = PSInvocationState.RUNNING + self.output = [] + self.streams.debug = [] + self.streams.error = [] + self.streams.information = [] + self.streams.progress = [] + self.streams.verbose = [] + self.streams.warning = [] + + def end_invoke(self): + while self.state == PSInvocationState.RUNNING: + self.poll_invoke() + + +def mock_powershell_factory(): + return MagicMock(return_value=MockPowerShell()) + + +@patch( + f"{PsrpHook.__module__}.{PsrpHook.__name__}.get_connection", + new=lambda _, conn_id: Connection( + conn_id=conn_id, + login='username', + password='password', + host='remote_host', + ), +) +@patch(f"{PsrpHook.__module__}.WSMan") +@patch(f"{PsrpHook.__module__}.PowerShell", new_callable=mock_powershell_factory) +@patch(f"{PsrpHook.__module__}.RunspacePool") +class TestPsrpHook(TestCase): + def test_get_conn(self, runspace_pool, powershell, ws_man): + hook = PsrpHook(CONNECTION_ID) + assert hook.get_conn() is runspace_pool.return_value + + def test_get_conn_unexpected_extra(self, runspace_pool, powershell, ws_man): + hook = PsrpHook(CONNECTION_ID) + conn = hook.get_connection(CONNECTION_ID) + + def get_connection(*args): + conn.extra = '{"foo": "bar"}' + return conn + + hook.get_connection = get_connection + with raises(AirflowException, match="Unexpected extra configuration keys: foo"): + hook.get_conn() + + @parameterized.expand([(None,), (ERROR,)]) + def test_invoke(self, runspace_pool, powershell, ws_man, logging_level): + runspace_options = {"connection_name": "foo"} + wsman_options = {"encryption": "auto"} + + options = {} + if logging_level is not None: + options["logging_level"] = logging_level + + on_output_callback = Mock() + + with PsrpHook( + CONNECTION_ID, + runspace_options=runspace_options, + wsman_options=wsman_options, + on_output_callback=on_output_callback, + **options, + ) as hook, patch.object(type(hook), "log") as logger: + try: + with hook.invoke() as ps: + assert ps.state == PSInvocationState.NOT_STARTED + + # We're simulating an error in order to test error + # handling as well as the logging of error exception + # details. + ps.had_errors = True + except AirflowException as exc: + assert str(exc) == 'Process had one or more errors' + else: + self.fail("Expected an error") + assert ps.state == PSInvocationState.COMPLETED + + assert on_output_callback.mock_calls == [call('output')] + assert runspace_pool.return_value.__exit__.mock_calls == [call(None, None, None)] assert ws_man().__exit__.mock_calls == [call(None, None, None)] + assert ws_man.call_args_list[0][1]["encryption"] == "auto" + assert logger.method_calls[0] == call.setLevel(logging_level or DEBUG) + + def assert_log(level, *args): + assert call.log(level, *args) in logger.method_calls + + assert_log(DEBUG, '%s: %s', 'command', 'debug1') + assert_log(DEBUG, '%s', 'debug2') + assert_log(ERROR, '%s: %s', 'command', 'error') + assert_log(INFO, '%s: %s', 'command', 'verbose') + assert_log(WARNING, '%s: %s', 'command', 'warning') + assert_log(INFO, 'Progress: %s (%s)', 'activity', 'description') + assert_log(INFO, '%s (%s): %s', 'computer', 'user', 'information') + assert_log(INFO, '%s: %s', 'reason', ps.streams.error[0]) + assert_log(INFO, DUMMY_STACKTRACE[0]) + assert_log(INFO, DUMMY_STACKTRACE[1]) + + assert call('Invocation state: %s', 'Completed') in logger.info.mock_calls + + assert runspace_pool.call_args == call(ws_man.return_value, connection_name='foo') + + def test_invoke_cmdlet(self, *mocks): + with PsrpHook(CONNECTION_ID) as hook: + ps = hook.invoke_cmdlet('foo', bar="1", baz="2") + assert [call('foo', use_local_scope=None)] == ps.add_cmdlet.mock_calls + assert [call({'bar': '1', 'baz': '2'})] == ps.add_parameters.mock_calls + + def test_invoke_powershell(self, *mocks): + with PsrpHook(CONNECTION_ID) as hook: + ps = hook.invoke_powershell('foo') + assert call('foo') in ps.add_script.mock_calls - assert call('%s', '') in log_info.mock_calls - assert call('Information: %s', '') in log_info.mock_calls - assert call('Invocation state: %s', 'Completed') in log_info.mock_calls + def test_invoke_local_context(self, *mocks): + hook = PsrpHook(CONNECTION_ID) + ps = hook.invoke_powershell('foo') + assert call('foo') in ps.add_script.mock_calls diff --git a/tests/providers/microsoft/psrp/operators/test_psrp.py b/tests/providers/microsoft/psrp/operators/test_psrp.py index df4b6fe54c3fb..288d630e8db11 100644 --- a/tests/providers/microsoft/psrp/operators/test_psrp.py +++ b/tests/providers/microsoft/psrp/operators/test_psrp.py @@ -16,40 +16,104 @@ # specific language governing permissions and limitations # under the License. -import unittest -from unittest.mock import patch +from itertools import product +from typing import Any, Dict, NamedTuple, Optional +from unittest import TestCase +from unittest.mock import Mock, call, patch import pytest +from jinja2.nativetypes import NativeEnvironment from parameterized import parameterized +from pypsrp.powershell import Command, PowerShell from airflow.exceptions import AirflowException -from airflow.providers.microsoft.psrp.operators.psrp import PSRPOperator +from airflow.models.baseoperator import BaseOperator +from airflow.providers.microsoft.psrp.operators.psrp import PsrpOperator +from airflow.settings import json CONNECTION_ID = "conn_id" -class TestPSRPOperator(unittest.TestCase): +class ExecuteParameter(NamedTuple): + name: str + expected_method: str + expected_parameters: Optional[Dict[str, Any]] + + +class TestPsrpOperator(TestCase): def test_no_command_or_powershell(self): - exception_msg = "Must provide either 'command' or 'powershell'" + exception_msg = "Must provide exactly one of 'command', 'powershell', or 'cmdlet'" with pytest.raises(ValueError, match=exception_msg): - PSRPOperator(task_id='test_task_id', psrp_conn_id=CONNECTION_ID) + PsrpOperator(task_id='test_task_id', psrp_conn_id=CONNECTION_ID) + + def test_cmdlet_task_id_default(self): + operator = PsrpOperator(cmdlet='Invoke-Foo', psrp_conn_id=CONNECTION_ID) + assert operator.task_id == 'Invoke-Foo' @parameterized.expand( - [ - (False,), - (True,), - ] + list( + product( + [ + # These tuples map the command parameter to an execution method + # and parameter set. + ExecuteParameter("command", call.add_script("cmd.exe /c @'\nfoo\n'@"), None), + ExecuteParameter("powershell", call.add_script("foo"), None), + ExecuteParameter("cmdlet", call.add_cmdlet("foo"), {"bar": "baz"}), + ], + [False, True], + [False, True], + ) + ) ) - @patch(f"{PSRPOperator.__module__}.PSRPHook") - def test_execute(self, had_errors, hook): - op = PSRPOperator(task_id='test_task_id', psrp_conn_id=CONNECTION_ID, command='dummy') - ps = hook.return_value.__enter__.return_value.invoke_powershell.return_value - ps.output = [""] - ps.had_errors = had_errors + @patch(f"{PsrpOperator.__module__}.PsrpHook") + def test_execute(self, parameter, had_errors, do_xcom_push, hook_impl): + kwargs = {parameter.name: "foo"} + if parameter.expected_parameters: + kwargs["parameters"] = parameter.expected_parameters + psrp_session_init = Mock(spec=Command) + op = PsrpOperator( + task_id='test_task_id', + psrp_conn_id=CONNECTION_ID, + psrp_session_init=psrp_session_init, + do_xcom_push=do_xcom_push, + **kwargs, + ) + ps = Mock(spec=PowerShell, output=[json.dumps("")], had_errors=had_errors) + hook_impl.configure_mock( + **{"return_value.__enter__.return_value.invoke.return_value.__enter__.return_value": ps} + ) if had_errors: exception_msg = "Process failed" with pytest.raises(AirflowException, match=exception_msg): op.execute(None) else: output = op.execute(None) - assert output == ps.output + assert output == [json.loads(output) for output in ps.output] if do_xcom_push else ps.output + is_logged = hook_impl.call_args[1]['on_output_callback'] == op.log.info + assert do_xcom_push ^ is_logged + expected_ps_calls = [ + call.add_command(psrp_session_init), + parameter.expected_method, + ] + if parameter.expected_parameters: + expected_ps_calls.extend([call.add_parameters({'bar': 'baz'})]) + if parameter.name in ("cmdlet", "powershell") and do_xcom_push: + expected_ps_calls.append( + call.add_cmdlet('ConvertTo-Json'), + ) + assert ps.mock_calls == expected_ps_calls + + def test_securestring_sandboxed(self): + op = PsrpOperator(psrp_conn_id=CONNECTION_ID, cmdlet='test') + template = op.get_template_env().from_string("{{ 'foo' | securestring }}") + with pytest.raises(AirflowException): + template.render() + + @patch.object(BaseOperator, "get_template_env") + def test_securestring_native(self, get_template_env): + op = PsrpOperator(psrp_conn_id=CONNECTION_ID, cmdlet='test') + get_template_env.return_value = NativeEnvironment() + template = op.get_template_env().from_string("{{ 'foo' | securestring }}") + rendered = template.render() + assert rendered.tag == "SS" + assert rendered.value == "foo"