Add Microsoft PSRP provider (#17361)

CONTRIBUTING.rst

@@ -592,10 +592,11 @@ cgroups, cloudant, cncf.kubernetes, crypto, dask, databricks, datadog, deprecate
 devel_all, devel_ci, devel_hadoop, dingding, discord, doc, docker, druid, elasticsearch, exasol,
 facebook, ftp, gcp, gcp_api, github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive,
 http, imap, jdbc, jenkins, jira, kerberos, kubernetes, ldap, leveldb, microsoft.azure,
-microsoft.mssql, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, opsgenie, oracle,
-pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole, rabbitmq, redis, s3,
-salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, snowflake, spark, sqlite,
-ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, winrm, yandex, zendesk
+microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas,
+opsgenie, oracle, pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole,
+rabbitmq, redis, s3, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack,
+snowflake, spark, sqlite, ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs,
+winrm, yandex, zendesk
 
   .. END EXTRAS HERE
 

INSTALL

@@ -96,10 +96,11 @@ cgroups, cloudant, cncf.kubernetes, crypto, dask, databricks, datadog, deprecate
 devel_all, devel_ci, devel_hadoop, dingding, discord, doc, docker, druid, elasticsearch, exasol,
 facebook, ftp, gcp, gcp_api, github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive,
 http, imap, jdbc, jenkins, jira, kerberos, kubernetes, ldap, leveldb, microsoft.azure,
-microsoft.mssql, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, opsgenie, oracle,
-pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole, rabbitmq, redis, s3,
-salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, snowflake, spark, sqlite,
-ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, winrm, yandex, zendesk
+microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas,
+opsgenie, oracle, pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole,
+rabbitmq, redis, s3, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack,
+snowflake, spark, sqlite, ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs,
+winrm, yandex, zendesk
 
 # END EXTRAS HERE
 

airflow/providers/microsoft/psrp/CHANGELOG.rst

@@ -0,0 +1,25 @@
+ .. 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.
+
+
+Changelog
+---------
+
+1.0.0
+.....
+
+Initial version of the provider.

airflow/providers/microsoft/psrp/__init__.py

@@ -0,0 +1,17 @@
+#
+# 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.

airflow/providers/microsoft/psrp/hooks/__init__.py

@@ -0,0 +1,17 @@
+#
+# 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.

airflow/providers/microsoft/psrp/hooks/psrp.py

@@ -0,0 +1,118 @@
+#
+# 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.
+
+from time import sleep
+
+from pypsrp.messages import ErrorRecord, InformationRecord, ProgressRecord
+from pypsrp.powershell import PowerShell, PSInvocationState, RunspacePool
+from pypsrp.wsman import WSMan
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class PSRPHook(BaseHook):
+    """
+    Hook for PowerShell Remoting Protocol execution.
+
+    The hook must be used as a context manager.
+    """
+
+    _client = None
+    _poll_interval = 1
+
+    def __init__(self, psrp_conn_id: str):
+        self.conn_id = psrp_conn_id
+
+    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__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        try:
+            self._client.__exit__()
+        finally:
+            self._client = None
+
+    def invoke_powershell(self, script: str) -> PowerShell:
+        with RunspacePool(self._client) as pool:
+            ps = PowerShell(pool)
+            ps.add_script(script)
+            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),
+            ]
+            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()
+
+                for (i, (stream, handler)) in enumerate(streams):
+                    offset = offsets[i]
+                    while len(stream) > offset:
+                        handler(stream[offset])
+                        offset += 1
+                    offsets[i] = offset
+
+            # For good measure, we'll make sure the process has
+            # stopped running.
+            ps.end_invoke()
+
+            if ps.streams.error:
+                raise AirflowException("Process had one or more errors")
+
+            self.log.info("Invocation state: %s", str(PSInvocationState(ps.state)))
+            return ps
+
+    def _log_output(self, message: str):
+        self.log.info("%s", message)
+
+    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
+
+        if isinstance(record, InformationRecord):
+            self.log.info("Information: %s", record.message_data)
+            return
+
+        if isinstance(record, ProgressRecord):
+            self.log.info("Progress: %s (%s)", record.activity, record.description)
+            return
+
+        self.log.info("Unsupported record type: %s", type(record).__name__)

airflow/providers/microsoft/psrp/operators/__init__.py

@@ -0,0 +1,17 @@
+#
+# 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.

airflow/providers/microsoft/psrp/operators/psrp.py

@@ -0,0 +1,66 @@
+#
+# 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.
+
+from typing import List, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.microsoft.psrp.hooks.psrp import PSRPHook
+
+
+class PSRPOperator(BaseOperator):
+    """PowerShell Remoting Protocol operator.
+
+    :param psrp_conn_id: connection id
+    :type psrp_conn_id: str
+    :param command: command to execute on remote host. (templated)
+    :type command: str
+    :param powershell: powershell to execute on remote host. (templated)
+    :type powershell: str
+    """
+
+    template_fields = (
+        "command",
+        "powershell",
+    )
+    template_fields_renderers = {"command": "powershell", "powershell": "powershell"}
+    ui_color = "#901dd2"
+
+    def __init__(
+        self,
+        *,
+        psrp_conn_id: str,
+        command: Optional[str] = None,
+        powershell: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        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
+
+    def execute(self, context: dict) -> 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

airflow/providers/microsoft/psrp/provider.yaml

@@ -0,0 +1,45 @@
+# 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.
+
+---
+package-name: apache-airflow-providers-microsoft-psrp
+name: PowerShell Remoting Protocol (PSRP)
+description: |
+    `PowerShell Remoting Protocol (PSRP)
+    <https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-psrp/>`__
+
+versions:
+  - 1.0.0
+
+additional-dependencies:
+  - pypsrp>=0.5.0
+
+integrations:
+  - integration-name: Windows Remote Management (WinRM)
+    external-doc-url: https://docs.microsoft.com/en-us/windows/win32/winrm/portal
+    logo: /integration-logos/winrm/WinRM.png
+    tags: [protocol]
+
+operators:
+  - integration-name: Windows Remote Management (WinRM)
+    python-modules:
+      - airflow.providers.microsoft.winrm.operators.winrm
+
+hooks:
+  - integration-name: Windows Remote Management (WinRM)
+    python-modules:
+      - airflow.providers.microsoft.winrm.hooks.winrm

docs/apache-airflow-providers-microsoft-psrp/commits.rst

@@ -0,0 +1,20 @@
+
+ .. 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.
+
+Package apache-airflow-providers-microsoft-psrp
+-----------------------------------------------