Added retry to initial pipe create in case service hadn't started yet

pypsexec/__init__.py

@@ -10,4 +10,4 @@ def emit(self, record):
 logger = logging.getLogger(__name__)
 logger.addHandler(NullHandler())
 
-__version__ = '0.0.1.dev6'
+__version__ = '0.0.1.dev7'

pypsexec/client.py

@@ -2,6 +2,7 @@
 import os
 import socket
 import sys
+import time
 import uuid
 
 from smbprotocol.connection import Connection, NtStatus
@@ -260,6 +261,9 @@ def run_executable(self, executable, arguments=None, processors=None,
         if run_elevated and run_limited:
             raise PypsexecException("Both run_elevated and run_limited are "
                                     "set, only 1 of these can be true")
+        if stdin is not None and (asynchronous or interactive):
+            raise PypsexecException("Cannot send stdin data on an interactive "
+                                    "or asynchronous process")
 
         log.debug("Making sure PAExec service is running")
         self._service.start()
@@ -293,12 +297,29 @@ def run_executable(self, executable, arguments=None, processors=None,
         input_data['buffer'] = settings
 
         # write the settings to the main PAExec pipe
-        main_pipe = open_pipe(smb_tree, self._exe_file,
-                              FilePipePrinterAccessMask.GENERIC_READ |
-                              FilePipePrinterAccessMask.GENERIC_WRITE |
-                              FilePipePrinterAccessMask.FILE_APPEND_DATA |
-                              FilePipePrinterAccessMask.READ_CONTROL |
-                              FilePipePrinterAccessMask.SYNCHRONIZE)
+        pipe_access_mask = FilePipePrinterAccessMask.GENERIC_READ | \
+            FilePipePrinterAccessMask.GENERIC_WRITE | \
+            FilePipePrinterAccessMask.FILE_APPEND_DATA | \
+            FilePipePrinterAccessMask.READ_CONTROL | \
+            FilePipePrinterAccessMask.SYNCHRONIZE
+        for i in range(0, 3):
+            try:
+                main_pipe = open_pipe(smb_tree, self._exe_file,
+                                      pipe_access_mask)
+            except SMBResponseException as exc:
+                if exc.status != NtStatus.STATUS_OBJECT_NAME_NOT_FOUND:
+                    raise exc
+                elif i == 2:
+                    raise PypsexecException("Failed to open main PAExec pipe "
+                                            "%s, no more attempts remaining"
+                                            % self._exe_file)
+                log.warning("Main pipe %s does not exist yet on attempt %d. "
+                            "Trying again in 5 seconds"
+                            % (self._exe_file, i + 1))
+                time.sleep(5)
+            else:
+                break
+
         log.info("Writing PAExecSettingsMsg to the main PAExec pipe")
         log.info(str(input_data))
         main_pipe.write(input_data.pack(), 0)
@@ -345,15 +366,22 @@ def run_executable(self, executable, arguments=None, processors=None,
                 pass
 
             # send any input if there was any
-            if stdin and isinstance(stdin, bytes):
-                log.info("Sending stdin bytes over stdin pipe: %s"
-                         % self._stdin_pipe_name)
-                stdin_pipe.write(stdin)
-            elif stdin:
-                log.info("Sending stdin generator bytes over stdin pipe: %s"
-                         % self._stdin_pipe_name)
-                for stdin_data in stdin():
-                    stdin_pipe.write(stdin_data)
+            try:
+                if stdin and isinstance(stdin, bytes):
+                    log.info("Sending stdin bytes over stdin pipe: %s"
+                             % self._stdin_pipe_name)
+                    stdin_pipe.write(stdin)
+                elif stdin:
+                    log.info("Sending stdin generator bytes over stdin pipe: "
+                             "%s" % self._stdin_pipe_name)
+                    for stdin_data in stdin():
+                        stdin_pipe.write(stdin_data)
+            except SMBResponseException as exc:
+                # if it fails with a STATUS_PIPE_BROKEN exception, continue as
+                # the actual error will be in the response (process failed)
+                if exc.status != NtStatus.STATUS_PIPE_BROKEN:
+                    raise exc
+                log.warning("Failed to send data through stdin: %s" % str(exc))
 
         # read the final response from the process
         log.info("Reading result of PAExec process")

pypsexec/pipe.py

@@ -1,5 +1,4 @@
 import logging
-import sys
 import threading
 import warnings
 

tests/test_client.py

@@ -73,6 +73,24 @@ def test_proc_both_elevated_and_limited_error(self):
         assert str(exc.value) == "Both run_elevated and run_limited are " \
                                  "set, only 1 of these can be true"
 
+    def test_proc_stdin_and_async(self):
+        client = Client("username", "password", "server")
+        with pytest.raises(PypsexecException) as exc:
+            client.run_executable("whoami",
+                                  asynchronous=True,
+                                  stdin=b"")
+        assert str(exc.value) == "Cannot send stdin data on an interactive " \
+                                 "or asynchronous process"
+
+    def test_proc_stdin_and_interactive(self):
+        client = Client("username", "password", "server")
+        with pytest.raises(PypsexecException) as exc:
+            client.run_executable("whoami",
+                                  interactive=True,
+                                  stdin=b"")
+        assert str(exc.value) == "Cannot send stdin data on an interactive " \
+                                 "or asynchronous process"
+
 
 # these are the functional, they only run in the presence of env vars
 class TestClientFunctional(object):