diff --git a/os_win/tests/unit/test_utilsfactory.py b/os_win/tests/unit/test_utilsfactory.py index 011987a8..0df0a73f 100644 --- a/os_win/tests/unit/test_utilsfactory.py +++ b/os_win/tests/unit/test_utilsfactory.py @@ -32,6 +32,7 @@ from os_win.utils.compute import rdpconsoleutils from os_win.utils.compute import vmutils from os_win.utils.dns import dnsutils from os_win.utils import hostutils +from os_win.utils.io import ioutils from os_win.utils.network import networkutils from os_win.utils import pathutils from os_win.utils import processutils @@ -145,6 +146,11 @@ class TestHyperVUtilsFactory(test_base.OsWinBaseTestCase): expected_class=processutils.ProcessUtils, class_type='processutils') + def test_get_ioutils(self): + self._check_get_class( + expected_class=ioutils.IOUtils, + class_type='ioutils') + def test_utils_public_signatures(self): for module_name in utilsfactory.utils_map.keys(): classes = utilsfactory.utils_map[module_name] diff --git a/os_win/tests/unit/utils/io/test_ioutils.py b/os_win/tests/unit/utils/io/test_ioutils.py index 5a1d3717..a14546a9 100644 --- a/os_win/tests/unit/utils/io/test_ioutils.py +++ b/os_win/tests/unit/utils/io/test_ioutils.py @@ -13,6 +13,7 @@ # License for the specific language governing permissions and limitations # under the License.import mock +import ddt import mock import six @@ -24,6 +25,7 @@ from os_win.utils.winapi import constants as w_const from os_win.utils.winapi import wintypes +@ddt.ddt class IOUtilsTestCase(test_base.BaseTestCase): _autospec_classes = [ @@ -63,6 +65,40 @@ class IOUtilsTestCase(test_base.BaseTestCase): **self._run_args) self.assertEqual(self._mock_run.return_value, ret_val) + @ddt.data({}, + {'inherit_handle': True}, + {'sec_attr': mock.sentinel.sec_attr}) + @ddt.unpack + @mock.patch.object(wintypes, 'HANDLE') + @mock.patch.object(wintypes, 'SECURITY_ATTRIBUTES') + def test_create_pipe(self, mock_sec_attr_cls, mock_handle_cls, + inherit_handle=False, sec_attr=None): + r, w = self._ioutils.create_pipe( + sec_attr, mock.sentinel.size, inherit_handle) + + exp_sec_attr = None + if sec_attr: + exp_sec_attr = sec_attr + elif inherit_handle: + exp_sec_attr = mock_sec_attr_cls.return_value + + self.assertEqual(mock_handle_cls.return_value.value, r) + self.assertEqual(mock_handle_cls.return_value.value, w) + + self._mock_run.assert_called_once_with( + ioutils.kernel32.CreatePipe, + self._ctypes.byref(mock_handle_cls.return_value), + self._ctypes.byref(mock_handle_cls.return_value), + self._ctypes.byref(exp_sec_attr) if exp_sec_attr else None, + mock.sentinel.size, + **self._run_args) + + if not sec_attr and exp_sec_attr: + self.assertEqual(inherit_handle, exp_sec_attr.bInheritHandle) + self.assertEqual(self._ctypes.sizeof.return_value, + exp_sec_attr.nLength) + self._ctypes.sizeof.assert_called_once_with(exp_sec_attr) + def test_wait_named_pipe(self): fake_timeout_s = 10 self._ioutils.wait_named_pipe(mock.sentinel.pipe_name, @@ -206,6 +242,26 @@ class IOUtilsTestCase(test_base.BaseTestCase): **self._run_args) mock_wait_io_completion.assert_called_once_with(mock_event) + @mock.patch.object(wintypes, 'DWORD') + def test_read_file(self, mock_dword): + num_bytes_read = mock_dword.return_value + + ret_val = self._ioutils.read_file( + mock.sentinel.handle, + mock.sentinel.buff, + mock.sentinel.num_bytes, + mock.sentinel.overlapped_struct) + + self.assertEqual(num_bytes_read.value, ret_val) + self._mock_run.assert_called_once_with( + ioutils.kernel32.ReadFile, + mock.sentinel.handle, + mock.sentinel.buff, + mock.sentinel.num_bytes, + self._ctypes.byref(num_bytes_read), + self._ctypes.byref(mock.sentinel.overlapped_struct), + **self._run_args) + @mock.patch.object(ioutils.IOUtils, '_reset_event') @mock.patch.object(ioutils.IOUtils, '_wait_io_completion') def test_write(self, mock_wait_io_completion, mock_reset_event): @@ -227,6 +283,25 @@ class IOUtilsTestCase(test_base.BaseTestCase): **self._run_args) mock_wait_io_completion.assert_called_once_with(mock_event) + @mock.patch.object(wintypes, 'DWORD') + def test_write_file(self, mock_dword): + num_bytes_written = mock_dword.return_value + ret_val = self._ioutils.write_file( + mock.sentinel.handle, + mock.sentinel.buff, + mock.sentinel.num_bytes, + mock.sentinel.overlapped_struct) + + self.assertEqual(num_bytes_written.value, ret_val) + self._mock_run.assert_called_once_with( + ioutils.kernel32.WriteFile, + mock.sentinel.handle, + mock.sentinel.buff, + mock.sentinel.num_bytes, + self._ctypes.byref(num_bytes_written), + self._ctypes.byref(mock.sentinel.overlapped_struct), + **self._run_args) + def test_buffer_ops(self): mock.patch.stopall() diff --git a/os_win/utils/io/ioutils.py b/os_win/utils/io/ioutils.py index 740e6b9c..cff859b6 100644 --- a/os_win/utils/io/ioutils.py +++ b/os_win/utils/io/ioutils.py @@ -61,6 +61,31 @@ class IOUtils(object): eventlet_nonblocking_mode=eventlet_blocking_mode) return self._win32_utils.run_and_check_output(*args, **kwargs) + def create_pipe(self, security_attributes=None, size=0, + inherit_handle=False): + """Create an anonymous pipe. + + The main advantage of this method over os.pipe is that it allows + creating inheritable pipe handles (which is flawed on most Python + versions). + """ + r = wintypes.HANDLE() + w = wintypes.HANDLE() + + if inherit_handle and not security_attributes: + security_attributes = wintypes.SECURITY_ATTRIBUTES() + security_attributes.bInheritHandle = inherit_handle + security_attributes.nLength = ctypes.sizeof(security_attributes) + + self._run_and_check_output( + kernel32.CreatePipe, + ctypes.byref(r), + ctypes.byref(w), + ctypes.byref(security_attributes) if security_attributes else None, + size) + + return r.value, w.value + @_utils.retry_decorator(exceptions=exceptions.Win32IOException, max_sleep_time=2) def wait_named_pipe(self, pipe_name, timeout=WAIT_PIPE_DEFAULT_TIMEOUT): @@ -155,6 +180,18 @@ class IOUtils(object): completion_routine) self._wait_io_completion(overlapped_structure.hEvent) + def read_file(self, handle, buff, num_bytes, overlapped_structure=None): + # Similar to IOUtils.read, but intended for synchronous operations. + num_bytes_read = wintypes.DWORD(0) + overlapped_structure_ref = ( + ctypes.byref(overlapped_structure) if overlapped_structure + else None) + self._run_and_check_output(kernel32.ReadFile, + handle, buff, num_bytes, + ctypes.byref(num_bytes_read), + overlapped_structure_ref) + return num_bytes_read.value + def write(self, handle, buff, num_bytes, overlapped_structure, completion_routine): self._reset_event(overlapped_structure.hEvent) @@ -164,6 +201,18 @@ class IOUtils(object): completion_routine) self._wait_io_completion(overlapped_structure.hEvent) + def write_file(self, handle, buff, num_bytes, overlapped_structure=None): + # Similar to IOUtils.write, but intended for synchronous operations. + num_bytes_written = wintypes.DWORD(0) + overlapped_structure_ref = ( + ctypes.byref(overlapped_structure) if overlapped_structure + else None) + self._run_and_check_output(kernel32.WriteFile, + handle, buff, num_bytes, + ctypes.byref(num_bytes_written), + overlapped_structure_ref) + return num_bytes_written.value + @classmethod def get_buffer(cls, buff_size, data=None): buff = (ctypes.c_ubyte * buff_size)() diff --git a/os_win/utils/winapi/libs/kernel32.py b/os_win/utils/winapi/libs/kernel32.py index 96bb69f2..5ec22b24 100644 --- a/os_win/utils/winapi/libs/kernel32.py +++ b/os_win/utils/winapi/libs/kernel32.py @@ -95,6 +95,14 @@ def register(): ] lib_handle.CreateFileW.restype = wintypes.HANDLE + lib_handle.CreatePipe.argtypes = [ + wintypes.PHANDLE, + wintypes.PHANDLE, + wintypes.PVOID, + wintypes.DWORD + ] + lib_handle.CreatePipe.restype = wintypes.BOOL + lib_handle.CreateSymbolicLinkW.argtypes = [ wintypes.LPCWSTR, wintypes.LPCWSTR, @@ -136,6 +144,15 @@ def register(): lib_handle.LocalFree.argtypes = [wintypes.HANDLE] lib_handle.LocalFree.restype = wintypes.HANDLE + lib_handle.ReadFile.argtypes = [ + wintypes.HANDLE, + wintypes.LPVOID, + wintypes.DWORD, + wintypes.LPDWORD, + wintypes.LPOVERLAPPED + ] + lib_handle.ReadFile.restype = wintypes.BOOL + lib_handle.ReadFileEx.argtypes = [ wintypes.HANDLE, wintypes.LPVOID, @@ -167,6 +184,15 @@ def register(): ] lib_handle.WaitNamedPipeW.restype = wintypes.BOOL + lib_handle.WriteFile.argtypes = [ + wintypes.HANDLE, + wintypes.LPCVOID, + wintypes.DWORD, + wintypes.LPDWORD, + wintypes.LPOVERLAPPED, + ] + lib_handle.WriteFile.restype = wintypes.BOOL + lib_handle.WriteFileEx.argtypes = [ wintypes.HANDLE, wintypes.LPCVOID, diff --git a/os_win/utils/winapi/wintypes.py b/os_win/utils/winapi/wintypes.py index 43fef84d..e1546b6e 100644 --- a/os_win/utils/winapi/wintypes.py +++ b/os_win/utils/winapi/wintypes.py @@ -101,3 +101,11 @@ if sys.platform == 'win32': None, DWORD, DWORD, LPOVERLAPPED) else: LPOVERLAPPED_COMPLETION_ROUTINE = PVOID + + +class SECURITY_ATTRIBUTES(ctypes.Structure): + _fields_ = [ + ('nLength', DWORD), + ('lpSecurityDescriptor', LPVOID), + ('bInheritHandle', BOOL) + ] diff --git a/os_win/utilsfactory.py b/os_win/utilsfactory.py index e30b96a5..d1b96fc8 100644 --- a/os_win/utilsfactory.py +++ b/os_win/utilsfactory.py @@ -132,7 +132,15 @@ utils_map = { 'ProcessUtils': { 'min_version': 6.2, 'max_version': None, - 'path': 'os_win.utils.processutils.ProcessUtils'}} + 'path': 'os_win.utils.processutils.ProcessUtils'}, + }, + 'ioutils': { + 'IOUtils': { + 'min_version': 6.2, + 'max_version': None, + 'path': 'os_win.utils.io.ioutils.IOUtils' + } + } } @@ -234,3 +242,7 @@ def get_migrationutils(): def get_processutils(): return _get_class(class_type='processutils') + + +def get_ioutils(): + return _get_class(class_type='ioutils')