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,
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:
ps.poll_invoke(timeout=self._operation_timeout)
for i, stream in enumerate(streams):
offset = offsets[i]
while len(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 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)
def check_creds(host, username, password, domain):
has_access = False
try:
username = '******' % (domain, username)
wsman = WSMan(host,
username=username,
password=password,
ssl=False,
auth='ntlm')
with RunspacePool(wsman) as pool:
ps = PowerShell(pool)
ps.add_cmdlet('Invoke-Expression')\
.add_parameter('Command', 'whoami')
ps.begin_invoke()
while ps.output is None or ps.output == '':
ps.poll_invoke()
ps.end_invoke()
if len(ps.output) > 0:
has_access = True
except:
has_access = False
return has_access
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
